// Copyright chenyy 2026. Licensed under the MIT License.
// TensorRT inference engine 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_YOLO_TENSORRT_HPP_
#define ARMOR_YOLO_DETECT_YOLO_TENSORRT_HPP_

#include <memory>
#include <string>
#include <vector>
#include <cmath>
#include <cstdint>

#include <NvInfer.h>
#include <cuda_runtime_api.h>
#include <opencv2/core.hpp>

#include "armor_yolo_detect/types.hpp"

namespace armor_yolo_detect {

class YoloTensorRT {
public:
  static constexpr int INPUT_WIDTH = 640;
  static constexpr int INPUT_HEIGHT = 640;
  static constexpr int INPUT_C = 3;
  static constexpr int OUTPUT_ROWS = 25200;
  static constexpr int OUTPUT_COLS = 23;

  static constexpr float CONF_THRESHOLD = 0.65f;
  static constexpr float NMS_THRESHOLD = 0.45f;

  YoloTensorRT() = default;
  ~YoloTensorRT();

  bool initFromEngine(const std::string& engine_path);

  bool buildFromOnnx(
      const std::string& onnx_path,
      const std::string& engine_path,
      int max_batch_size = 1);

  std::vector<YoloObject> infer(const cv::Mat& img, int detect_color);

  void setPostprocessOptions(float conf_threshold, float nms_threshold, int max_detections) {
    conf_threshold_ = std::max(0.0f, conf_threshold);
    nms_threshold_ = std::max(0.0f, nms_threshold);
    max_detections_ = std::max(1, max_detections);
  }

  void setDebugOptions(bool dump_model_output, bool dump_all_rows, int dump_topk, int dump_interval) {
    dump_model_output_ = dump_model_output;
    dump_all_rows_ = dump_all_rows;
    dump_topk_ = std::max(1, dump_topk);
    dump_interval_ = std::max(1, dump_interval);
  }

  bool isInitialized() const { return engine_ != nullptr && context_ != nullptr; }

private:
  cv::Mat preprocess(const cv::Mat& img);

  bool inferInternal(const float* input, float* output, int batch_size);

  std::vector<YoloObject> postprocess(const float* output, int detect_color, bool dump_this_frame);

  std::vector<int> applyNMS(
      const std::vector<cv::Rect>& boxes,
      const std::vector<float>& scores,
      float nms_threshold);

  float sigmoid(float x) const {
    if (x > 0) {
      return 1.0f / (1.0f + std::exp(-x));
    } else {
      return std::exp(x) / (1.0f + std::exp(x));
    }
  }

  nvinfer1::IRuntime* runtime_ = nullptr;
  nvinfer1::ICudaEngine* engine_ = nullptr;
  nvinfer1::IExecutionContext* context_ = nullptr;

  void* device_buffer_[2] = {nullptr, nullptr};
  int input_index_ = 0;
  int output_index_ = 0;
  int io_tensor_count_ = 0;
  int output_rows_ = OUTPUT_ROWS;
  int output_cols_ = OUTPUT_COLS;
  std::string input_tensor_name_;
  std::string output_tensor_name_;

  bool is_input_half_ = false;
  std::vector<uint16_t> half_input_buffer_;
  std::vector<float> input_buffer_;
  std::vector<float> output_buffer_;

  cudaStream_t cuda_stream_ = nullptr;

  bool dump_model_output_ = false;
  bool dump_all_rows_ = false;
  int dump_topk_ = 5;
  int dump_interval_ = 30;
  uint64_t infer_counter_ = 0;

  float conf_threshold_ = CONF_THRESHOLD;
  float nms_threshold_ = NMS_THRESHOLD;
  int max_detections_ = 20;
};

}  // namespace armor_yolo_detect
#endif  // ARMOR_YOLO_DETECT_YOLO_TENSORRT_HPP_
