#include "./interface/DecoderManager.h"
#include <mutex>
#include <thread>
#include <chrono>
#include <string>

#include "acl/acl.h"
#include "acl/ops/acl_dvpp.h"

#include <opencv2/opencv.hpp>

#include "../ai_engine_module/face_det_ai_engine.h"

using namespace std;

struct decode_cbk_userdata{
  string task_id;
  void* opaque;
  void* opaque1;
};

deque<DeviceMemory*> m_RgbDataList;
mutex m_DataListMtx; 

thread* m_pAlgorthimThread{nullptr};
bool m_bfinish{false};
int m_devId = 0;
const char* task_id = "test0";
int skip_frame_ = 5;
int m_batch_size = 10;

face_det_ai_engine m_face_det_ai_engine;

aclrtContext ctx;

void algorthim_process_thread();
void algorthim_face_detect(vector<DeviceMemory*> vec_gpuMem);

static long long get_cur_time_ms(){
    chrono::time_point<chrono::system_clock, chrono::milliseconds> tpMicro
        = chrono::time_point_cast<chrono::milliseconds>(chrono::system_clock::now());

    return tpMicro.time_since_epoch().count();
}

static void post_decod_cbk(const void * userPtr, DeviceMemory* devFrame){
    do{
        if(m_bfinish){
        break;
        }
        m_DataListMtx.lock();
        if(m_RgbDataList.size() >= 30){
        m_DataListMtx.unlock();
        std::this_thread::sleep_for(std::chrono::milliseconds(3));
        continue;
        }
        m_RgbDataList.push_back(devFrame);
        m_DataListMtx.unlock();
        break;
    }while (true);
}

static void decode_finished_cbk(const void * userPtr){
    decode_cbk_userdata* ptr = (decode_cbk_userdata*)userPtr;
    if (ptr!= nullptr){
        printf("task finished: %s \n", ptr->task_id.c_str());
    }
    delete ptr;
    ptr = nullptr;
}

