Hu Chunming / Project_VideoStructure

#include "DxDecoderWrap.h"

#include <thread>
#include <chrono>
#include "logger.hpp"

void decoded_cbk(const void * userPtr, GPUFrame * gpuFrame){
	DxDecoderWrap* _this = (DxDecoderWrap*)userPtr;
	if(nullptr != _this){
		_this->post_decode_thread(gpuFrame);
	}
}

void decode_finished_cbk(const void * userPtr){
	DxDecoderWrap* _this = (DxDecoderWrap*)userPtr;
	if(nullptr != _this){
		_this->decode_finished_thread();
	}
}

DxDecoderWrap::DxDecoderWrap( const DxConfig * cfg )
{
	m_bClose = false;
	m_pDec = nullptr;
	m_pHwData = nullptr;

	m_cfg.post_decoded_cbk = decoded_cbk;
    m_cfg.decode_finished_cbk = decode_finished_cbk;
    m_cfg.force_tcp = cfg->forceTcp;  // rtsp用tcp
    m_cfg.gpuid = std::to_string(cfg->devId);

	m_name = cfg->name;
	m_decMode = cfg->decMode;

	CUdevice cuDevice;
	cuDeviceGet(&cuDevice, cfg->devId);
	cuCtxCreate(&m_cuda_cxt, 0, cuDevice);

	return;
}


DxDecoderWrap::~DxDecoderWrap()
{
	DxCloseDecoder();

	return;
}


int DxDecoderWrap::DxOpenDecoder( const char * uri, unsigned int skip )
{
	int ret = 0;

	m_pDec = new FFNvDecoder();
	if ( nullptr == m_pDec ) {
		return -1;
	}

	m_cfg.uri = uri;
	m_cfg.skip_frame = skip;

	bool bRet= m_pDec->init(m_cfg);
    if (bRet)
    {
        m_pDec->setName(m_name) ;
		m_pDec->m_postDecArg = this;
		m_pDec->m_finishedDecArg = this;
		if(1 == m_decMode) {
			m_pDec->setDecKeyframe(true);
		}
		m_pDec->start();
        LOG_INFO("[{}][{}]- 解码器初始化成功", m_name.c_str(), uri);
        return 0;
    }
    
    // 创建失败，关闭解码器
    m_pDec->close();
    delete m_pDec;
	m_pDec = nullptr;

    LOG_ERROR("[{}][{}]- 解码器初始化失败！", m_name.c_str(), uri);
    return -1;
}


int DxDecoderWrap::DxCloseDecoder()
{
	m_bClose = true;

	if(m_pDec) {
		m_pDec->close();
		delete m_pDec;
		m_pDec = nullptr;
	}

	m_queue_frames_mutex.lock();
	while (m_queue_frames.size() > 0) {
		GPUFrame * decodedFrame = m_queue_frames.front();
		m_queue_frames.pop();
		delete decodedFrame;
		decodedFrame = nullptr;
	}
	m_queue_frames_mutex.unlock();

	if(m_pHwData != nullptr){
		cudaFree(m_pHwData);
	}

	cuCtxPopCurrent(nullptr);
	cuCtxDestroy(m_cuda_cxt);

	return 0;
}


int DxDecoderWrap::DxGetResolution( int &width, int &height )
{
	if(m_pDec) {
		m_pDec->getResolution(width, height);
		return 0;
	}

	return -1;
}

int DxDecoderWrap::DxGetFrameCount()
{
	if(m_pDec != nullptr){
		return m_pDec->GetFrameCount();
	}
	return -1;
}


int DxDecoderWrap::PauseDecoder()
{
	if(m_pDec) {
		m_pDec->pause();
	}
	return 0;
}


int DxDecoderWrap::ResumeDecoder()
{
	if(m_pDec) {
		m_pDec->resume();
	}

	return 0;
}


bool DxDecoderWrap::DxDecoderIsRun() const
{
	if(m_pDec) {
		return m_pDec->isRunning();
	}

	return false;
}


bool DxDecoderWrap::DxFrameIsEmpty()
{
	m_queue_frames_mutex.lock();
	int count = m_queue_frames.size();
	m_queue_frames_mutex.unlock();

	return 0 == count;
}


int DxDecoderWrap::DxLockFrame( DxGPUFrame * frame )
{
	if ( NULL == frame )
	{
		return -1;
	}

	std::lock_guard<std::mutex> l(m_queue_frames_mutex);
	
	if(m_queue_frames.size() <= 0) {
		return -1;
	}
	GPUFrame * decodedFrame = m_queue_frames.front();
	m_queue_frames.pop();

	AVFrame* gpuFrame = decodedFrame->gpuFrame;

	cudaSetDevice(atoi(m_cfg.gpuid.c_str()));
	cuCtxPushCurrent(m_cuda_cxt);
	cudaError_t cudaStatus;
	size_t  rgbSize = 3 * gpuFrame->width * gpuFrame->height * sizeof(unsigned char);
	if(m_pHwData == nullptr){
		cudaStatus = cudaMalloc((void **)&m_pHwData, rgbSize);
		if (cudaStatus != cudaSuccess) {
			LOG_ERROR("[{}]- cudaMalloc failed !!!", m_name.c_str());
			cuCtxPopCurrent(&m_cuda_cxt);
			return -1;
		}
	}
	cuda_common::setColorSpace( ITU_709, 0 );
	cudaStatus = cuda_common::CUDAToBGR((CUdeviceptr)gpuFrame->data[0],(CUdeviceptr)gpuFrame->data[1], gpuFrame->linesize[0], gpuFrame->linesize[1], m_pHwData, gpuFrame->width, gpuFrame->height);
	cuCtxPopCurrent(&m_cuda_cxt);
	if (cudaStatus != cudaSuccess) {
		LOG_ERROR("[{}]- CUDAToBGR failed !!!", m_name.c_str());
		return -1;
	}

	frame->width = gpuFrame->width;
	frame->height = gpuFrame->height;
	frame->size = rgbSize;
	frame->frame = m_pHwData;
	frame->timestamp = decodedFrame->ts;

	delete decodedFrame;
	decodedFrame = nullptr;

	return 0;
}

void DxDecoderWrap::post_decode_thread(GPUFrame * decodedFrame) {
	while(!m_bClose) {
		m_queue_frames_mutex.lock();
		if(m_queue_frames.size() < 3) {
			// 入队
			m_queue_frames.push(decodedFrame);
			m_queue_frames_mutex.unlock();
			break;
		} else {
			m_queue_frames_mutex.unlock();
			std::this_thread::sleep_for(std::chrono::milliseconds(10));
		}
	}
}

void DxDecoderWrap::decode_finished_thread() {
	m_bClose = true;
}