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_callback(gpuFrame);
	}
}

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

DxDecoderWrap::DxDecoderWrap( const DxConfig * cfg )
{
	m_bClose = false;
	m_pDec = 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;

	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_bReal = m_pDec->isRealStream();
		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) {
		DxGPUFrame decodedFrame = m_queue_frames.front();
		if (decodedFrame.frame) {
			cudaFree(decodedFrame.frame);
			decodedFrame.frame = nullptr;
		}
		m_queue_frames.pop();
	}
	m_queue_frames_mutex.unlock();

	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->isFinished();
	}

	return false;
}

bool DxDecoderWrap::DxDecoderIsFinished() {
	if(m_pDec) {
		return m_pDec->isFinished();
	}

	return true;
}

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)
{
	std::lock_guard<std::mutex> l(m_queue_frames_mutex);
	
	if(m_queue_frames.size() <= 0) {
		return -1;
	}
	frame = m_queue_frames.front();
	m_queue_frames.pop();

	return 0;
}

DxGPUFrame DxDecoderWrap::DxGetFrame()
{
	std::lock_guard<std::mutex> l(m_queue_frames_mutex);
	
	DxGPUFrame frame;
	if(m_queue_frames.size() <= 0) {
		return frame;
	}
	frame = m_queue_frames.front();
	m_queue_frames.pop();

	return frame;
}

void DxDecoderWrap::post_decode_callback(GPUFrame * decodedFrame) {
	while(!m_bClose) {
		m_queue_frames_mutex.lock();
		if(m_queue_frames.size() < 3) {
			// 入队
			AVFrame* gpuFrame = decodedFrame->gpuFrame;

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

			DxGPUFrame frame;
			frame.width = gpuFrame->width;
			frame.height = gpuFrame->height;
			frame.size = gpuFrame->width;
			frame.frame = pHwData;
			frame.timestamp = decodedFrame->ts;
			frame.dec_name = m_name;

			m_queue_frames.push(frame);
			m_queue_frames_mutex.unlock();
			break;
		} else {
			m_queue_frames_mutex.unlock();
			if (m_bReal) {
				// 实时流不等待，直接弃帧
				return;
			}
			std::this_thread::sleep_for(std::chrono::milliseconds(10));
		}
	}
}

void DxDecoderWrap::decode_finished_callback() {
	// m_bClose = true;
}