#include "FFNvDecoder.h"

#include <chrono>
#include <thread>
#include <fstream>

#include <chrono>

#include "FFCuContextManager.h"

#include "common_header.h"

#include "GpuRgbMemory.hpp"
#include "cuda_kernels.h"

using namespace std;

// 参考博客： https://blog.csdn.net/qq_40116098/article/details/120704340

static AVPixelFormat get_hw_format(AVCodecContext *avctx, const AVPixelFormat *pix_fmts)
{
	FFNvDecoder* _this = (FFNvDecoder*)avctx->opaque;

	const AVPixelFormat *p;

	for (p = pix_fmts; *p != -1; p++) {
		if (*p == _this->getHwPixFmt())
			return *p;
	}

	LOG_ERROR("Failed to get HW surface format");
	return AV_PIX_FMT_NONE;
}

FFNvDecoder::FFNvDecoder()
{
	// 初始化解码对象
	fmt_ctx = nullptr;
	avctx = nullptr;
	m_bRunning = false;

	stream = nullptr;
    stream_index = -1;
    hw_pix_fmt = AV_PIX_FMT_NONE;
    m_dec_name = "";

	m_bPause = false;
	m_bReal = true;

	m_decode_thread = 0;
	m_post_decode_thread = 0;

	m_bFinished = false;
	m_dec_keyframe = false;
	m_fps = 0.0;
}

FFNvDecoder::~FFNvDecoder()
{
	m_dec_keyframe = false;
}

bool FFNvDecoder::init(FFDecConfig& cfg)
{
	m_cfg = cfg;
	m_dec_name = cfg.dec_name;

	fstream infile(cfg.uri);
	if (infile.is_open()){
		m_bReal = false;
		infile.close();
	}else {
		m_bReal = true;
	}

	post_decoded_cbk = cfg.post_decoded_cbk;
    decode_finished_cbk = cfg.decode_finished_cbk;

	return init(cfg.uri.c_str(), cfg.gpuid.c_str(),cfg.force_tcp);
}

bool FFNvDecoder::init(const char* uri, const char* gpuid, bool force_tcp)
{
	// av_log_set_level(AV_LOG_DEBUG);

	avformat_network_init();

	// 打开输入视频文件
	AVDictionary *options = nullptr;
	av_dict_set( &options, "bufsize", "655360", 0 );
	av_dict_set( &options, "rtsp_transport", force_tcp ? "tcp" : "udp", 0 );
	// av_dict_set( &options, "listen_timeout", "30", 0 ); // 单位为s
	av_dict_set( &options, "stimeout", "30000000", 0 ); // 单位为 百万分之一秒
	
	fmt_ctx = avformat_alloc_context();
	const char* input_file = uri;
	if (avformat_open_input(&fmt_ctx, input_file, nullptr, &options) != 0) {
		LOG_ERROR("Cannot open input file:{}",input_file);
		return false;
	}

	// 查找流信息
	if (avformat_find_stream_info(fmt_ctx, nullptr) < 0) {
		LOG_ERROR("Cannot find input stream information");
		return false;
	}

	// 查找视频流信息
	AVCodec *decoder = nullptr;
	stream_index = av_find_best_stream(fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, &decoder, 0);
	if (stream_index < 0) {
		LOG_ERROR("Cannot find a video stream in the input file");
		return false;
	}

	string cuvid_dec_name = string(decoder->name) + "_cuvid";
	AVCodec *vcodec = avcodec_find_decoder_by_name(cuvid_dec_name.c_str());
	if (!(avctx = avcodec_alloc_context3(vcodec)))
		return (bool)AVERROR(ENOMEM);

	// 得到视频流对象
	stream = fmt_ctx->streams[stream_index];
	if (avcodec_parameters_to_context(avctx, stream->codecpar) < 0)
		return false;

	m_fps = av_q2d(stream ->avg_frame_rate);

	avctx->opaque = this;
	// 设置解码器管理器的像素格式回调函数
	avctx->get_format = get_hw_format;

	hw_pix_fmt = AV_PIX_FMT_CUDA;

	FFCuContextManager* pCtxMgr = FFCuContextManager::getInstance();

	AVBufferRef *hw_device_ctx = pCtxMgr->getCuCtx(gpuid);
	if(nullptr == hw_device_ctx){
		av_log(nullptr, AV_LOG_ERROR, "create CUDA context failed ! \n");
		return false;
	}
	avctx->hw_device_ctx = av_buffer_ref(hw_device_ctx);
	if (nullptr == avctx->hw_device_ctx)
	{
		return false;
	}

	// 打开解码器流
	AVDictionary *op = nullptr;
	av_dict_set( &op, "gpu", gpuid, 0 );
	// av_dict_set( &op, "surfaces", "5", 0 );
	if (avcodec_open2(avctx, vcodec, &op) < 0) {
		LOG_ERROR("Failed to open codec for stream");
		return false;
	}
	
	return true;
}

