#include "FFNvDecoderManager.h"
#include <iostream>

#include <thread>

#include <chrono>

#include "check_tool.h"

#include "cuda_kernels.h"
#include "NvJpegEncoder.h"

#include "opencv2\opencv.hpp"

#include "DogPoseDetectorOnnx.h"


using namespace std;
using namespace cv;

unsigned char *pHwRgb[2] = {nullptr, nullptr};

int sum1 = 0;
int sum2 = 0;

DogPoseDetectorOnnx poseDetector;


mutex m_mutex;
void saveFrame(AVFrame * gpuFrame, string file_name) {
	std::lock_guard<std::mutex> l(m_mutex);


	unsigned char *pHwData = nullptr;
	cudaError_t cudaStatus = cudaMalloc((void **)&pHwData, 3 * gpuFrame->width * gpuFrame->height * sizeof(unsigned char));

	cuda_common::setColorSpace(ITU709, 0);
	cudaStatus = cuda_common::CUDAToBGR((CUdeviceptr)gpuFrame->data[0], (CUdeviceptr)gpuFrame->data[1], gpuFrame->linesize[0], gpuFrame->linesize[1], pHwData, gpuFrame->width, gpuFrame->height);
	cudaDeviceSynchronize();
	if (cudaStatus != cudaSuccess) {
		cout << "CUDAToBGR failed !!!" << endl;
		return;
	}

	string path = file_name + ".jpg";
	saveJpeg(path.c_str(), pHwData, gpuFrame->width, gpuFrame->height, nullptr);  // 验证 CUDAToRGB 

	cudaDeviceSynchronize();
	cudaFree(pHwData);
	pHwData = nullptr;
}

static void showResult(unsigned char *pGpuBgr, int src_width, int src_height, std::vector<DogPoseResult> vec_result) {
	int buf_size = 3 * src_width * src_height;
	unsigned char* pBuf = new unsigned char[buf_size];
	cudaMemcpy(pBuf, pGpuBgr, buf_size * sizeof(unsigned char), cudaMemcpyDeviceToHost);
	cv::Mat image(src_height, src_width, CV_8UC3, pBuf);

	for (size_t i = 0; i < vec_result.size(); i++)
	{
		DogPoseResult poseResult = vec_result[i];
		std::cout << poseResult.x << std::endl;
		std::cout << poseResult.y << std::endl;
		std::cout << poseResult.width << std::endl;
		std::cout << poseResult.height << std::endl;
		std::cout << poseResult.confidence << std::endl;
		std::cout << poseResult.classId << std::endl;
		std::cout << poseResult.className << std::endl;

		cv::Rect position_boxe;
		position_boxe.x = poseResult.x;
		position_boxe.y = poseResult.y;
		position_boxe.width = poseResult.width;
		position_boxe.height = poseResult.height;
		cv::rectangle(image, position_boxe, cv::Scalar(0, 0, 255), 2, 8);
		cv::rectangle(image, cv::Point(position_boxe.x, position_boxe.y - 20), cv::Point(position_boxe.x, position_boxe.y), cv::Scalar(0, 255, 255), -1);
		cv::putText(image, poseResult.className, cv::Point(position_boxe.x, position_boxe.y - 10), cv::FONT_HERSHEY_SIMPLEX, .5, cv::Scalar(0, 0, 0));
	}

	cv::namedWindow("show", cv::WINDOW_NORMAL);
	cv::imwrite("result.jpg", image);
	cv::imshow("show", image);
	cv::waitKey(1);

	delete pBuf;
	pBuf = nullptr;
}

mutex m_mutex_show;
unsigned char *pShowData = nullptr;

void showFrame(AVFrame * gpuFrame) {
	std::lock_guard<std::mutex> l(m_mutex_show);

	cudaError_t cudaStatus = cudaSuccess;
	if (pShowData == nullptr)
	{
		cudaError_t cudaStatus = cudaMalloc((void **)&pShowData, 3 * gpuFrame->width * gpuFrame->height * sizeof(unsigned char));
	}
	
	cuda_common::setColorSpace(ITU709, 0);
	cudaStatus = cuda_common::CUDAToBGR((CUdeviceptr)gpuFrame->data[0], (CUdeviceptr)gpuFrame->data[1], gpuFrame->linesize[0], gpuFrame->linesize[1], pShowData, gpuFrame->width, gpuFrame->height);
	if (cudaStatus != cudaSuccess) {
		cout << "CUDAToBGR failed !!!" << endl;
		return;
	}

	int channel = 3;
	int width = gpuFrame->width;
	int height = gpuFrame->height;

	if (pShowData != nullptr && channel > 0 && width > 0 && height > 0) {
		std::vector<DogPoseResult> vec_result = poseDetector.detect(pShowData, width, height);

		showResult(pShowData, width, height, vec_result);

		//int nSize = channel * height * width;
		//unsigned char* cpu_data = new unsigned char[nSize];

		//cudaMemcpy(cpu_data, pShowData, nSize * sizeof(unsigned char), cudaMemcpyDeviceToHost);
		//cudaDeviceSynchronize();

		//cv::Mat img_(height, width, CV_8UC3, cpu_data);
		//imshow("show", img_);
		//waitKey(1);

		//delete[] cpu_data;
		//cpu_data = nullptr;

		

	}
}

/**
 * 注意： gpuFrame 在解码器设置的显卡上，后续操作要十分注意这一点，尤其是多线程情况
 * */
