// Copyright chenyy 2026. Licensed under the MIT License.
// ROS2 node implementation for YOLOv5 armor detection
//
// Copyright (C) FYT Vision Group. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "armor_yolo_detect/armor_yolo_detector_node.hpp"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstddef>
#include <filesystem>
#include <functional>
#include <iomanip>
#include <map>
#include <memory>
#include <limits>
#include <numeric>
#include <sstream>
#include <string>
#include <vector>

#include <cv_bridge/cv_bridge.h>
#include <tf2/LinearMath/Matrix3x3.h>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2/convert.h>
#include <tf2/time.h>
#include <tf2_ros/buffer.h>
#include <tf2_ros/create_timer_ros.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>

#include <image_transport/image_transport.hpp>
#include <rclcpp/qos.hpp>

#include <opencv2/core.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>

#include "armor_yolo_detect/types.hpp"
#include "armor_yolo_detect/detector_mode.hpp"
#include "armor_yolo_detect/async_detector.hpp"
#include "armor_detector/ba_solver.hpp"
#include "rm_utils/assert.hpp"
#include "rm_utils/common.hpp"
#include "rm_utils/logger/log.hpp"
#include "rm_utils/math/pnp_solver.hpp"
#include "rm_utils/math/utils.hpp"
#include "rm_utils/url_resolver.hpp"