bool FFNvDecoder::isSurport(FFDecConfig& cfg)
{
	bool bRet = init(cfg);
    decode_finished();
    return bRet;
}

bool FFNvDecoder::start(){

	m_bRunning = true;

	pthread_create(&m_decode_thread,0,
        [](void* arg)
        {
            FFNvDecoder* a=(FFNvDecoder*)arg;
            a->decode_thread();
            return (void*)0;
        }
    ,this);

	return true;
}

void FFNvDecoder::decode_thread()
{
	AVPacket* pkt ;
	pkt = av_packet_alloc();
	av_init_packet( pkt );

	pthread_create(&m_post_decode_thread,0,
        [](void* arg)
        {
            FFNvDecoder* a=(FFNvDecoder*)arg;
            a->post_decode_thread();
            return (void*)0;
        }
    ,this);

	// long start_time = UtilTools::get_cur_time_ms();

	while (m_bRunning)
	{
		if (!m_bReal)
		{
			if (m_bPause)
			{
				std::this_thread::sleep_for(std::chrono::milliseconds(3));
				continue;
			}
		}
		
		int result = av_read_frame(fmt_ctx, pkt);
		if (result == AVERROR_EOF || result < 0)
		{
			LOG_ERROR("Failed to read frame!");
			break;
		}

		if (m_dec_keyframe && !(pkt->flags & AV_PKT_FLAG_KEY)) {
			av_packet_unref(pkt);
			continue;
		}

		if (stream_index == pkt->stream_index){
			result = avcodec_send_packet(avctx, pkt);
			if (result < 0){
				av_packet_unref(pkt);
				LOG_ERROR("{} - Failed to send pkt: {}", m_dec_name, result);
				continue;
			}

			AVFrame* gpuFrame = av_frame_alloc();
			result = avcodec_receive_frame(avctx, gpuFrame);
			if ((result == AVERROR(EAGAIN) || result == AVERROR_EOF) || result < 0){
				LOG_ERROR("{} - Failed to receive frame: {}", m_dec_name, result);
				av_frame_free(&gpuFrame); 
				av_packet_unref(pkt);
				continue;
			}
			av_packet_unref(pkt);

			if (m_bReal){
				if (m_bPause){
					av_frame_free(&gpuFrame); 
					std::this_thread::sleep_for(std::chrono::milliseconds(3));
					continue;
				}
			}

			if(gpuFrame != nullptr){
				m_queue_mutex.lock();
				if(mFrameQueue.size() <= 10){
					mFrameQueue.push(gpuFrame);
				}else{
					av_frame_free(&gpuFrame); 
				}
				m_queue_mutex.unlock();
			}
		}
		av_packet_unref(pkt);
	}

	m_bRunning = false;
	av_packet_free(&pkt);

	// long end_time = UtilTools::get_cur_time_ms();
	// cout << "解码用时：" << end_time - start_time << endl;

	if (m_post_decode_thread != 0)
	{
		pthread_join(m_post_decode_thread,0);
	}

	decode_finished_cbk(m_finishedDecArg);

	decode_finished();

	// 清空队列
	while(mFrameQueue.size() > 0){
		AVFrame * gpuFrame = mFrameQueue.front();
		av_frame_free(&gpuFrame); 
		mFrameQueue.pop();
	}

	LOG_INFO("{} - decode thread exited.", m_dec_name);
}

