#include <stdio.h>
// #include <mutex> // 已移除
// #include <chrono> // 已移除
#include <string>
#include <vector>
// #include <algorithm> // 已移除

#include "postprocess.h"
#include "preprocess.h"
#include "rkYolov5s.hpp"
#include "rknn/coreNum.hpp"

#include "opencv2/core/core.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/imgproc/imgproc.hpp"
#include "rknn/rknn_api.h"

// trigger_alarm 和 get_current_time_seconds 已被移至 video_service.cc

static void dump_tensor_attr(rknn_tensor_attr *attr) {
  std::string shape_str = attr->n_dims < 1 ? "" : std::to_string(attr->dims[0]);
  for (int i = 1; i < attr->n_dims; ++i) {
    shape_str += ", " + std::to_string(attr->dims[i]);
  }
}

static unsigned char *load_data(FILE *fp, size_t ofst, size_t sz) {
  unsigned char *data;
  int ret;

  data = NULL;

  if (NULL == fp) {
    return NULL;
  }

  ret = fseek(fp, ofst, SEEK_SET);
  if (ret != 0) {
    printf("blob seek failure.\n");
    return NULL;
  }

  data = (unsigned char *)malloc(sz);
  if (data == NULL) {
    printf("buffer malloc failure.\n");
    return NULL;
  }
  ret = fread(data, 1, sz, fp);
  return data;
}

static unsigned char *load_model(const char *filename, int *model_size) {
  FILE *fp;
  unsigned char *data;

  fp = fopen(filename, "rb");
  if (NULL == fp) {
    printf("Open file %s failed.\n", filename);
    return NULL;
  }

  fseek(fp, 0, SEEK_END);
  int size = ftell(fp);

  data = load_data(fp, 0, size);

  fclose(fp);

  *model_size = size;
  return data;
}

static int saveFloat(const char *file_name, float *output, int element_size) {
  FILE *fp;
  fp = fopen(file_name, "w");
  for (int i = 0; i < element_size; i++) {
    fprintf(fp, "%.6f\n", output[i]);
  }
  fclose(fp);
  return 0;
}

rkYolov5s::rkYolov5s(const std::string &model_path,
                     const std::string &label_path, int class_num) {
  this->model_path = model_path;
  this->m_label_path = label_path; // [新增]
  this->m_class_num = class_num;   // [新增]

  nms_threshold = NMS_THRESH;
  box_conf_threshold = BOX_THRESH;
}

// set_intrusion_zone 已被移除

int rkYolov5s::init(rknn_context *ctx_in, bool share_weight) {
  printf("Loading model...\n");
  int model_data_size = 0;
  model_data = load_model(model_path.c_str(), &model_data_size);
  if (share_weight == true)
    ret = rknn_dup_context(ctx_in, &ctx);
  else
    ret = rknn_init(&ctx, model_data, model_data_size, 0, NULL);
  if (ret < 0) {
    printf("rknn_init error ret=%d\n", ret);
    return -1;
  }

  rknn_core_mask core_mask;
  switch (get_core_num()) {
  case 0:
    core_mask = RKNN_NPU_CORE_0;
    break;
  case 1:
    core_mask = RKNN_NPU_CORE_1;
    break;
  case 2:
    core_mask = RKNN_NPU_CORE_2;
    break;
  }
  ret = rknn_set_core_mask(ctx, core_mask);
  if (ret < 0) {
    printf("rknn_init core error ret=%d\n", ret);
    return -1;
  }

  rknn_sdk_version version;
  ret = rknn_query(ctx, RKNN_QUERY_SDK_VERSION, &version,
                   sizeof(rknn_sdk_version));
  if (ret < 0) {
    printf("rknn_init error ret=%d\n", ret);
    return -1;
  }
  printf("sdk version: %s driver version: %s\n", version.api_version,
         version.drv_version);

  ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
  if (ret < 0) {
    printf("rknn_init error ret=%d\n", ret);
    return -1;
  }
  printf("model input num: %d, output num: %d\n", io_num.n_input,
         io_num.n_output);

  input_attrs =
      (rknn_tensor_attr *)calloc(io_num.n_input, sizeof(rknn_tensor_attr));
  for (int i = 0; i < io_num.n_input; i++) {
    input_attrs[i].index = i;
    ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]),
                     sizeof(rknn_tensor_attr));
    if (ret < 0) {
      printf("rknn_init error ret=%d\n", ret);
      return -1;
    }
    dump_tensor_attr(&(input_attrs[i]));
  }

  output_attrs =
      (rknn_tensor_attr *)calloc(io_num.n_output, sizeof(rknn_tensor_attr));
  for (int i = 0; i < io_num.n_output; i++) {
    output_attrs[i].index = i;
    ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]),
                     sizeof(rknn_tensor_attr));
    dump_tensor_attr(&(output_attrs[i]));
  }

  if (input_attrs[0].fmt == RKNN_TENSOR_NCHW) {
    printf("model is NCHW input fmt\n");
    channel = input_attrs[0].dims[1];
    height = input_attrs[0].dims[2];
    width = input_attrs[0].dims[3];
  } else {
    printf("model is NHWC input fmt\n");
    height = input_attrs[0].dims[1];
    width = input_attrs[0].dims[2];
    channel = input_attrs[0].dims[3];
  }
  printf("model input height=%d, width=%d, channel=%d\n", height, width,
         channel);

  memset(inputs, 0, sizeof(inputs));
  inputs[0].index = 0;
  inputs[0].type = RKNN_TENSOR_UINT8;
  inputs[0].size = width * height * channel;
  inputs[0].fmt = RKNN_TENSOR_NHWC;
  inputs[0].pass_through = 0;

  static std::mutex postprocess_init_mutex;
  static bool postprocess_initialized = false;

  std::lock_guard<std::mutex> lock(postprocess_init_mutex);
  if (!postprocess_initialized) {
    // 调用我们新定义的 initPostProcess
    if (initPostProcess(m_label_path.c_str(), m_class_num) == 0) {
      postprocess_initialized = true;
      printf("PostProcess initialized successfully with %d classes from %s\n",
             m_class_num, m_label_path.c_str());
    } else {
      printf("Failed to initialize PostProcess!\n");
      return -1; // 初始化失败
    }
  }

  return 0; // rknn 初始化成功
}

