#include "sdkconfig.h"

#include 
#include 
#include 
#include 
#include 
#include 
#include 

#include 
#include 

#include 
#include 
#include 
#include 
#include 

#include "cli.hpp"
#include "heap_monitor.hpp"
#include "mjpeg_packetizer.hpp"
#include "nvs.hpp"
#include "rtsp_server.hpp"
#include "task.hpp"
#include "task_monitor.hpp"
#include "wifi_sta.hpp"
#include "wifi_sta_menu.hpp"

#include "esp_camera.h"

#if CONFIG_CAMERA_STREAMER_BOARD_XIAO_ESP32S3_SENSE
#include 
#include 

#include "generic_packetizer.hpp"
#include "xiao-esp32s3-sense.hpp"
#else
#include "esp32-timer-cam.hpp"
#endif

using namespace std::chrono_literals;

#if CONFIG_CAMERA_STREAMER_BOARD_XIAO_ESP32S3_SENSE
using Board = espp::XiaoEsp32S3Sense;
static constexpr const char *board_name = "Seeed Studio XIAO ESP32S3 Sense";
static constexpr const char *mdns_hostname_prefix = "xiao-camera-streamer";
static constexpr const char *mdns_instance_name = "XIAO Camera Streamer";
#else
using Board = espp::EspTimerCam;
static constexpr const char *board_name = "M5Stack ESP32 TimerCam";
static constexpr const char *mdns_hostname_prefix = "camera-streamer";
static constexpr const char *mdns_instance_name = "Camera Streamer";
#endif

static espp::Logger logger({.tag = "Camera Streamer", .level = espp::Logger::Verbosity::INFO});

namespace {
constexpr auto idle_capture_poll_period = 250ms;
constexpr auto dma_pressure_backoff = 250ms;
constexpr auto camera_capture_error_backoff = 100ms;
constexpr auto target_video_period = 100ms;
constexpr size_t min_dma_free_bytes_for_streaming = 12 * 1024;
constexpr size_t min_dma_largest_block_for_streaming = 4 * 1024;

constexpr int video_track_id = 0;
constexpr char rtsp_path[] = "mjpeg/1";

#if CONFIG_CAMERA_STREAMER_BOARD_XIAO_ESP32S3_SENSE
constexpr auto audio_capture_timeout = 50ms;
constexpr int audio_track_id = 1;
constexpr uint32_t microphone_sample_rate_hz =
    espp::XiaoEsp32S3Sense::microphone_default_sample_rate_hz();
constexpr size_t microphone_frame_samples = microphone_sample_rate_hz / 50; // 20 ms
constexpr size_t microphone_frame_bytes = microphone_frame_samples * sizeof(int16_t);
#endif

// Wait until `deadline`, returning true if the task was asked to stop while waiting.
template 
bool wait_until_or_stop(std::mutex &m, std::condition_variable &cv, bool &task_notified,
                        const std::chrono::time_point &deadline) {
  std::unique_lock lk(m);
  auto stop_requested = cv.wait_until(lk, deadline, [&task_notified] { return task_notified; });
  task_notified = false;
  return stop_requested;
}
} // namespace

std::recursive_mutex server_mutex;
std::unique_ptr camera_task;
std::unique_ptr memory_monitor_task;
std::unique_ptr task_monitor;
std::shared_ptr rtsp_server;
bool mdns_started{false};
std::atomic video_frames_streamed{0};

#if CONFIG_CAMERA_STREAMER_BOARD_XIAO_ESP32S3_SENSE
std::unique_ptr audio_task;
i2s_chan_handle_t microphone_rx_channel{nullptr};
std::atomic audio_frames_streamed{0};
#endif

esp_err_t initialize_camera(void);
bool start_rtsp_server(std::string_view server_address, int server_port);
void stop_streaming(void);
bool camera_task_fn(std::mutex &m, std::condition_variable &cv, bool &task_notified);
bool memory_monitor_task_fn(std::mutex &m, std::condition_variable &cv, bool &task_notified);

#if CONFIG_CAMERA_STREAMER_BOARD_XIAO_ESP32S3_SENSE
esp_err_t initialize_microphone(void);
void deinitialize_microphone(void);
bool initialize_and_test_sdcard(void);
bool audio_task_fn(std::mutex &m, std::condition_variable &cv, bool &task_notified);
#endif

