Hu Chunming / FFNvDecoder

#include "FFNvDecoderManager.h"

#include "FFNvDecoder.h"
#include "./gb28181/FFGB28181Decoder.h"

#include "logger.hpp"

using namespace std;


AbstractDecoder* FFNvDecoderManager::createDecoder(MgrDecConfig config){

    closeAllFinishedDecoder();

    if (config.cfg.post_decoded_cbk == nullptr || config.cfg.decode_finished_cbk== nullptr){
        return nullptr;
    }
    

    std::lock_guard<std::mutex> l(m_mutex);

    auto it = decoderMap.find(config.name);
    if (it != decoderMap.end()){
        LOG_ERROR("已存在name所标记的解码器");
        return nullptr;
    }

    AbstractDecoder* dec = nullptr;
    if(DECODER_TYPE_FFMPEG == config.dec_type){
        dec = new FFNvDecoder();
    }else if(DECODER_TYPE_GB28181 == config.dec_type){
        dec = new FFGB28181Decoder();
    }
    
    if (dec == nullptr){
        LOG_ERROR("没有指定解码器类型");
        return nullptr;
    }
    
    bool bRet= dec->init(config.cfg);
    if (bRet)
    {
        dec->setName(config.name) ;
        dec->post_decoded_cbk = config.cfg.post_decoded_cbk;
        dec->decode_finished_cbk = config.cfg.decode_finished_cbk;
        decoderMap[config.name] = dec;

        LOG_INFO("[{}][{}]- 解码器初始化成功",config.name, config.cfg.uri);
        return dec;
    }
    
    // 创建失败，关闭解码器
    dec->close();
    delete dec;

    LOG_ERROR("[{}][{}]- 解码器初始化失败！",config.name, config.cfg.uri);
    return nullptr;
}

bool FFNvDecoderManager::setPostDecArg(const string name, const void * userPtr)
{
    if (name.empty())
    {
        LOG_ERROR("name 为空!");
        return false;
    }

    std::lock_guard<std::mutex> l(m_mutex);

    auto dec = decoderMap.find(name);
    if (dec != decoderMap.end())
    {
        dec->second->m_postDecArg = userPtr;
        return true;
    }

    LOG_ERROR("没有找到name为{}的解码器",name);
    return false;
}

bool FFNvDecoderManager::setFinishedDecArg(const string name, const void * userPtr)
{
    if (name.empty())
    {
        LOG_ERROR("name 为空!");
        return false;
    }

    std::lock_guard<std::mutex> l(m_mutex);

    auto dec = decoderMap.find(name);
    if (dec != decoderMap.end())
    {
        dec->second->m_finishedDecArg = userPtr;
        return true;
    }

    LOG_ERROR("没有找到name为{}的解码器",name);
    return false;
}

AbstractDecoder* FFNvDecoderManager::getDecoderByName(const string name)
{
    if (name.empty())
    {
        LOG_ERROR("name 为空!");
        return nullptr;
    }
    
    std::lock_guard<std::mutex> l(m_mutex);

    auto dec = decoderMap.find(name);
    if (dec != decoderMap.end())
    {
        return dec->second;
    }

    LOG_ERROR("没有找到name为{}的解码器",name);
    return nullptr;
}

void FFNvDecoderManager::startDecode(AbstractDecoder* dec){
    if (dec != nullptr && !dec->isRunning())
    {
        dec->start();
    }
}

bool FFNvDecoderManager::startDecodeByName(const string name){
     if (name.empty())
    {
        LOG_ERROR("name 为空!");
        return false;
    }

    std::lock_guard<std::mutex> l(m_mutex);

    auto dec = decoderMap.find(name);
    if (dec != decoderMap.end())
    {
        return dec->second->start();
    }

    LOG_ERROR("没有找到name为{}的解码器",name);
    return false;
}

void FFNvDecoderManager::startAllDecode(){

    std::lock_guard<std::mutex> l(m_mutex);

    for(auto iter = decoderMap.begin(); iter != decoderMap.end(); iter++){
        if (!iter->second->isRunning())
        {
            iter->second->start();
        }
    }
}

bool FFNvDecoderManager::closeDecoderByName(const string name){
    if (name.empty())
    {
        LOG_ERROR("name 为空!");
        return false;
    }

    std::lock_guard<std::mutex> l(m_mutex);

    auto dec = decoderMap.find(name);
    if (dec != decoderMap.end())
    {
        dec->second->close();
        delete dec->second;
        dec->second = nullptr;
        decoderMap.erase(dec);

        return true;
    }
    
    LOG_ERROR("没有找到name为{}的解码器",name);
    return false;
}

void FFNvDecoderManager::closeAllDecoder()
{
    std::lock_guard<std::mutex> l(m_mutex);

    for(auto iter = decoderMap.begin(); iter != decoderMap.end(); iter++){
        iter->second->close();
        delete iter->second;
        iter->second = nullptr;
    }
    decoderMap.clear();
}