rknn_context *rkYolov5s::get_pctx() { return &ctx; }

detect_result_group_t rkYolov5s::infer(const cv::Mat &orig_img) {
  // std::lock_guard<std::mutex> lock(mtx); // 已移除
  cv::Mat img;
  cv::cvtColor(orig_img, img, cv::COLOR_BGR2RGB);
  img_width = img.cols;
  img_height = img.rows;

  BOX_RECT pads;
  memset(&pads, 0, sizeof(BOX_RECT));
  cv::Size target_size(width, height);
  cv::Mat resized_img(target_size.height, target_size.width, CV_8UC3);

  float scale_w = (float)target_size.width / img.cols;
  float scale_h = (float)target_size.height / img.rows;

  if (img_width != width || img_height != height) {
    rga_buffer_t src;
    rga_buffer_t dst;
    memset(&src, 0, sizeof(src));
    memset(&dst, 0, sizeof(dst));
    ret = resize_rga(src, dst, img, resized_img, target_size);
    if (ret != 0) {
      fprintf(stderr, "resize with rga error\n");
    }
    inputs[0].buf = resized_img.data;
  } else {
    inputs[0].buf = img.data;
  }

  rknn_inputs_set(ctx, io_num.n_input, inputs);

  rknn_output outputs[io_num.n_output];
  memset(outputs, 0, sizeof(outputs));
  for (int i = 0; i < io_num.n_output; i++) {
    outputs[i].want_float = 0;
  }

  ret = rknn_run(ctx, NULL);
  ret = rknn_outputs_get(ctx, io_num.n_output, outputs, NULL);

  detect_result_group_t detect_result_group;
  std::vector<float> out_scales;
  std::vector<int32_t> out_zps;
  for (int i = 0; i < io_num.n_output; ++i) {
    out_scales.push_back(output_attrs[i].scale);
    out_zps.push_back(output_attrs[i].zp);
  }
  post_process((int8_t *)outputs[0].buf, (int8_t *)outputs[1].buf,
               (int8_t *)outputs[2].buf, height, width, box_conf_threshold,
               nms_threshold, pads, scale_w, scale_h, out_zps, out_scales,
               m_class_num, &detect_result_group);

  ret = rknn_outputs_release(ctx, io_num.n_output, outputs);

  // 返回原始检测结果
  return detect_result_group;
}

rkYolov5s::~rkYolov5s() {
  // deinitPostProcess();

  ret = rknn_destroy(ctx);

  if (model_data)
    free(model_data);

  if (input_attrs)
    free(input_attrs);
  if (output_attrs)
    free(output_attrs);
}