int main(){

    // 算法初始化
    string models_dir = ".";
    
    aclInit(nullptr);
    aclrtSetDevice(m_devId);

    // 人脸检测初始化
    facedet_ai_engine_param fd_param;
    char model_path_yolov5s[100];
    strcpy(model_path_yolov5s, (models_dir + "/models/face_detect/face_det_yolov5s_310p.om").c_str());
    fd_param.sdk_param.det_modelNames = model_path_yolov5s;
    char model_path_ldmk[100];
    strcpy(model_path_ldmk, (models_dir + "/models/face_detect/face_ldmk_310p.om").c_str());
    fd_param.sdk_param.ldmk_modelNames = model_path_ldmk;
    char model_path_pose[100];
    strcpy(model_path_pose, (models_dir + "/models/face_detect/face_pose_310p.om").c_str());
    fd_param.sdk_param.pose_modelNames = model_path_pose;
    char model_path_score[100];
    strcpy(model_path_score, (models_dir + "/models/face_detect/face_score_310p.om").c_str());
    fd_param.sdk_param.score_modelNames = model_path_score;
    char model_path_fuzzy[100];
    strcpy(model_path_fuzzy, (models_dir + "/models/face_detect/face_fuzzy_310p.om").c_str());
    fd_param.sdk_param.fuzzy_modelNames = model_path_fuzzy;
    char model_path_occlusion[100];
    strcpy(model_path_occlusion, (models_dir + "/models/face_detect/face_occlusion_310p.om").c_str());
    fd_param.sdk_param.occlusion_modelNames = model_path_occlusion;
    fd_param.sdk_param.thresld = 0.6;
    fd_param.sdk_param.devId = m_devId;
    fd_param.sdk_param.auth_license = "sy_tongtu_aiplatform_sdk_2023";
    fd_param.sdk_param.facial_fea_point_config = SY_CONFIG_OPEN;		//是否启动关键点检测
    fd_param.sdk_param.pose_config = SY_CONFIG_OPEN;					//是否启动姿态角
    fd_param.sdk_param.quality_config = SY_CONFIG_OPEN;				//是否启动质量检测
    fd_param.sdk_param.score_config = SY_CONFIG_OPEN;				//是否启动人脸置信度  //SY_CONFIG_OPEN  SY_CONFIG_CLOSE
    fd_param.sdk_param.max_result_count = 50;
    int ret = m_face_det_ai_engine.init_ai_engine(fd_param);
    if (ret < 0 ) {
      printf("Init face detection failed");
      return ret;
    }

    // facedet_ai_engine_param fd_param;
    // fd_param.sdk_param.det_modelNames = "./models/face_detect/face_det_yolov5s_310p.om";
    // fd_param.sdk_param.ldmk_modelNames = "./models/face_detect/face_ldmk_310p.om";
    // fd_param.sdk_param.pose_modelNames = "./models/face_detect/face_pose_310p.om";
    // fd_param.sdk_param.score_modelNames = "./models/face_detect/face_score_310p.om";
    // fd_param.sdk_param.fuzzy_modelNames = "./models/face_detect/face_fuzzy_310p.om";
    // fd_param.sdk_param.occlusion_modelNames = "./models/face_detect/face_occlusion_310p.om";
    // fd_param.sdk_param.thresld = 0.6;
    // fd_param.sdk_param.devId = m_devId;
    // fd_param.sdk_param.auth_license = "sy_tongtu_aiplatform_sdk_2023";
    // fd_param.sdk_param.facial_fea_point_config = SY_CONFIG_OPEN;		//是否启动关键点检测
    // fd_param.sdk_param.pose_config = SY_CONFIG_OPEN;					//是否启动姿态角
    // fd_param.sdk_param.quality_config = SY_CONFIG_OPEN;				//是否启动质量检测
    // fd_param.sdk_param.score_config = SY_CONFIG_OPEN;				//是否启动人脸置信度  //SY_CONFIG_OPEN  SY_CONFIG_CLOSE
    // fd_param.sdk_param.max_result_count = 50;
    // int ret = m_face_det_ai_engine.init_ai_engine(fd_param);
    // if (ret < 0 ) {
    //     printf("Init face detection failed \n");
    //     return ret;
    // }
    m_face_det_ai_engine.add_tracker(task_id, skip_frame_); // 跳帧数暂时写死


    // 创建解码任务
    DecoderManager* pDecManager = DecoderManager::getInstance();

    MgrDecConfig config;
    config.name = task_id;
    config.cfg.uri = "/opt/share/data/caishenkezhan.mp4";
    config.cfg.post_decoded_cbk = post_decod_cbk;
    config.cfg.decode_finished_cbk = decode_finished_cbk;
    config.cfg.force_tcp = true;  // rtsp用tcp
    config.cfg.gpuid = to_string(m_devId);
    config.cfg.skip_frame = skip_frame_;

    config.dec_type = DECODER_TYPE_DVPP;
    
    AbstractDecoder* dec = pDecManager->createDecoder(config);
    if (!dec){
        printf("创建解码器失败 \n");
        return false;
    }

    decode_cbk_userdata* userPtr = new decode_cbk_userdata;
    userPtr->task_id = string(task_id);
    pDecManager->setPostDecArg(config.name, userPtr);
    pDecManager->setFinishedDecArg(config.name, userPtr);


    int input_image_width = 0;
    int input_image_height = 0;
    pDecManager->getResolution(config.name, input_image_width, input_image_height);

    
    // 创建算法线程
    m_pAlgorthimThread = new thread([](void* arg) {
      algorthim_process_thread();
      return (void*)0;
    }
    , nullptr);

    pDecManager->startDecodeByName(config.name);

    while (getchar() != 'q');
}

