// Copyright chenyy 2026. Licensed under the MIT License.
// ROS2 node 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.

#ifndef ARMOR_YOLO_DETECT_ARMOR_YOLO_DETECTOR_NODE_HPP_
#define ARMOR_YOLO_DETECT_ARMOR_YOLO_DETECTOR_NODE_HPP_

#include <memory>
#include <mutex>
#include <thread>
#include <deque>
#include <atomic>
#include <condition_variable>

#include <rclcpp/rclcpp.hpp>
#include <rclcpp/parameter_event_handler.hpp>
#include <sensor_msgs/msg/image.hpp>
#include <sensor_msgs/msg/camera_info.hpp>
#include <geometry_msgs/msg/pose_stamped.hpp>
#include <visualization_msgs/msg/marker_array.hpp>
#include <image_transport/image_transport.hpp>
#include <rm_interfaces/msg/armors.hpp>
#include <rm_interfaces/srv/set_mode.hpp>

#include <cv_bridge/cv_bridge.h>
#include <tf2_ros/buffer.h>
#include <tf2_ros/transform_listener.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>

#include "armor_yolo_detect/armor_yolo_detector.hpp"
#include "armor_yolo_detect/types.hpp"
#include "armor_yolo_detect/detector_mode.hpp"
#include "armor_detector/armor_pose_estimator.hpp"
#include "rm_utils/logger/log.hpp"
#include "rm_utils/heartbeat.hpp"

namespace armor_yolo_detect {

struct PendingFrame {
  sensor_msgs::msg::Image::ConstSharedPtr img_msg;
  std::chrono::steady_clock::time_point frame_start;
};

class ArmorYoloDetectorNode : public rclcpp::Node {
public:
  explicit ArmorYoloDetectorNode(const rclcpp::NodeOptions& options);
  ~ArmorYoloDetectorNode();

private:
  // Initialization
  std::unique_ptr<Detector> initDetector();

  // Image processing
  void imageCallback(const sensor_msgs::msg::Image::ConstSharedPtr img_msg);
  void processImageLoop();
  void processFrame(PendingFrame frame);
  std::vector<fyt::auto_aim::Armor> detectArmors(
      const sensor_msgs::msg::Image::ConstSharedPtr& img_msg,
      double* detect_ms,
      double* debug_publish_ms);

  // Mode switching
  void setModeCallback(
      const std::shared_ptr<rm_interfaces::srv::SetMode::Request> request,
      const std::shared_ptr<rm_interfaces::srv::SetMode::Response> response);

  // Debug publishers
  void createDebugPublishers();
  void destroyDebugPublishers();
  void refreshDebugPublishersIfNeeded();
  void publishMarkers();

  // Parameter callbacks
  rcl_interfaces::msg::SetParametersResult onSetParameters(
      std::vector<rclcpp::Parameter> parameters);

  // Convert YOLO armor to auto_aim armor for pose estimation
  fyt::auto_aim::Armor convertToAutoAimArmor(const Armor& yolo_armor);

  // ROS interfaces
  rclcpp::Subscription<sensor_msgs::msg::Image>::SharedPtr img_sub_;
  rclcpp::Subscription<sensor_msgs::msg::CameraInfo>::SharedPtr cam_info_sub_;
  rclcpp::Publisher<rm_interfaces::msg::Armors>::SharedPtr armors_pub_;
  rclcpp::Publisher<visualization_msgs::msg::MarkerArray>::SharedPtr marker_pub_;
  rclcpp::Service<rm_interfaces::srv::SetMode>::SharedPtr set_mode_srv_;

  // Image transport for debug
  image_transport::Publisher result_img_pub_;

  // TF
  std::shared_ptr<tf2_ros::Buffer> tf2_buffer_;
  std::shared_ptr<tf2_ros::TransformListener> tf2_listener_;
  std::string odom_frame_;

  // Detector
  std::unique_ptr<Detector> detector_;
  cv::Point2f cam_center_;

  // Detection mode
  DetectorMode detector_mode_ = DetectorMode::HYBRID_SYNC;

  // Camera info for pose estimation
  sensor_msgs::msg::CameraInfo::SharedPtr cam_info_;
  std::unique_ptr<fyt::auto_aim::ArmorPoseEstimator> armor_pose_estimator_;

  // Processing thread
  std::thread image_processing_thread_;
  std::deque<PendingFrame> frame_queue_;
  std::mutex queue_mutex_;
  std::condition_variable queue_cv_;
  bool stop_processing_ = false;
  size_t max_queue_size_;
  int process_every_n_frames_ = 1;
  std::atomic<uint64_t> input_frame_counter_{0};

  // Debug
  bool debug_ = false;
  bool debug_markers_ = true;
  bool debug_result_img_ = true;
  bool debug_terminal_log_ = true;
  int debug_log_interval_frames_ = 30;
  std::atomic<uint64_t> debug_log_counter_{0};
  rcl_interfaces::msg::SetParametersResult::SharedPtr result_;
  rclcpp::node_interfaces::OnSetParametersCallbackHandle::SharedPtr
      on_set_parameters_callback_handle_;
  std::shared_ptr<rclcpp::ParameterEventHandler> debug_param_sub_;
  rclcpp::ParameterCallbackHandle::SharedPtr debug_cb_handle_;
  rclcpp::TimerBase::SharedPtr debug_publishers_refresh_timer_;
  std::atomic<bool> debug_publishers_refresh_requested_{false};
  rclcpp::TimerBase::SharedPtr stage_timer_;

  // BA
  bool use_ba_ = true;

  // Message
  rm_interfaces::msg::Armors armors_msg_;
  visualization_msgs::msg::MarkerArray marker_array_;
  visualization_msgs::msg::Marker armor_marker_;
  visualization_msgs::msg::Marker text_marker_;

  // Processing mutex for detector and pose estimator
  std::mutex processing_mutex_;

  // Heartbeat
  fyt::HeartBeatPublisher::SharedPtr heartbeat_;
};

}  // namespace armor_yolo_detect
#endif  // ARMOR_YOLO_DETECT_ARMOR_YOLO_DETECTOR_NODE_HPP_