extern "C" void app_main(void) {
  logger.info("Bootup");
  logger.info("Board: {}", board_name);

#if CONFIG_ESP32_WIFI_NVS_ENABLED
  std::error_code ec;
  // initialize NVS, needed for WiFi
  espp::Nvs nvs;
  nvs.init(ec);
#endif

  static constexpr float disconnected_led_breathing_period = 1.0f;
  static constexpr float connected_led_breathing_period = 3.5f;

  auto &board = Board::get();

#if !CONFIG_CAMERA_STREAMER_BOARD_XIAO_ESP32S3_SENSE
  // initialize RTC (TimerCam has an onboard BM8563)
  if (!board.initialize_rtc()) {
    logger.error("Could not initialize RTC");
    return;
  }
#endif

  // initialize LED
  if (!board.initialize_led(disconnected_led_breathing_period)) {
    logger.error("Could not initialize LED");
    return;
  }
  board.start_led_breathing();

#if CONFIG_CAMERA_STREAMER_BOARD_XIAO_ESP32S3_SENSE
  // initialize the (optional) microSD card; continues on failure
  initialize_and_test_sdcard();
#endif

  // initialize camera
  logger.info("Initializing camera");
  auto err = initialize_camera();
  if (err != ESP_OK) {
    logger.error("Could not initialize camera: {} '{}'", err, esp_err_to_name(err));
    return;
  }

#if CONFIG_CAMERA_STREAMER_BOARD_XIAO_ESP32S3_SENSE
  // initialize microphone (XIAO ESP32S3 Sense has an onboard PDM microphone)
  logger.info("Initializing microphone");
  err = initialize_microphone();
  if (err != ESP_OK) {
    logger.error("Could not initialize microphone: {} '{}'", err, esp_err_to_name(err));
    deinitialize_microphone();
    return;
  }
#endif

  logger.info("Starting memory monitors");
  task_monitor = std::make_unique(espp::TaskMonitor::Config{.period = 30s});
  memory_monitor_task = espp::Task::make_unique(espp::Task::Config{
      .callback = memory_monitor_task_fn,
      .task_config =
          {
              .name = "Memory Monitor",
              .stack_size_bytes = 6 * 1024,
          },
      .log_level = espp::Logger::Verbosity::WARN,
  });
  memory_monitor_task->start();

  // initialize WiFi
  logger.info("Initializing WiFi");
  espp::WifiSta wifi_sta(
      {.ssid = CONFIG_ESP_WIFI_SSID,
       .password = CONFIG_ESP_WIFI_PASSWORD,
       .num_connect_retries = CONFIG_ESP_MAXIMUM_RETRY,
       .on_connected =
           []() { Board::get().set_led_breathing_period(connected_led_breathing_period); },
       .on_disconnected =
           []() {
             Board::get().set_led_breathing_period(disconnected_led_breathing_period);
             logger.info("WiFi disconnected, stopping RTSP streaming");
             stop_streaming();
           },
       .on_got_ip =
           [](ip_event_got_ip_t *eventdata) {
             auto server_address = fmt::format("{}.{}.{}.{}", IP2STR(&eventdata->ip_info.ip));
             logger.info("got IP: {}", server_address);
             // start from a clean slate in case a previous session is still up
             stop_streaming();

             if (esp_wifi_set_ps(WIFI_PS_NONE) != ESP_OK) {
               logger.warn(
                   "Could not disable WiFi power save; RTSP streaming may hit TX backpressure");
             } else {
               logger.info("Disabled WiFi power save for RTSP streaming");
             }

             if (!start_rtsp_server(server_address, CONFIG_RTSP_SERVER_PORT)) {
               logger.error("RTSP server failed to start, not starting capture tasks");
               return;
             }

             logger.info("Creating camera task");
             auto new_camera_task = espp::Task::make_unique(espp::Task::Config{
                 .callback = camera_task_fn,
                 .task_config = {.name = "Camera Task", .priority = 10},
                 .log_level = espp::Logger::Verbosity::WARN,
             });
             if (!new_camera_task->start()) {
               logger.error("Could not start camera task");
               stop_streaming();
               return;
             }

#if CONFIG_CAMERA_STREAMER_BOARD_XIAO_ESP32S3_SENSE
             logger.info("Creating audio task");
             auto new_audio_task = espp::Task::make_unique(espp::Task::Config{
                 .callback = audio_task_fn,
                 .task_config =
                     {
                         .name = "Audio Task",
                         .stack_size_bytes = 6 * 1024,
                         .priority = 9,
                     },
                 .log_level = espp::Logger::Verbosity::WARN,
             });
             if (!new_audio_task->start()) {
               logger.error("Could not start audio task");
               new_camera_task.reset();
               stop_streaming();
               return;
             }
#endif

             std::lock_guard lock(server_mutex);
             camera_task = std::move(new_camera_task);
#if CONFIG_CAMERA_STREAMER_BOARD_XIAO_ESP32S3_SENSE
             audio_task = std::move(new_audio_task);
#endif
           }});

  espp::WifiStaMenu sta_menu(wifi_sta);
  auto root_menu = sta_menu.get();
  root_menu->Insert(
      "memory",
      [](std::ostream &out) {
        out