void FFNvDecoderManager::closeAllFinishedDecoder()
{
    std::lock_guard<std::mutex> l(m_mutex);

     for(auto iter = decoderMap.begin(); iter != decoderMap.end(); ){
        if (iter->second->isFinished())
        {
            delete iter->second;
            iter->second = nullptr;
            iter = decoderMap.erase(iter);
        }
        else
        {
            iter++ ;
        }
    }
}

int FFNvDecoderManager::count()
{
    closeAllFinishedDecoder();

    std::lock_guard<std::mutex> l(m_mutex);
    return decoderMap.size();
}

bool FFNvDecoderManager::pauseDecoder(const string name)
{
    if (name.empty())
    {
        LOG_ERROR("name 为空!");
        return false;
    }

    std::lock_guard<std::mutex> l(m_mutex);

    auto dec = decoderMap.find(name);
    if (dec != decoderMap.end())
    {
        dec->second->pause();
        return true;
    }
    
    LOG_ERROR("没有找到name为{}的解码器",name);
    return false;
}

bool FFNvDecoderManager::resumeDecoder(const string name)
{
    if (name.empty())
    {
        LOG_ERROR("name 为空!");
        return false;
    }

    std::lock_guard<std::mutex> l(m_mutex);

    auto dec = decoderMap.find(name);
    if (dec != decoderMap.end())
    {
        dec->second->resume();
        return true;
    }
    
    LOG_ERROR("没有找到name为{}的解码器",name);
    return false;
}

bool FFNvDecoderManager::isSurport(MgrDecConfig& config)
{
    {
        std::lock_guard<std::mutex> l(m_mutex);

        auto it = decoderMap.find(config.name);
        if (it != decoderMap.end()){
            LOG_ERROR("已存在name所标记的解码器");
            return false;
        }
    }

    AbstractDecoder* dec = nullptr;
    if(config.dec_type = DECODER_TYPE_FFMPEG){
        dec = new FFNvDecoder();
    }else if(config.dec_type = DECODER_TYPE_GB28181){
        dec = new FFGB28181Decoder();
    }
    
    if (dec == nullptr){
        LOG_ERROR("没有指定解码器类型");
        return false;
    }

    bool bRet = dec->isSurport(config.cfg);
    delete dec;
    dec = nullptr;

    return bRet;
}

bool FFNvDecoderManager::isRunning(const string name){
    if (name.empty())
    {
        LOG_ERROR("name 为空!");
        return false;
    }

    std::lock_guard<std::mutex> l(m_mutex);

    auto dec = decoderMap.find(name);
    if (dec != decoderMap.end())
    {
        return dec->second->isRunning();
    }
    
    LOG_ERROR("没有找到name为{}的解码器",name);
    return false;
}

bool FFNvDecoderManager::isFinished(const string name){
    if (name.empty())
    {
        LOG_ERROR("name 为空!");
        return false;
    }

    std::lock_guard<std::mutex> l(m_mutex);

    auto dec = decoderMap.find(name);
    if (dec != decoderMap.end())
    {
        return dec->second->isFinished();
    }
    
    LOG_ERROR("没有找到name为{}的解码器",name);
    return false;
}

bool FFNvDecoderManager::isPausing(const string name){
    if (name.empty())
    {
        LOG_ERROR("name 为空!");
        return false;
    }

    std::lock_guard<std::mutex> l(m_mutex);

    auto dec = decoderMap.find(name);
    if (dec != decoderMap.end())
    {
        return dec->second->isPausing();
    }
    
    LOG_ERROR("没有找到name为{}的解码器",name);
    return false;
}

bool FFNvDecoderManager::setDecKeyframe(const string name, bool bKeyframe)
{
    if (name.empty())
    {
        LOG_ERROR("name 为空!");
        return false;
    }

    std::lock_guard<std::mutex> l(m_mutex);

    auto dec = decoderMap.find(name);
    if (dec != decoderMap.end())
    {
        dec->second->setDecKeyframe(bKeyframe);
        return true;
    }
    
    LOG_ERROR("没有找到name为{}的解码器",name);
    return false;
}

bool FFNvDecoderManager::getResolution(const string name, int &width, int &height)
{
    if (name.empty())
    {
        LOG_ERROR("name 为空!");
        return false;
    }

    std::lock_guard<std::mutex> l(m_mutex);

    auto dec = decoderMap.find(name);
    if (dec != decoderMap.end())
    {
        dec->second->getResolution(width, height);
        return true;
    }
    
    LOG_ERROR("没有找到name为{}的解码器",name);
    return false;
}

vector<string> FFNvDecoderManager::getAllDecodeName(){
    
    closeAllFinishedDecoder();

    std::lock_guard<std::mutex> l(m_mutex);

    vector<string> decode_names;
    for(auto it = decoderMap.begin(); it != decoderMap.end(); ++it){
        decode_names.push_back(it->first);
    }
    return decode_names;
}