void FFNvDecoder::decode_finished(){
	if (avctx)
	{
		avcodec_free_context(&avctx);
	}
	
	if (fmt_ctx)
	{
		avformat_close_input(&fmt_ctx);
	}

	m_bFinished = true;
	m_dec_keyframe = false;
}

void FFNvDecoder::post_decode_thread(){
	int skip_frame = m_cfg.skip_frame;
	if (skip_frame <= 0){
		skip_frame = 1;
	}
	
	int index = 0;
	while (m_bRunning)
	{
		if(mFrameQueue.size() > 0){
			std::lock_guard<std::mutex> l(m_snapshot_mutex);
			// 取队头数据
			m_queue_mutex.lock();
			AVFrame * gpuFrame = mFrameQueue.front();
			mFrameQueue.pop();
			m_queue_mutex.unlock();
			// 跳帧
			if (skip_frame == 1 || index % skip_frame == 0){
				post_decoded_cbk(m_postDecArg, convert2bgr(gpuFrame));
				index = 0;
			}

			av_frame_free(&gpuFrame); 

			index++;
		}
	}

	LOG_INFO("post decode thread exited.");
}

void FFNvDecoder::close(){
	m_bRunning=false;
	if(m_decode_thread != 0){
		pthread_join(m_decode_thread,0);
	}
	m_dec_keyframe = false;
}

AVPixelFormat FFNvDecoder::getHwPixFmt(){
	return hw_pix_fmt;
}

bool FFNvDecoder::isRunning(){
	return m_bRunning;
}

bool FFNvDecoder::isFinished(){
	return m_bFinished;
}

bool FFNvDecoder::isPausing(){
	return m_bPause;
}

bool FFNvDecoder::getResolution( int &width, int &height ){
	if (avctx != nullptr)
	{
		width = avctx->width;
		height = avctx->height;
		return true;
	}
	
	return false;
}

void FFNvDecoder::pause(){
	m_bPause = true;
}

void FFNvDecoder::resume(){
	m_bPause = false;
}

void FFNvDecoder::setDecKeyframe(bool bKeyframe)
{
	m_dec_keyframe = bKeyframe;
}

int FFNvDecoder::getCachedQueueLength(){
	m_queue_mutex.lock();
	int queue_size = mFrameQueue.size(); 
	m_queue_mutex.lock();
	return queue_size;
}

float FFNvDecoder::fps(){
	return m_fps;
}

void FFNvDecoder::setPostDecArg(const void* postDecArg){
	m_postDecArg = postDecArg;
}

void FFNvDecoder::setFinishedDecArg(const void* finishedDecArg){
	m_finishedDecArg = finishedDecArg;
}