void postDecoded(const void * userPtr, AVFrame * gpuFrame){
    FFNvDecoder* decoder = (FFNvDecoder*)userPtr;
    if (decoder!= nullptr)
    {
        // cout << "decode name: " << decoder->getName() << endl;

        // const char* gpu_pixfmt = av_get_pix_fmt_name((AVPixelFormat)gpuFrame->format);
        // cout << "pixfmt: " << gpu_pixfmt << endl;
        // cout << "keyframe: " << gpuFrame->key_frame << " width: " << gpuFrame->width << " height: "<< gpuFrame->height << endl;
        // cout << "decode successed ✿✿ヽ(°▽°)ノ✿ " << endl;

		if (gpuFrame->format == AV_PIX_FMT_CUDA)
		{
			// cout << "gpuid = " << atoi(decoder->m_cfg.gpuid.c_str()) << endl;
			cudaSetDevice(atoi(decoder->m_cfg.gpuid.c_str()));

			//saveFrame(gpuFrame, decoder->getName());
			showFrame(gpuFrame);
		}
    }
}

long long start_time = 0;
long long end_time = 0;
bool count_flag = false;
int count_num = 0;
int count_std = 100;

static long long get_cur_time(){
    // 获取操作系统当前时间点（精确到微秒）
    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 int suming = 0;
unsigned char *pHwData = nullptr;

void postDecoded0(const void * userPtr, AVFrame * gpuFrame){
    // std::this_thread::sleep_for(std::chrono::milliseconds(30000));

    FFNvDecoder* decoder = (FFNvDecoder*)userPtr;
    if (decoder!= nullptr)
    {
        // cout << "decode name: " << decoder->getName() << endl;
        if (decoder->getName() == "dec")
        {
            if (! count_flag)
            {
                count_flag = true;
                count_num = 0;
                end_time = start_time = get_cur_time();
            }
            count_num++;
            suming ++ ;
            if (count_num >= count_std)
            {
                // end_time = get_cur_time();
                // long time_using = end_time - start_time;
                // double time_per_frame = double(time_using)/count_std ;
                // cout << count_std << "帧用时:" << time_using << "ms 每帧用时：" << time_per_frame << "ms" << endl;
                cout << "keyframe: " << gpuFrame->key_frame << " width: " << gpuFrame->width << " height: "<< gpuFrame->height << endl;

                count_flag = false;
            }
            cout << "帧数：" << suming << endl;
        }
    }
}

void decode_finished_cbk(const void* userPtr){
    cout << "decode_finish timestamp: " << get_cur_time() << endl;
}

// string test_uri = "rtmp://192.168.10.56:1935/objecteye/1";
// string test_uri = "/home/cmhu/data/output_800x480.mp4";
// string test_uri = "/home/cmhu/data/output_1920x1080.mp4";
// string test_uri = "rtsp://176.10.0.2:8554/stream";
// string test_uri = "/mnt/f/fiss/test_data/h265.mp4";
//string test_uri = "rtsp://176.10.0.4:8554/stream";
string test_uri = "f://data/5min.mp4";

void createDecode(int index){
    FFNvDecoderManager* pDecManager = FFNvDecoderManager::getInstance();
    MgrDecConfig config;
    config.name = "dec" + to_string(index);
    config.cfg.uri = test_uri;
    config.cfg.post_decoded_cbk = postDecoded;
    config.cfg.decode_finished_cbk = decode_finished_cbk;
    config.cfg.force_tcp = true;

    if (index % 2 == 0)
    {
        config.cfg.gpuid = "0";
    }
    else
    {
        config.cfg.gpuid = "0";
    }
    
    FFNvDecoder* decoder = pDecManager->createDecoder(config);
    if (!decoder)
    {
        return ;
    }
    pDecManager->setUserPtr(config.name, decoder);
    pDecManager->startDecodeByName(config.name);
}

void logFF(void *, int level, const char *fmt, va_list ap)
{
    vfprintf(stdout, fmt, ap);
}


int main(int argc, char* argv[]) {

    printf("start \n");
    //if (argc != 3) {
    //    fprintf(stderr, "./xxx uri gpu_id\n");
    //    return -1;
    //}

	//char* uri = argv[1];
	char* gpuid = "0";//argv[2];

    cout << av_version_info() << endl;

	poseDetector.init();

	//namedWindow("show", WINDOW_NORMAL);

    //evalQuality(uri, gpuid);


    // av_log_set_callback(&logFF);

    FFNvDecoderManager* pDecManager = FFNvDecoderManager::getInstance();

    // int count = 99;
    // for (size_t i = 0; i < count ; i++)
    // {
    //     createDecode(i);
    // }

    MgrDecConfig config;
    config.name = "dec";
    config.cfg.uri = test_uri;
    config.cfg.post_decoded_cbk = postDecoded;
    config.cfg.decode_finished_cbk = decode_finished_cbk;
    config.cfg.force_tcp = true;
    config.cfg.gpuid = "0";
    FFNvDecoder* dec2 = pDecManager->createDecoder(config);
    if (!dec2)
    {
        return 1;
    }
    pDecManager->setUserPtr(config.name, dec2);
    // pDecManager->setDecKeyframe(config.name, true);
    pDecManager->startDecodeByName(config.name);

    int w,h;
    pDecManager->getResolution(config.name, w,h);
    printf( "%s : %dx%d\n", config.name.c_str() , w,h );

    //thread* m_thread = new thread([](void* arg)
    //    {
    //        while (true)
    //        {
    //            std::this_thread::sleep_for(std::chrono::milliseconds(5000));
    //            FFNvDecoderManager* pDecManager = FFNvDecoderManager::getInstance();
    //            int count = pDecManager->count();
    //            cout << "当前运行路数： " << pDecManager->count() << endl;
    //        }

    //        return (void*)0;
    //    }
    //, nullptr);


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

    cout << "总共帧数：" << sum << endl;

    pDecManager->closeAllDecoder();
}