Hu Chunming / FFNvDecoder

//#include "LOG_manager.h"
#include <iostream>
#include "FFGB28181Decoder.h"



extern "C" {
#include "libavutil/avstring.h"
#include "libavformat/avformat.h"
#include "libswscale/swscale.h"
}

#include"RTPTcpReceiver.h"
#include"RTPUdpReceiver.h"

#include <cuda_runtime.h>

#include "common_header.h"

#include "../nvdec/FFCuContextManager.h"
#include "../nvdec/GpuRgbMemory.hpp"
#include "../nvdec/cuda_kernels.h"

#define ECLOSED 0
#define ECLOSING 1
#define ERUNNING 2
#define EPAUSE 3

static void RTP_Stream_CallBack(void* userdata, int videoType, char* data, int len, int isKey, uint64_t pts, uint64_t localPts)
{
    FFGB28181Decoder* decoder = (FFGB28181Decoder*)userdata;
    decoder->stream_callback(videoType, data, len, isKey, pts, localPts);
}

static void RTP_Stream_End_CallBack(void* userdata)
{
    FFGB28181Decoder* decoder = (FFGB28181Decoder*)userdata;
    decoder->stream_end_callback();
}

FFGB28181Decoder::FFGB28181Decoder() {
	m_frameSkip = 1;
    m_port = -1;
    m_dec_keyframe = false;
    m_post_decode_thread = 0;
}

FFGB28181Decoder::~FFGB28181Decoder()
{
    close();
	
    if (m_pAVCodecCtx) {
        avcodec_close(m_pAVCodecCtx);
        avcodec_free_context(&m_pAVCodecCtx);
    }

    m_dec_keyframe = false;
    
    LOG_INFO("destroy OK--{}", m_dec_name);
}

void FFGB28181Decoder::close(){
    if (m_status == ECLOSED || m_status == ECLOSING) return ;
    
    m_status = ECLOSING;

    LOG_DEBUG("real decode thread exit success 1--{}", m_dec_name);

    if(nullptr != m_rtpPtr){
        if (m_rtpPtr->IsOpened()) {
            m_rtpPtr->Close();
            LOG_DEBUG("real decode thread exit success 2--{}", m_dec_name);
        }

        delete m_rtpPtr;
        m_rtpPtr = nullptr;
    }

    if (gpu_options) av_dict_free(&gpu_options);

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

    while(mFrameQueue.size() > 0){
		AVFrame * gpuFrame = mFrameQueue.front();
		av_frame_free(&gpuFrame); 
		mFrameQueue.pop();
	}

    m_status = ECLOSED;

    LOG_INFO("解码器关闭成功 --{}", m_dec_name);
}

bool FFGB28181Decoder::init(FFDecConfig& cfg){
    if(cfg.force_tcp){
        m_rtpPtr = new RTPTcpReceiver();
    }else{
        m_rtpPtr = new RTPUdpReceiver();
    }
    if(nullptr == m_rtpPtr){
        return false;
    }

    m_dec_name = cfg.uri;
	m_frameSkip = cfg.skip_frame;
    if (m_frameSkip < 1) m_frameSkip = 1;

    m_gpuid = atoi(cfg.gpuid.c_str());

    m_rtpPtr->SetDeviceID(m_dec_name);

    if(cfg.request_stream_cbk == nullptr){
        LOG_INFO("request_stream_cbk 为 nullptr -- {}", m_dec_name);
        return false;
    }

    post_decoded_cbk = cfg.post_decoded_cbk;
    decode_finished_cbk = cfg.decode_finished_cbk;
    m_rtpPtr->SetRequestStreamCallback(cfg.request_stream_cbk);

    m_port = cfg.port;

    m_cfg = cfg;

    LOG_INFO("init - {} : ", m_dec_name, m_port);

    return true;
}

bool FFGB28181Decoder::start() {

    m_status = ERUNNING;

	m_rtpPtr->SetOutputCallback(RTP_Stream_CallBack, this);
	m_rtpPtr->SetVodEndCallback(RTP_Stream_End_CallBack, this);

    LOG_INFO("start - {} {}: ", m_dec_name, m_port);

    bool bRet = m_rtpPtr->Open((uint16_t)m_port);
    if(bRet){
        pthread_create(&m_post_decode_thread,0,
            [](void* arg)
            {
                FFGB28181Decoder* a=(FFGB28181Decoder*)arg;
                a->post_decode_thread();
                return (void*)0;
            }
        ,this);
    }
    return bRet;
}

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

static AVPixelFormat get_hw_format(AVCodecContext *avctx, const AVPixelFormat *pix_fmts){
	return AV_PIX_FMT_CUDA;
}

