#include"check_tool.h"

#include "FFCuContextManager.h"

extern "C"
{
	#include <libavcodec/avcodec.h> 
	#include <libavdevice/avdevice.h> 
	#include <libavformat/avformat.h> 
	#include <libavfilter/avfilter.h> 
	#include <libavutil/avutil.h> 
    #include <libavutil/pixdesc.h> 
	#include <libswscale/swscale.h>
    #include <libavutil/imgutils.h>
}

#include <chrono>

static AVPixelFormat get_hw_format(AVCodecContext *avctx, const AVPixelFormat *pix_fmts){

	const AVPixelFormat *p;

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

	av_log(nullptr, AV_LOG_ERROR, "Failed to get HW surface format. \n");
	return AV_PIX_FMT_NONE;
}

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());
    // (微秒精度的)时间点 => (微秒精度的)时间戳
    time_t currentTime = tpMicro.time_since_epoch().count();

    return (long long )currentTime;
}

void evalQuality(const string& uri, const char* gpuid){
	AVFormatContext* ifmt_ctx = nullptr;
	AVCodecContext* codec_ctx = nullptr;
	AVCodec *decoder = nullptr;
	AVCodec* codec = nullptr;
	AVPacket* pkt = nullptr;
	int video_index = -1;
	AVStream* st = nullptr;
	SwsContext *img_convert_ctx = nullptr;
	uint8_t *buffer = nullptr;
    int numBytes = 0;

	long start_time = 0;
	long end_time = 0;
	long s_time = 0;
	long e_time = 0;
	long long sum = 0;
	double avg_time = 0.0;

	int sep = 0;

	string cuvid_dec_name = "" ;
	FFCuContextManager* pCtxMgr = FFCuContextManager::getInstance();
	AVDictionary *op = nullptr;
	AVBufferRef *hw_device_ctx = nullptr;
 

	avformat_network_init();

	// 打开输入视频文件
	AVDictionary *options = nullptr;
	av_dict_set( &options, "bufsize", "655360", 0 );
	av_dict_set( &options, "rtsp_transport", "tcp", 0 );
	// av_dict_set( &options, "listen_timeout", "30", 0 ); // 单位为s
	av_dict_set( &options, "stimeout", "30000000", 0 ); // 单位为 百万分之一秒
	
	///打开输入的流
	ifmt_ctx = avformat_alloc_context();
	int ret = avformat_open_input(&ifmt_ctx, uri.c_str(), nullptr, &options);
	if (ret != 0){
		av_log(nullptr, AV_LOG_ERROR, "Couldn't open input stream ! \n");
		goto end_flag ;
	}
 
	//查找流信息
	if (avformat_find_stream_info(ifmt_ctx, nullptr) < 0){
		av_log(nullptr, AV_LOG_ERROR, "Couldn't find stream information ! \n");
		goto end_flag ;
	}

	// 查找视频流信息
	video_index = av_find_best_stream(ifmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, &decoder, 0);
	if (video_index < 0) {
		av_log(nullptr, AV_LOG_ERROR, "Cannot find a video stream in the input file ! \n");
		goto end_flag ;
	}

	cuvid_dec_name = string(decoder->name) + "_cuvid";
	codec = avcodec_find_decoder_by_name(cuvid_dec_name.c_str());
	if (!codec){
		av_log(nullptr, AV_LOG_ERROR, "Codec not found ! \n");
		goto end_flag ;
	}
    
    //申请AVCodecContext
    codec_ctx = avcodec_alloc_context3(codec);
    if (!codec_ctx){
        goto end_flag ;
    }
 
	st = ifmt_ctx->streams[video_index];
	if(avcodec_parameters_to_context(codec_ctx, st->codecpar) < 0){
		goto end_flag ;
	}

	{
		av_log(nullptr, AV_LOG_INFO, "path: %s \n", ifmt_ctx->url);
		av_log(nullptr, AV_LOG_INFO, "format: %s \n", ifmt_ctx->iformat->name);
		const char* profile = avcodec_profile_name(codec_ctx->codec_id, codec_ctx->profile);
		av_log(nullptr, AV_LOG_INFO, "codec: %s(%s) \n", decoder->name, profile);
		const char* pix_fmt_name = av_get_pix_fmt_name(codec_ctx->pix_fmt);
		av_log(nullptr, AV_LOG_INFO, "pix_fmt_name: %s \n", pix_fmt_name);
		av_log(nullptr, AV_LOG_INFO, "width x height: %dX%d \n", st->codecpar->width, st->codecpar->height);

		av_log(NULL, AV_LOG_INFO, "frame_rate: %1.0f ", av_q2d(st->r_frame_rate));
		av_log(NULL, AV_LOG_INFO, "  -->  reference PPS: %f ms\n", 1/av_q2d(st->r_frame_rate) * 1000);
		
	}

	// 设置解码器管理器的像素格式回调函数
	codec_ctx->get_format = get_hw_format;
	
	hw_device_ctx = pCtxMgr->getCuCtx(gpuid);
	if(nullptr == hw_device_ctx){
		av_log(nullptr, AV_LOG_ERROR, "create CUDA context failed ! \n");
		goto end_flag ;
	}
	codec_ctx->hw_device_ctx = av_buffer_ref(hw_device_ctx);
	if (nullptr == codec_ctx->hw_device_ctx){
		goto end_flag ;
	}

	// 打开解码器流
	av_dict_set( &op, "gpu", gpuid, 0 );
	if (avcodec_open2(codec_ctx, codec, &op) < 0) {
		av_log(nullptr, AV_LOG_ERROR, "Failed to open codec for stream ! \n");
		goto end_flag ;
	}
	
	s_time = get_cur_time();
	start_time = get_cur_time();
	pkt = av_packet_alloc();
	while (av_read_frame(ifmt_ctx, pkt) >= 0){
		end_time = get_cur_time();
		sum ++ ;
		sep ++ ;
		if(end_time - start_time > 1000){
			avg_time = double(end_time - start_time) / sep;
			start_time = get_cur_time();
			sep = 0;
			av_log(nullptr, AV_LOG_INFO, "PPS: %f ms\n", avg_time);
		}
		av_packet_unref(pkt);
	}

	e_time = get_cur_time();
	avg_time = double(e_time - s_time) / sum;
	av_log(nullptr, AV_LOG_INFO, "TOOTAL PPS: %f ms\n", avg_time);

end_flag:
	if (codec_ctx != nullptr){
		avcodec_close(codec_ctx);
		avcodec_free_context(&codec_ctx);
	}
	
	if (ifmt_ctx != nullptr){
		avformat_close_input(&ifmt_ctx);
	}

	if (pkt != nullptr){
		av_packet_free(&pkt);
	}
}