namespace armor_yolo_detect {
namespace {

double durationMs(const std::chrono::steady_clock::duration& duration) {
  return std::chrono::duration<double, std::milli>(duration).count();
}

}  // namespace

ArmorYoloDetectorNode::ArmorYoloDetectorNode(const rclcpp::NodeOptions& options)
    : Node("armor_yolo_detector", options) {

  FYT_REGISTER_LOGGER("armor_yolo_detect", "~/fyt2024-log", INFO);
  FYT_INFO("armor_yolo_detect", "Starting ArmorYoloDetectorNode!");

  // Initialize detector
  detector_ = initDetector();

  const int max_queue_size_param = static_cast<int>(this->declare_parameter("max_queue_size", 3));
  max_queue_size_ = static_cast<size_t>(std::max(1, max_queue_size_param));
  process_every_n_frames_ =
      std::max(1, static_cast<int>(this->declare_parameter("process_every_n_frames", 1)));

  // BA setting
  use_ba_ = this->declare_parameter("use_ba", true);

  // Publishers
  armors_pub_ = this->create_publisher<rm_interfaces::msg::Armors>(
      "armor_detector/armors", rclcpp::SensorDataQoS());

  // Transform
  odom_frame_ = this->declare_parameter("target_frame", "odom");

  // Visualization markers
  armor_marker_.ns = "armors";
  armor_marker_.action = visualization_msgs::msg::Marker::ADD;
  armor_marker_.type = visualization_msgs::msg::Marker::CUBE;
  armor_marker_.scale.x = 0.03;
  armor_marker_.scale.y = 0.15;
  armor_marker_.scale.z = 0.12;
  armor_marker_.color.a = 1.0;
  armor_marker_.color.r = 1.0;
  armor_marker_.lifetime = rclcpp::Duration::from_seconds(0.1);

  text_marker_.ns = "classification";
  text_marker_.action = visualization_msgs::msg::Marker::ADD;
  text_marker_.type = visualization_msgs::msg::Marker::TEXT_VIEW_FACING;
  text_marker_.scale.z = 0.1;
  text_marker_.color.a = 1.0;
  text_marker_.color.r = 1.0;
  text_marker_.color.g = 1.0;
  text_marker_.color.b = 1.0;
  text_marker_.lifetime = rclcpp::Duration::from_seconds(0.1);

  marker_pub_ =
      this->create_publisher<visualization_msgs::msg::MarkerArray>("armor_detector/marker", 10);

  // Debug
  debug_ = this->declare_parameter("debug", false);
  debug_log_interval_frames_ =
      std::max(1, static_cast<int>(this->declare_parameter("debug_log_interval_frames", 30)));
  debug_terminal_log_ = this->declare_parameter("debug.enable_terminal_log", true);
  debug_markers_ = this->declare_parameter("debug.enable_markers", true);
  debug_result_img_ = this->declare_parameter("debug.enable_result_img", true);

  if (debug_) {
    createDebugPublishers();
  }

  debug_publishers_refresh_timer_ = this->create_wall_timer(
      std::chrono::milliseconds(100),
      [this]() { refreshDebugPublishersIfNeeded(); });

  debug_param_sub_ = std::make_shared<rclcpp::ParameterEventHandler>(this);
  debug_cb_handle_ =
      debug_param_sub_->add_parameter_callback("debug", [this](const rclcpp::Parameter& p) {
        std::lock_guard<std::mutex> lock(processing_mutex_);
        debug_ = p.as_bool();
        debug_ ? createDebugPublishers() : destroyDebugPublishers();
      });

  // Camera info
  cam_info_sub_ = this->create_subscription<sensor_msgs::msg::CameraInfo>(
      "camera_info", rclcpp::SensorDataQoS(),
      [this](sensor_msgs::msg::CameraInfo::SharedPtr camera_info) {
        std::lock_guard<std::mutex> lock(processing_mutex_);
        cam_center_ = cv::Point2f(camera_info->k[2], camera_info->k[5]);
        cam_info_ = std::make_shared<sensor_msgs::msg::CameraInfo>(*camera_info);
        armor_pose_estimator_ = std::make_unique<fyt::auto_aim::ArmorPoseEstimator>(cam_info_);
        armor_pose_estimator_->enableBA(use_ba_);
        cam_info_sub_.reset();
      });

  // Image subscription
  img_sub_ = this->create_subscription<sensor_msgs::msg::Image>(
      "image_raw", rclcpp::SensorDataQoS(),
      std::bind(&ArmorYoloDetectorNode::imageCallback, this, std::placeholders::_1));

  // TF
  tf2_buffer_ = std::make_shared<tf2_ros::Buffer>(this->get_clock());
  auto timer_interface = std::make_shared<tf2_ros::CreateTimerROS>(
      this->get_node_base_interface(), this->get_node_timers_interface());
  tf2_buffer_->setCreateTimerInterface(timer_interface);
  tf2_listener_ = std::make_shared<tf2_ros::TransformListener>(*tf2_buffer_);

  // Mode service
  set_mode_srv_ = this->create_service<rm_interfaces::srv::SetMode>(
      "armor_detector/set_mode",
      std::bind(&ArmorYoloDetectorNode::setModeCallback, this,
                std::placeholders::_1, std::placeholders::_2));

  // Image processing thread
  image_processing_thread_ = std::thread(&ArmorYoloDetectorNode::processImageLoop, this);
  FYT_INFO("armor_yolo_detect", "Image worker started: threads=1, max_queue_size={}", max_queue_size_);

  heartbeat_ = fyt::HeartBeatPublisher::create(this);
}

ArmorYoloDetectorNode::~ArmorYoloDetectorNode() {
  {
    std::lock_guard<std::mutex> lock(queue_mutex_);
    stop_processing_ = true;
  }
  queue_cv_.notify_one();
  if (image_processing_thread_.joinable()) {
    image_processing_thread_.join();
  }
}

void ArmorYoloDetectorNode::imageCallback(const sensor_msgs::msg::Image::ConstSharedPtr img_msg) {
  if ((input_frame_counter_.fetch_add(1, std::memory_order_relaxed) + 1) %
          static_cast<uint64_t>(process_every_n_frames_) !=
      0) {
    return;
  }

  auto frame_start = std::chrono::steady_clock::now();
  {
    std::lock_guard<std::mutex> lock(queue_mutex_);
    if (frame_queue_.size() >= max_queue_size_) {
      frame_queue_.pop_front();
    }
    frame_queue_.push_back(PendingFrame{.img_msg = img_msg, .frame_start = frame_start});
  }
  queue_cv_.notify_one();
}

void ArmorYoloDetectorNode::processImageLoop() {
  while (true) {
    PendingFrame frame;
    {
      std::unique_lock<std::mutex> lock(queue_mutex_);
      queue_cv_.wait(lock, [this]() { return stop_processing_ || !frame_queue_.empty(); });
      if (stop_processing_ && frame_queue_.empty()) {
        return;
      }
      frame = std::move(frame_queue_.front());
      frame_queue_.pop_front();
    }
    processFrame(std::move(frame));
  }
}

void ArmorYoloDetectorNode::processFrame(PendingFrame frame) {
  bool debug_enabled = false;
  bool debug_terminal_log_enabled = false;
  bool debug_marker_enabled = false;
  {
    std::lock_guard<std::mutex> lock(processing_mutex_);
    debug_enabled = debug_;
    debug_terminal_log_enabled = debug_terminal_log_;
    debug_marker_enabled = debug_markers_;
  }

  auto stage_start = std::chrono::steady_clock::now();
  Eigen::Matrix3d imu_to_camera = Eigen::Matrix3d::Identity();

  // Get TF transform
  try {
    rclcpp::Time target_time = frame.img_msg->header.stamp;
    auto odom_to_gimbal = tf2_buffer_->lookupTransform(
        odom_frame_, frame.img_msg->header.frame_id, target_time,
        rclcpp::Duration::from_seconds(0.01));
    auto msg_q = odom_to_gimbal.transform.rotation;
    tf2::Quaternion tf_q;
    tf2::fromMsg(msg_q, tf_q);
    tf2::Matrix3x3 tf2_matrix = tf2::Matrix3x3(tf_q);
    imu_to_camera << tf2_matrix.getRow(0)[0], tf2_matrix.getRow(0)[1], tf2_matrix.getRow(0)[2],
        tf2_matrix.getRow(1)[0], tf2_matrix.getRow(1)[1], tf2_matrix.getRow(1)[2],
        tf2_matrix.getRow(2)[0], tf2_matrix.getRow(2)[1], tf2_matrix.getRow(2)[2];
  } catch (const tf2::TransformException& ex) {
    FYT_ERROR("armor_yolo_detect", "lookUpTransform failed: {}", ex.what());
    return;
  }

  // Detect armors
  double detect_ms = 0.0;
  double debug_image_ms = 0.0;
  stage_start = std::chrono::steady_clock::now();
  auto armors = detectArmors(frame.img_msg, &detect_ms, &debug_image_ms);
  const double detect_total_ms = durationMs(std::chrono::steady_clock::now() - stage_start);

  // Init message
  rm_interfaces::msg::Armors armors_msg;
  armors_msg.header = frame.img_msg->header;
  armors_msg.armors.clear();

  // Extract armor poses
  stage_start = std::chrono::steady_clock::now();
  {
    std::lock_guard<std::mutex> lock(processing_mutex_);
    if (armor_pose_estimator_ != nullptr) {
      armors_msg.armors = armor_pose_estimator_->extractArmorPoses(armors, imu_to_camera);
    } else {
      FYT_WARN("armor_yolo_detect", "PnP Failed!");
    }
  }
  const double pose_ms = durationMs(std::chrono::steady_clock::now() - stage_start);

  // Publish markers
  stage_start = std::chrono::steady_clock::now();
  {
    std::lock_guard<std::mutex> lock(processing_mutex_);
    if (debug_enabled && debug_marker_enabled) {
      armors_msg_ = armors_msg;
      marker_array_.markers.clear();
      armor_marker_.id = 0;
      text_marker_.id = 0;
      armor_marker_.header = text_marker_.header = armors_msg_.header;

      for (const auto& armor : armors_msg_.armors) {
        armor_marker_.pose = armor.pose;
        armor_marker_.id++;
        text_marker_.pose.position = armor.pose.position;
        text_marker_.id++;
        text_marker_.pose.position.y -= 0.1;
        text_marker_.text = armor.number;
        marker_array_.markers.emplace_back(armor_marker_);
        marker_array_.markers.emplace_back(text_marker_);
      }
      publishMarkers();
    }
  }
  const double marker_ms = durationMs(std::chrono::steady_clock::now() - stage_start);

  // Publish detected armors
  stage_start = std::chrono::steady_clock::now();
  armors_pub_->publish(armors_msg);
  const double publish_ms = durationMs(std::chrono::steady_clock::now() - stage_start);

  if (debug_enabled && debug_terminal_log_enabled &&
      (debug_log_counter_.fetch_add(1, std::memory_order_relaxed) + 1) %
              static_cast<uint64_t>(debug_log_interval_frames_) ==
          0) {
    const double total_ms = durationMs(std::chrono::steady_clock::now() - frame.frame_start);
    FYT_INFO("armor_yolo_detect",
             "Detector Latency: total={:.3f}ms detect={:.3f}ms pose={:.3f}ms "
             "marker={:.3f}ms publish={:.3f}ms",
             total_ms, detect_total_ms, pose_ms, marker_ms, publish_ms);
  }
}

std::unique_ptr<Detector> ArmorYoloDetectorNode::initDetector() {
  Detector::Params params;

  // Load engine path (0526 model, TensorRT .trt format)
  namespace fs = std::filesystem;
  const std::string engine_url =
      this->declare_parameter("engine_path", "package://armor_yolo_detect/model/yolov5_0526.trt");
  fs::path engine_path = fyt::utils::URLResolver::getResolvedPath(engine_url);

  if (!fs::exists(engine_path)) {
    FYT_ERROR("armor_yolo_detect", "Engine file not found: {}", engine_path.string());
    return nullptr;
  }

  params.conf_threshold = this->declare_parameter("conf_threshold", 0.65);
  params.nms_threshold = this->declare_parameter("nms_threshold", 0.50);
  params.detect_color = this->declare_parameter("detect_color", -1);
  params.roi_expand_pixel = this->declare_parameter("roi_expand_pixel", 75);

  auto detector = std::make_unique<Detector>();
  if (!detector->init(engine_path.string(), params)) {
    FYT_ERROR("armor_yolo_detect", "Failed to initialize detector!");
    return nullptr;
  }

  // Binary threshold for light detection
  detector->binary_thres = this->declare_parameter("binary_thres", 160);

  // Light detection params
  detector->light_params.min_ratio = this->declare_parameter("light.min_ratio", 0.08);
  detector->light_params.max_ratio = this->declare_parameter("light.max_ratio", 0.4);
  detector->light_params.max_angle = this->declare_parameter("light.max_angle", 40.0);
  detector->light_params.use_fit_line = this->declare_parameter("light.use_fit_line", true);
  detector->light_params.color_diff_thresh =
      static_cast<int>(this->declare_parameter("light.color_diff_thresh", 25));

  // Armor matching params
  detector->armor_params.min_light_ratio =
      this->declare_parameter("armor.min_light_ratio", 0.6);
  detector->armor_params.min_small_center_distance =
      this->declare_parameter("armor.min_small_center_distance", 0.8);
  detector->armor_params.max_small_center_distance =
      this->declare_parameter("armor.max_small_center_distance", 3.2);
  detector->armor_params.min_large_center_distance =
      this->declare_parameter("armor.min_large_center_distance", 3.2);
  detector->armor_params.max_large_center_distance =
      this->declare_parameter("armor.max_large_center_distance", 5.0);
  detector->armor_params.max_angle = this->declare_parameter("armor.max_angle", 35.0);

  // Detection mode
  int mode_int = this->declare_parameter("detector_mode", 0);
  detector_mode_ = static_cast<DetectorMode>(mode_int);
  detector->setMode(detector_mode_);

  // ROI update interval (for roi_cache and async modes)
  int roi_update_interval = this->declare_parameter("roi_update_interval", 5);
  detector->setRoiUpdateInterval(roi_update_interval);

  FYT_INFO("armor_yolo_detect", "Inference backend: TensorRT ({})", engine_path.string());
  FYT_INFO("armor_yolo_detect", "Detector mode: {} (roi_update_interval={})",
           detectorModeToString(detector_mode_), roi_update_interval);

  // Set dynamic parameter callback
  on_set_parameters_callback_handle_ = this->add_on_set_parameters_callback(
      std::bind(&ArmorYoloDetectorNode::onSetParameters, this, std::placeholders::_1));

  return detector;
}

std::vector<fyt::auto_aim::Armor> ArmorYoloDetectorNode::detectArmors(
    const sensor_msgs::msg::Image::ConstSharedPtr& img_msg, double* detect_ms,
    double* debug_publish_ms) {
  std::lock_guard<std::mutex> lock(processing_mutex_);

  auto stage_start = std::chrono::steady_clock::now();

  // Convert ROS img to cv::Mat
  auto img = cv_bridge::toCvShare(img_msg, "bgr8")->image;

  // Run detection: YOLO for ROI + traditional vision for refinement
  auto armors = detector_->detect(img);

  if (detect_ms != nullptr) {
    *detect_ms = durationMs(std::chrono::steady_clock::now() - stage_start);
  }

  // Debug image publishing
  stage_start = std::chrono::steady_clock::now();
  if (debug_ && debug_result_img_) {
    detector_->debug_img = img.clone();

    cv::putText(detector_->debug_img,
                "det=" + std::to_string(armors.size()),
                cv::Point(20, 40), cv::FONT_HERSHEY_SIMPLEX, 1.0,
                cv::Scalar(255, 255, 0), 2);

    // Draw detection results using armor_detector's drawResults style
    detector_->drawResults(detector_->debug_img);

    // Publish result image
    result_img_pub_.publish(
      cv_bridge::CvImage(img_msg->header, "bgr8", detector_->debug_img).toImageMsg());
  }

  if (debug_publish_ms != nullptr) {
    *debug_publish_ms = durationMs(std::chrono::steady_clock::now() - stage_start);
  }

  return armors;
}

fyt::auto_aim::Armor ArmorYoloDetectorNode::convertToAutoAimArmor(const Armor& yolo_armor) {
  return yolo_armor.toAutoAimArmor();
}

rcl_interfaces::msg::SetParametersResult ArmorYoloDetectorNode::onSetParameters(
    std::vector<rclcpp::Parameter> parameters) {
  rcl_interfaces::msg::SetParametersResult result;
  result.successful = true;

  std::lock_guard<std::mutex> lock(processing_mutex_);

  for (const auto& param : parameters) {
    if (param.get_name() == "conf_threshold") {
      detector_->params.conf_threshold = param.as_double();
    } else if (param.get_name() == "nms_threshold") {
      detector_->params.nms_threshold = param.as_double();
    } else if (param.get_name() == "detect_color") {
      detector_->params.detect_color = param.as_int();
    } else if (param.get_name() == "roi_expand_pixel") {
      detector_->params.roi_expand_pixel = param.as_int();
    } else if (param.get_name() == "binary_thres") {
      detector_->binary_thres = param.as_int();
    } else if (param.get_name() == "light.min_ratio") {
      detector_->light_params.min_ratio = param.as_double();
    } else if (param.get_name() == "light.max_ratio") {
      detector_->light_params.max_ratio = param.as_double();
    } else if (param.get_name() == "light.max_angle") {
      detector_->light_params.max_angle = param.as_double();
    } else if (param.get_name() == "light.use_fit_line") {
      detector_->light_params.use_fit_line = param.as_bool();
    } else if (param.get_name() == "light.color_diff_thresh") {
      detector_->light_params.color_diff_thresh = param.as_int();
    } else if (param.get_name() == "armor.min_light_ratio") {
      detector_->armor_params.min_light_ratio = param.as_double();
    } else if (param.get_name() == "armor.min_small_center_distance") {
      detector_->armor_params.min_small_center_distance = param.as_double();
    } else if (param.get_name() == "armor.max_small_center_distance") {
      detector_->armor_params.max_small_center_distance = param.as_double();
    } else if (param.get_name() == "armor.min_large_center_distance") {
      detector_->armor_params.min_large_center_distance = param.as_double();
    } else if (param.get_name() == "armor.max_large_center_distance") {
      detector_->armor_params.max_large_center_distance = param.as_double();
    } else if (param.get_name() == "armor.max_angle") {
      detector_->armor_params.max_angle = param.as_double();
    } else if (param.get_name() == "detector_mode") {
      detector_mode_ = static_cast<DetectorMode>(param.as_int());
      detector_->setMode(detector_mode_);
      FYT_INFO("armor_yolo_detect", "Detector mode changed to: {}",
               detectorModeToString(detector_mode_));
    } else if (param.get_name() == "roi_update_interval") {
      detector_->setRoiUpdateInterval(param.as_int());
    } else if (param.get_name() == "debug.enable_terminal_log") {
      debug_terminal_log_ = param.as_bool();
    } else if (param.get_name() == "debug.enable_markers") {
      debug_markers_ = param.as_bool();
    } else if (param.get_name() == "debug.enable_result_img") {
      debug_result_img_ = param.as_bool();
    }
  }

  return result;
}

void ArmorYoloDetectorNode::setModeCallback(
    const std::shared_ptr<rm_interfaces::srv::SetMode::Request> request,
    const std::shared_ptr<rm_interfaces::srv::SetMode::Response> response) {
  response->success = true;
  response->message = "0";

  fyt::VisionMode mode = static_cast<fyt::VisionMode>(request->mode);
  std::string mode_name = fyt::visionModeToString(mode);

  if (mode_name == "UNKNOWN") {
    FYT_ERROR("armor_yolo_detect", "Invalid mode: {}", request->mode);
    return;
  }

  auto createImageSub = [this]() {
    if (img_sub_ == nullptr) {
      img_sub_ = this->create_subscription<sensor_msgs::msg::Image>(
          "image_raw", rclcpp::SensorDataQoS(),
          std::bind(&ArmorYoloDetectorNode::imageCallback, this, std::placeholders::_1));
    }
  };

  switch (mode) {
    case fyt::VisionMode::AUTO_AIM_RED: {
      detector_->params.detect_color = 0;  // RED
      createImageSub();
      break;
    }
    case fyt::VisionMode::AUTO_AIM_BLUE: {
      detector_->params.detect_color = 1;  // BLUE
      createImageSub();
      break;
    }
    default: {
      img_sub_.reset();
      std::lock_guard<std::mutex> lock(queue_mutex_);
      frame_queue_.clear();
    }
  }

  FYT_WARN("armor_yolo_detect", "Set mode to {}", mode_name);
}

void ArmorYoloDetectorNode::createDebugPublishers() {
  result_img_pub_ = image_transport::create_publisher(this, "armor_detector/result_img");
}

void ArmorYoloDetectorNode::destroyDebugPublishers() {
  result_img_pub_.shutdown();
}

void ArmorYoloDetectorNode::refreshDebugPublishersIfNeeded() {
  if (!debug_publishers_refresh_requested_.exchange(false, std::memory_order_relaxed)) {
    return;
  }

  std::lock_guard<std::mutex> lock(processing_mutex_);
  destroyDebugPublishers();
  if (debug_) {
    createDebugPublishers();
  }
}

void ArmorYoloDetectorNode::publishMarkers() {
  using Marker = visualization_msgs::msg::Marker;
  armor_marker_.action = armors_msg_.armors.empty() ? Marker::DELETEALL : Marker::ADD;
  marker_array_.markers.emplace_back(armor_marker_);
  marker_pub_->publish(marker_array_);
}

}  // namespace armor_yolo_detect

#include "rclcpp_components/register_node_macro.hpp"

RCLCPP_COMPONENTS_REGISTER_NODE(armor_yolo_detect::ArmorYoloDetectorNode)