int FFNvDecoderManager::getCachedQueueLength(const string name){
    if (name.empty()){
        LOG_ERROR("name 为空!");
        return -1;
    }

    std::lock_guard<std::mutex> l(m_mutex);

    auto dec = decoderMap.find(name);
    if (dec != decoderMap.end()){
        return dec->second->getCachedQueueLength();
    }
    
    LOG_ERROR("没有找到name为{}的解码器",name);
    return -1;
}

FFImgInfo* FFNvDecoderManager::snapshot(const string& uri){
    if (uri.empty()){
        return nullptr;
    }
 
    AVFormatContext* ifmt_ctx = nullptr;
	AVCodecContext* codec_ctx = nullptr;
	AVCodec* codec = nullptr;
	AVPacket* pkt = nullptr;
	AVFrame *frame = nullptr;
	AVFrame *pFrameRGB = nullptr;	
	int video_index = -1;
	AVStream* st = nullptr;
	SwsContext *img_convert_ctx = nullptr;
	uint8_t *buffer = nullptr;
    int numBytes = 0;
	int index = 0;

	FFImgInfo* imgInfo = nullptr;
 
 	//av_register_all();
	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, "stimeout", "30000000", 0 ); // 单位为 百万分之一秒
	
	///打开输入的流
	int ret = avformat_open_input(&ifmt_ctx, uri.c_str(), nullptr, &options);
	if (ret != 0){
		printf("Couldn't open input stream.\n");
		goto end_flag ;
	}
 
	//查找流信息
	if (avformat_find_stream_info(ifmt_ctx, nullptr) < 0){
		printf("Couldn't find stream information.\n");
		goto end_flag ;
	}
 
	//找到视频流索引
    video_index = av_find_best_stream(ifmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
 
    st = ifmt_ctx->streams[video_index];
    
    //找到解码器
    codec = avcodec_find_decoder(st->codecpar->codec_id);
    if (!codec){
        fprintf(stderr, "Codec not found\n");
        goto end_flag ;
    }
 
    //申请AVCodecContext
    codec_ctx = avcodec_alloc_context3(codec);
    if (!codec_ctx){
        goto end_flag ;
    }
 
	avcodec_parameters_to_context(codec_ctx, ifmt_ctx->streams[video_index]->codecpar);
 
    //打开解码器
    if ((ret = avcodec_open2(codec_ctx, codec, nullptr) < 0)){
        goto end_flag ;
    }
	
    // 计算解码后原始数据所需缓冲区大小，并分配内存空间 Determine required buffer size and allocate buffer
    numBytes = av_image_get_buffer_size(AV_PIX_FMT_RGB24, codec_ctx->width, codec_ctx->height, 1);
    buffer = (uint8_t *)av_malloc(numBytes * sizeof(uint8_t));
    
	pFrameRGB = av_frame_alloc();
    av_image_fill_arrays(pFrameRGB->data, pFrameRGB->linesize, buffer, AV_PIX_FMT_BGR24, codec_ctx->width, codec_ctx->height, 1);

	img_convert_ctx = sws_getContext(codec_ctx->width, codec_ctx->height,codec_ctx->pix_fmt, codec_ctx->width, codec_ctx->height, AV_PIX_FMT_BGR24, 
		   SWS_BICUBIC, nullptr, nullptr, nullptr);
 
	pkt = av_packet_alloc();
	frame = av_frame_alloc();
	while (av_read_frame(ifmt_ctx, pkt) >= 0){
		if (pkt->stream_index == video_index){
			int ret = avcodec_send_packet(codec_ctx, pkt);
			if (ret >= 0){
				ret = avcodec_receive_frame(codec_ctx, frame);
				if ((ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) || ret < 0){
					LOG_ERROR("Failed to receive frame: {}",ret);
					continue;
				}

				index ++ ;

				if (index >= 5){
					// 取解码出来的第三帧，应该可以一定程度优化花屏问题
					sws_scale(img_convert_ctx, (const unsigned char* const*)frame->data, frame->linesize, 0, codec_ctx->height, pFrameRGB->data, pFrameRGB->linesize);

					imgInfo = new FFImgInfo();
					imgInfo->pData = buffer;
					imgInfo->height = codec_ctx->height;
					imgInfo->width = codec_ctx->width;

					break;
				}
			}
		}
		av_packet_unref(pkt);
	}

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 (frame != nullptr){
		av_frame_free(&frame);
	}

	if (pFrameRGB != nullptr){
		av_frame_free(&pFrameRGB);
	}

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

	return imgInfo;
}

void FFNvDecoderManager::releaseFFImgInfo(FFImgInfo* info){
	if(nullptr != info){
		if(info->pData != nullptr){
			av_free(info->pData);
			info->pData = nullptr;
		}
		delete info;
		info = nullptr;
	}
}