DeviceRgbMemory* FFNvDecoder::convert2bgr(AVFrame * gpuFrame){
	if (gpuFrame != nullptr && gpuFrame->format == AV_PIX_FMT_CUDA ){
		LOG_DEBUG("decode task: gpuid: {}  width: {} height: {}", m_cfg.gpuid, gpuFrame->width, gpuFrame->height);
		GpuRgbMemory* gpuMem = new GpuRgbMemory(3, gpuFrame->width, gpuFrame->height, getName(), m_cfg.gpuid, false, true);

		do{
			if (gpuMem->getMem() == nullptr){
				LOG_ERROR("new GpuRgbMemory failed !!!");
				break;
			}
			
			cudaSetDevice(atoi(m_cfg.gpuid.c_str()));
			cuda_common::setColorSpace( ITU_709, 0 );
			cudaError_t cudaStatus = cuda_common::CUDAToBGR((CUdeviceptr)gpuFrame->data[0],(CUdeviceptr)gpuFrame->data[1], gpuFrame->linesize[0], gpuFrame->linesize[1], gpuMem->getMem(), gpuFrame->width, gpuFrame->height);
			cudaDeviceSynchronize();
			if (cudaStatus != cudaSuccess) {
				LOG_ERROR("CUDAToBGR failed failed !!!");
				break;
			}

			return gpuMem;
		}while(0);

		delete gpuMem;
		gpuMem = nullptr;
	}

	return nullptr;
}

FFImgInfo* FFNvDecoder::snapshot(){

	// 锁住停止队列消耗
	std::lock_guard<std::mutex> l(m_snapshot_mutex);

	AVFrame * gpuFrame = nullptr;

	bool bFirst = true;
	while(true){
		m_queue_mutex.lock();
		if(mFrameQueue.size() <= 0){
			m_queue_mutex.unlock();
			if(bFirst){
				std::this_thread::sleep_for(std::chrono::milliseconds(100));
				bFirst = false;
				continue;
			}else{
				// 再进来说明前面已经等了 100 ms
				// 100 ms都没有等到解码数据，则退出
				return nullptr;
			}
		}

		// 队列中数据大于1 
		gpuFrame = mFrameQueue.front();
		m_queue_mutex.unlock();
		break;
	}

	if (gpuFrame != nullptr && gpuFrame->format == AV_PIX_FMT_CUDA ){
		LOG_DEBUG("decode task: gpuid: {}  width: {} height: {}", m_cfg.gpuid, gpuFrame->width, gpuFrame->height);
		GpuRgbMemory* gpuMem = new GpuRgbMemory(3, gpuFrame->width, gpuFrame->height, getName(), m_cfg.gpuid, false, true);

		if (gpuMem->getMem() == nullptr){
			LOG_ERROR("new GpuRgbMemory failed !!!");
			return nullptr;
		}
		
		cudaSetDevice(atoi(m_cfg.gpuid.c_str()));
		cuda_common::setColorSpace( ITU_709, 0 );
		cudaError_t cudaStatus = cuda_common::CUDAToBGR((CUdeviceptr)gpuFrame->data[0],(CUdeviceptr)gpuFrame->data[1], gpuFrame->linesize[0], gpuFrame->linesize[1], gpuMem->getMem(), gpuFrame->width, gpuFrame->height);
		cudaDeviceSynchronize();
		if (cudaStatus != cudaSuccess) {
			LOG_ERROR("CUDAToBGR failed failed !!!");
			return nullptr;
		}

		unsigned char * pHwRgb = gpuMem->getMem();
		int channel = gpuMem->getChannel();
		int width = gpuMem->getWidth();
		int height = gpuMem->getHeight();

		if (pHwRgb != nullptr && channel > 0 && width > 0 && height > 0){
			int nSize = channel * height * width;

			LOG_INFO("channel:{} height:{} width:{}", channel, height, width);
			// unsigned char* cpu_data = new unsigned char[nSize];

            unsigned char* cpu_data = (unsigned char *)av_malloc(nSize * sizeof(unsigned char));

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

			delete gpuMem;
			gpuMem = nullptr;

			FFImgInfo* imgInfo = new FFImgInfo();
			imgInfo->dec_name = m_dec_name;
			imgInfo->pData = cpu_data;
			imgInfo->height = height;
			imgInfo->width = width;
			imgInfo->timestamp = UtilTools::get_cur_time_ms();
			imgInfo->index = m_index;

			m_index++;

			return imgInfo;
		}

		delete gpuMem;
		gpuMem = nullptr;
	}

	return nullptr;
}