#include <fmt/core.h>
#include <yaml-cpp/yaml.h>

#include <filesystem>
#include <fstream>
#include <opencv2/opencv.hpp>

#include "calibration/board_pattern.hpp"
#include "src/component/img_tools.hpp"
#include "src/component/logger.hpp"
#include "src/component/math_tools.hpp"
#include "src/device/camera.hpp"
#include "src/device/gimbal/gimbal.hpp"

const std::string keys =
  "{help h usage ?  |                          | 输出命令行参数说明}"
  "{@config-path c  | configs/calibration.yaml | yaml配置文件路径 }"
  "{output-folder o |      assets/img_with_q   | 输出文件夹路径   }";

void write_q(const std::string q_path, const Eigen::Quaterniond & q)
{
  std::ofstream q_file(q_path);
  Eigen::Vector4d xyzw = q.coeffs();
  // 输出顺序为wxyz
  q_file << fmt::format("{} {} {} {}", xyzw[3], xyzw[0], xyzw[1], xyzw[2]);
  q_file.close();
}

void capture_loop(
  const std::string & config_path, const std::string & output_folder)
{
  // 从配置文件加载标定板参数（支持 circles_grid 和 chessboard）
  auto yaml = YAML::LoadFile(config_path);
  auto board_pattern = calibration::load_board_pattern(yaml);

  device::Gimbal gimbal(config_path);
  device::Camera camera(config_path);
  cv::Mat img;
  std::chrono::steady_clock::time_point timestamp;

  int count = 0;
  while (true) {
    camera.read(img, timestamp);
    Eigen::Quaterniond q = gimbal.q(timestamp);

    // 在图像上显示欧拉角，用来判断imuabs系的xyz正方向，同时判断imu是否存在零漂
    auto img_with_ypr = img.clone();
    Eigen::Vector3d zyx = component::eulers(q, 2, 1, 0) * 57.3;  // degree
    component::draw_text(img_with_ypr, fmt::format("Z {:.2f}", zyx[0]), {40, 40}, {0, 0, 255});
    component::draw_text(img_with_ypr, fmt::format("Y {:.2f}", zyx[1]), {40, 80}, {0, 0, 255});
    component::draw_text(img_with_ypr, fmt::format("X {:.2f}", zyx[2]), {40, 120}, {0, 0, 255});

    std::vector<cv::Point2f> centers_2d;
    bool success;
    if (board_pattern.pattern_type == calibration::PatternType::chessboard) {
      // 棋盘格检测很慢，先在缩小图上快速检测，再映射回原图做亚像素精化
      cv::Mat small;
      double scale = 0.5;
      cv::resize(img, small, {}, scale, scale);
      std::vector<cv::Point2f> small_pts;
      success = calibration::find_pattern_points(small, board_pattern, small_pts);
      if (success) {
        for (auto & p : small_pts) { p.x /= scale; p.y /= scale; }
        cv::Mat gray;
        cv::cvtColor(img, gray, cv::COLOR_BGR2GRAY);
        cv::cornerSubPix(
          gray, small_pts, cv::Size(11, 11), cv::Size(-1, -1),
          cv::TermCriteria(cv::TermCriteria::COUNT + cv::TermCriteria::EPS, 30, 1e-3));
        centers_2d = std::move(small_pts);
      }
    } else {
      success = calibration::find_pattern_points(img, board_pattern, centers_2d);
    }
    cv::drawChessboardCorners(img_with_ypr, board_pattern.pattern_size, centers_2d, success);
    cv::resize(img_with_ypr, img_with_ypr, {}, 0.5, 0.5);  // 显示时缩小图片尺寸

    // 按“s”保存图片和对应四元数，按“q”退出程序
    cv::imshow("Press s to save, q to quit", img_with_ypr);
    auto key = cv::waitKey(1);
    if (key == 'q')
      break;
    else if (key != 's')
      continue;

    // 保存图片和四元数
    count++;
    auto img_path = fmt::format("{}/{}.jpg", output_folder, count);
    auto q_path = fmt::format("{}/{}.txt", output_folder, count);
    cv::imwrite(img_path, img);
    write_q(q_path, q);
    component::logger()->info("[{}] Saved in {}", count, output_folder);
  }

  // 离开该作用域时，camera和gimbal会自动关闭
}

int main(int argc, char * argv[])
{
  // 读取命令行参数
  cv::CommandLineParser cli(argc, argv, keys);
  if (cli.has("help")) {
    cli.printMessage();
    return 0;
  }
  auto config_path = cli.get<std::string>(0);
  auto output_folder = cli.get<std::string>("output-folder");

  // 新建输出文件夹
  std::filesystem::create_directory(output_folder);

  // 从配置文件读取标定板类型和尺寸
  auto yaml = YAML::LoadFile(config_path);
  auto board_pattern = calibration::load_board_pattern(yaml);
  component::logger()->info(
    "标定板类型: {}, 尺寸: {}列{}行",
    calibration::pattern_name(board_pattern.pattern_type),
    board_pattern.pattern_size.width, board_pattern.pattern_size.height);

  // 主循环，保存图片和对应四元数
  capture_loop(config_path, output_folder);

  component::logger()->warn("注意四元数输出顺序为wxyz");

  return 0;
}