void algorthim_process_thread(){
    
    aclrtSetDevice(m_devId);
    aclrtCreateContext(&ctx, m_devId);

    while (true){
        if(m_bfinish){
            break;
        }

        vector<DeviceMemory*> vec_gpuMem;
        m_DataListMtx.lock();
        while (!m_RgbDataList.empty()){
            DeviceMemory* gpuMem = m_RgbDataList.front();
            if(gpuMem->getMem() == nullptr){
              // 错误数据，直接删除
              delete gpuMem;
              gpuMem = nullptr;
              printf("mem is null \n");
            } else {
              vec_gpuMem.push_back(gpuMem);
            }
            m_RgbDataList.pop_front();
            if(vec_gpuMem.size() >= m_batch_size){
                break;
            }
        }
        m_DataListMtx.unlock();

        if(vec_gpuMem.size() <= 0){
            std::this_thread::sleep_for(std::chrono::milliseconds(3));
            continue;
        }

        aclrtSetCurrentContext(ctx);
        algorthim_face_detect(vec_gpuMem);

        for(int i=0;i < vec_gpuMem.size(); i++){
            DeviceMemory* mem = vec_gpuMem[i];
            if(mem->getSize() <= 0){
              continue;
            }
            delete mem;
            mem = nullptr;
        }
        vec_gpuMem.clear();

    }
    
    aclrtDestroyContext(ctx);

    printf("algorthim_process_thread exit. \n");
}

void algorthim_face_detect(vector<DeviceMemory*> vec_gpuMem) {
 
  vector<string> interest_task_list;
  vector<sy_img> interest_imgs;
  vector<DeviceMemory*> vec_vptMem;

  for (int i = 0; i < vec_gpuMem.size(); i++) {
      DeviceMemory* mem = vec_gpuMem[i];

      sy_img img;
      img.w_ = mem->getWidth();
      img.h_ = mem->getHeight();
      img.c_ = mem->getChannel();
      img.data_ = mem->getMem();
      interest_imgs.push_back(img);
      interest_task_list.push_back(mem->getId());
      vec_vptMem.push_back(mem);
  }

  cv::Scalar color;
  color[0]=0;
  color[1]=0;
  color[2]=255;
  if (!interest_imgs.empty()) {

    unsigned image_size = interest_imgs.size();

    // 人脸检测、跟踪
    std::vector<onelevel_det_result> facedet_result(image_size);
    std::vector<std::vector<int>> face_deleteObjectID(image_size);

    int ret = m_face_det_ai_engine.ai_engine_process_batch(interest_task_list, interest_imgs.data(), facedet_result, face_deleteObjectID);
    if(ret <= 0){
      printf("face detect error!!! \n");
      return;
    }



    aclrtSetCurrentContext(ctx);
    for(int idx=0; idx < vec_vptMem.size(); idx++){

      DeviceMemory* memPtr = vec_vptMem[idx];
      string task_id = memPtr->getId();
      int channel = memPtr->getChannel();
      int height = memPtr->getHeight();
      int width = memPtr->getWidth();

      if (0 == facedet_result[idx].obj_count) {
        continue;
      }

      int nSize = channel * height * width;
      unsigned char* cpu_data = (unsigned char *)malloc(nSize * sizeof(unsigned char));
      aclError aclRet = aclrtMemcpy(cpu_data, nSize, memPtr->getMem(), nSize, ACL_MEMCPY_DEVICE_TO_HOST);
      cv::Mat img_(height, width, CV_8UC3, cpu_data);

      for (int c = 0; c < facedet_result[idx].obj_count; c++) {
        
        det_objinfo obj_info = facedet_result[idx].obj[c];
        string str_obj_id = to_string(obj_info.id);
        cv::rectangle(img_,cv::Point(obj_info.left,obj_info.top),cv::Point(obj_info.right,obj_info.bottom),color,2);
        cv::putText(img_,str_obj_id.c_str(),cv::Point(obj_info.center_x,obj_info.center_y),cv::FONT_HERSHEY_SIMPLEX,2,cv::Scalar(0,0,255),4,8);
      }

      string file_name = "./res/recode/" + task_id + "_" + to_string(memPtr->getFrameNb()) + ".jpg";
      bool bWrite = cv::imwrite(file_name, img_);

      free(cpu_data);
    }





    for (int i = 0; i < face_deleteObjectID.size(); ++i){
      std::vector<int>().swap(face_deleteObjectID[i]);
    }
    std::vector<std::vector<int>>().swap(face_deleteObjectID);
    std::vector<onelevel_det_result>().swap(facedet_result);
  }

  aclFinalize();
}