// Copyright chenyy 2026. Licensed under the MIT License.
// YOLOv5-based armor detector with traditional vision refinement
//
// Uses YOLO to provide ROI regions, then applies traditional binary + light detection
// for precise corner localization. Input/output compatible with armor_detector.
//
// Supports multiple detection modes:
// - HYBRID_SYNC: YOLO runs every frame (original)
// - HYBRID_ROI_CACHE: YOLO runs every N frames, detection runs every frame
// - HYBRID_ASYNC: Async queue with ROI cache
//
// 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_HPP_
#define ARMOR_YOLO_DETECT_ARMOR_YOLO_DETECTOR_HPP_

#include <memory>
#include <string>
#include <vector>
#include <deque>
#include <atomic>
#include <mutex>
#include <condition_variable>

#include <opencv2/core.hpp>
// project
#include "armor_yolo_detect/types.hpp"
#include "armor_yolo_detect/yolo_tensorrt.hpp"
#include "armor_yolo_detect/detector_mode.hpp"
#include "armor_detector/types.hpp"
#include "armor_detector/light_corner_corrector.hpp"
#include "rm_utils/common.hpp"

namespace armor_yolo_detect {

class Detector {
public:
  struct Params {
    float conf_threshold = 0.65f;
    float nms_threshold = 0.45f;
    int detect_color = 0;  // 0=red, 1=blue
    // ROI expansion pixels
    int roi_expand_pixel = 75;
  };

  // Light detection params (from armor_detector)
  struct LightParams {
    double min_ratio = 0.08;
    double max_ratio = 0.4;
    double max_angle = 40.0;
    bool use_fit_line = true;
    int color_diff_thresh = 25;
  };

  // Armor matching params (from armor_detector)
  struct ArmorParams {
    double min_light_ratio = 0.6;
    double min_small_center_distance = 0.8;
    double max_small_center_distance = 3.2;
    double min_large_center_distance = 3.2;
    double max_large_center_distance = 5.0;
    double max_angle = 35.0;
  };

  Detector() = default;
  ~Detector() = default;

  /**
   * @brief Initialize detector with TensorRT engine and traditional vision params
   * @param engine_path Path to TensorRT engine file
   * @param params Detection parameters
   */
  bool init(const std::string& engine_path, const Params& params);

  /**
   * @brief Detect armors using YOLO for ROI + traditional vision for refinement
   * @param input Input image (BGR)
   * @return Vector of detected Armors (fyt::auto_aim::Armor type for compatibility)
   */
  std::vector<fyt::auto_aim::Armor> detect(const cv::Mat& input);

  /**
   * @brief Detect armors and return raw YOLO objects
   * @param input Input image (RGB)
   * @return Vector of detected YoloObjects
   */
  std::vector<YoloObject> detectRaw(const cv::Mat& input);

  // Set detection mode
  void setMode(DetectorMode mode);

  // Get current mode
  DetectorMode getMode() const { return mode_; }

  // Set ROI update interval (frames between YOLO updates)
  void setRoiUpdateInterval(int interval);

  // Parameters
  Params params;
  LightParams light_params;
  ArmorParams armor_params;
  int binary_thres = 160;

  // Debug
  bool enable_debug = false;
  cv::Mat debug_img;

  // Store ROI rectangles for debug drawing
  std::vector<cv::Rect> debug_rois_;

  // Draw detection results for debug visualization
  void drawResults(cv::Mat& img) noexcept;

private:
  // Traditional vision methods
  std::vector<fyt::auto_aim::Light> findLightsInROI(
      const cv::Mat& rgb_img, const cv::Mat& gray_img, const cv::Mat& binary_img,
      const cv::Rect& roi);
  bool isLight(const fyt::auto_aim::Light& light);
  bool containLight(const int i, const int j, const std::vector<fyt::auto_aim::Light>& lights);
  fyt::auto_aim::ArmorType isArmor(const fyt::auto_aim::Light& light_1, const fyt::auto_aim::Light& light_2);
  std::vector<fyt::auto_aim::Armor> matchLightsInROI(const std::vector<fyt::auto_aim::Light>& lights);
  void correctCorners(fyt::auto_aim::Armor& armor, const cv::Mat& gray_img);

  // Expand YOLO bbox to ROI
  cv::Rect expandBBox(const cv::Rect& bbox, const cv::Size& img_size, int pixel_expand);

  // Process ROIs to detect armors (common logic for all modes)
  std::vector<fyt::auto_aim::Armor> processROIs(
      const cv::Mat& input,
      const cv::Mat& gray_img,
      const cv::Mat& binary_img);

  std::unique_ptr<YoloTensorRT> tensorrt_;
  std::vector<YoloObject> yolo_objects_;
  std::vector<fyt::auto_aim::Armor> armors_;
  std::unique_ptr<fyt::auto_aim::LightCornerCorrector> corner_corrector_;

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

  // ROI cache for HYBRID_ROI_CACHE and HYBRID_ASYNC modes
  std::vector<YoloObject> cached_yolo_objects_;
  int roi_update_interval_ = 5;
  std::atomic<uint64_t> frame_counter_{0};
};

}  // namespace armor_yolo_detect
#endif  // ARMOR_YOLO_DETECT_ARMOR_YOLO_DETECTOR_HPP_