void FFGB28181Decoder::stream_callback(int videoType, char* data, int len, int isKey, uint64_t pts, uint64_t localPts) {
    if (m_status == EPAUSE) return;

    // 若设置为关键帧解码，非关键帧数据直接返回
    if(m_dec_keyframe && !isKey) return;
    
    if (len <= 0) {
        if (data == nullptr && pts == -1) {
            LOG_INFO("frame callback EOF!");
			post_decoded_cbk(m_postDecArg, nullptr);
            return ;
        }
        LOG_INFO("frame data is zero --{}", m_dec_name);
        return;
    }

    AVPacket* pkt = av_packet_alloc();
	av_init_packet(pkt);

    pkt->size = len;
    pkt->data = (uint8_t*)data;

	if (m_pAVCodecCtx == nullptr) {
        LOG_INFO("frame data is zero --{}", m_dec_name);
		if (VIDEO_TYPE_H264 == videoType) {
            if (m_gpuid >= 0){
                m_pAVCodec = avcodec_find_decoder_by_name("h264_cuvid");
            }
            else{
				m_pAVCodec = avcodec_find_decoder(AV_CODEC_ID_H264);
			}
			LOG_INFO("m_avCodecName is H264");		
		}
		else if (VIDEO_TYPE_H265 == videoType)
		{
            if (m_gpuid >= 0){
                m_pAVCodec = avcodec_find_decoder_by_name("hevc_cuvid");
            }
            else{
                m_pAVCodec = avcodec_find_decoder(AV_CODEC_ID_H265);
            }
			LOG_INFO("m_avCodecName is H265");
		}
		else{
			LOG_INFO("m_avCodecName is unknown, videoType is {}", videoType);
		}

		if (!m_pAVCodec)
		{
			LOG_ERROR("frameCallback frame decode error, ERROR_DECODER_NOT_FOUND");
			return;
		}

        m_pAVCodecCtx = avcodec_alloc_context3(m_pAVCodec);

        if (m_gpuid >= 0) {
            char gpui[8] = { 0 };
            sprintf(gpui, "%d", m_gpuid);
            av_dict_set(&gpu_options, "gpu", gpui, 0);

            m_pAVCodecCtx->get_format = get_hw_format;

            FFCuContextManager* pCtxMgr = FFCuContextManager::getInstance();
            m_pAVCodecCtx->hw_device_ctx = av_buffer_ref(pCtxMgr->getCuCtx(gpui));
            if (nullptr == m_pAVCodecCtx->hw_device_ctx){
                // TODO 这里应该抛出错误
                return ;
            }
        }

        if (avcodec_open2(m_pAVCodecCtx, m_pAVCodec, &gpu_options) < 0)
			return;
	}

    //开始解码
    int ret = avcodec_send_packet(m_pAVCodecCtx, pkt);
    if (ret < 0) {
        //send_exception(RunMessageType::E2002, e_msg);
        LOG_ERROR("Real stream视频解码失败，请检查视频设备{}: avcodec_send_packet failed. ret={}", m_dec_name, ret);
        av_packet_free(&pkt);
        pkt = nullptr;
        return;
    }
 
    if (frameW < 1) {
        frameW = m_pAVCodecCtx->width;
        frameH = m_pAVCodecCtx->height;
        if (frameW <= 0 || frameH <= 0) {
            LOG_ERROR("[{}] frame W or H is error! ({},{})", m_dec_name, frameW, frameH);
            av_packet_free(&pkt);
            pkt = nullptr;
            return;
        }	
    }
	// m_fps = m_pAVCodecCtx->pkt_timebase.den == 0 ? 25.0 : av_q2d(m_pAVCodecCtx->pkt_timebase);
    m_fps = av_q2d(m_pAVCodecCtx->framerate);
	// LOG_DEBUG("frameW {}--frameH {}", frameW, frameH);

    AVFrame* gpuFrame = av_frame_alloc();
    ret = avcodec_receive_frame(m_pAVCodecCtx, gpuFrame);
    if ((ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) || ret < 0){
        LOG_ERROR("{} - Failed to receive frame: {}", m_dec_name, ret);
        av_packet_free(&pkt);
        pkt = nullptr;
        av_frame_free(&gpuFrame);
        gpuFrame = nullptr;
        return;
    }

    av_packet_free(&pkt);
    pkt = nullptr;
    
    if (gpuFrame->width != frameW || gpuFrame->height != frameH){
        LOG_INFO("AVFrame is inconsistent: width is {}, height is {}; original frameW is {}, frameH is {}--{}", gpuFrame->width, gpuFrame->height, frameW, frameH , m_dec_name);
        av_frame_free(&gpuFrame);
        gpuFrame = nullptr;
        return;
    }

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

void FFGB28181Decoder::post_decode_thread(){
	
	int index = 0;
	while (isRunning())
	{
		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 (m_frameSkip == 1 || index % m_frameSkip == 0){
				post_decoded_cbk(m_postDecArg, convert2bgr(gpuFrame));
			}

			av_frame_free(&gpuFrame); 
            gpuFrame = nullptr;

			index++;
            if(index >= 100000){
                index = 0;
            }
		}
	}

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

void FFGB28181Decoder::stream_end_callback()
{
    LOG_INFO("send_video_eof--{}", m_dec_name);

	decode_finished_cbk(m_finishedDecArg);

    return;
}

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

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

void FFGB28181Decoder::pause() {
    m_status = EPAUSE;
    LOG_INFO("pause --{}", m_dec_name);
}

void FFGB28181Decoder::resume() {
    m_status = ERUNNING;
    LOG_INFO("resume --{}", m_dec_name);
}

bool FFGB28181Decoder::isRunning(){
    if (m_status == ECLOSED || m_status == ECLOSING){
        return false;
    }
    return true;
}

bool FFGB28181Decoder::isFinished(){
    if (m_status == ECLOSED || m_status == ECLOSING){
        return true;
    }
    return false;
}

bool FFGB28181Decoder::isPausing(){
    if (m_status == EPAUSE){
        return true;
    }
    return false;
}

bool FFGB28181Decoder::getResolution( int &width, int &height ){
    width = frameW;
    height = frameH;
    return true;
}

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

bool FFGB28181Decoder::isSurport(FFDecConfig& cfg){
    // 由于是否支持需要在拿到数据后才能断定，无法事先判断，所以这个地方默认返回true
    return true;
}

int FFGB28181Decoder::getCachedQueueLength(){
    return m_rtpPtr->GetPsFrameListSize();
}

DeviceRgbMemory* FFGB28181Decoder::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* FFGB28181Decoder::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;
}