Hu Chunming / FFNvDecoder

#include "FFReceiver.h"
#include <fstream>

const int g_pkt_size = 1024 * 1024;  // 单个AVPacket大小的最大值

FFReceiver::FFReceiver(/* args */)
{
    fmt_ctx = nullptr;
	m_bRunning = false;

	stream = nullptr;
    stream_index = -1;
    pix_fmt = AV_PIX_FMT_NONE;
    m_dec_name = "";

	m_bPause = false;
	m_bReal = true;

	m_bFinished = false;
	m_dec_keyframe = false;
	m_fps = 0.0;

    m_read_thread = 0;
}

FFReceiver::~FFReceiver()
{
    releaseFFmpeg();

	// 这个只能放在析构函数中，因为会影响到解码类中的m_pktQueueptr队列
	// 所以应当确保在所有工作线程都退出后才释放
	for(int i = 0; i < m_vec_pkt.size(); i++){
		av_packet_free(&m_vec_pkt[i]);
	}
}

AVCodecContext* FFReceiver::init_FFmpeg(ReceiverConfig config){

#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(58, 9, 100)
    av_register_all();
#endif
#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(58, 10, 100)
    avcodec_register_all();
#endif

    avformat_network_init();

	const char* uri = config.uri;
    fstream infile(uri);
	if (infile.is_open()){
		m_bReal = false;
		infile.close();
	}else {
		m_bReal = true;
	}

    m_dec_name = config.dec_name;
	m_pktQueueptr = config.pktQueueptr;
	receiver_finished_cbk = config.receiver_finished_cbk;

	// 打开输入视频文件
	AVDictionary *options = nullptr;
	av_dict_set( &options, "bufsize", "655360", 0 );
	av_dict_set( &options, "rtsp_transport", config.force_tcp ? "tcp" : "udp", 0 );
	av_dict_set( &options, "stimeout", "30000000", 0 ); // 单位为 百万分之一秒
	
	fmt_ctx = avformat_alloc_context();
	const char* input_file = uri;
	if (avformat_open_input(&fmt_ctx, input_file, nullptr, &options) != 0) {
        LOG_ERROR("[{}]- Cannot open input file: {}", m_dec_name, input_file);
		return nullptr;
	}
    av_dump_format(fmt_ctx, 0, input_file, 0);

	// 查找流信息
	if (avformat_find_stream_info(fmt_ctx, nullptr) < 0) {
        LOG_ERROR("[{}]- Cannot find input stream information!", m_dec_name);
		return nullptr;
	}

	// 查找视频流信息
	AVCodec *decoder = nullptr;
	stream_index = av_find_best_stream(fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, &decoder, 0);
	if (stream_index < 0) {
        LOG_ERROR("[{}]- Cannot find a video stream in the input file!", m_dec_name);
		return nullptr;
	}
	AVCodec *vcodec = avcodec_find_decoder(decoder->id);

	avctx = avcodec_alloc_context3(vcodec);
	if(avctx == nullptr){
        LOG_ERROR("[{}]- alloc AVCodecContext failed!", m_dec_name);
		return nullptr;
	}

	do{
		// 得到视频流对象
		AVStream* stream = fmt_ctx->streams[stream_index];
		AVCodecParameters *codecpar = stream->codecpar;
		if (avcodec_parameters_to_context(avctx, codecpar) < 0)
			break;

		const AVBitStreamFilter * filter = nullptr;
		if(codecpar->codec_id == AV_CODEC_ID_H264){
			filter = av_bsf_get_by_name("h264_mp4toannexb");
		}else if(codecpar->codec_id == AV_CODEC_ID_HEVC){
			filter = av_bsf_get_by_name("hevc_mp4toannexb");
		}else {
            LOG_ERROR("[{}]- codec_id is not supported!", m_dec_name);
			break;
		}

		int ret = av_bsf_alloc(filter, &h264bsfc);
		if (ret < 0){
			break;
		}
		
		avcodec_parameters_copy(h264bsfc->par_in, codecpar);
		av_bsf_init(h264bsfc);

		frame_width = codecpar->width;
		frame_height = codecpar->height;
		pix_fmt = (AVPixelFormat)codecpar->format;
		m_fps = av_q2d(stream ->avg_frame_rate);

        LOG_INFO("[{}]- init ffmpeg success! input:{} frame_width:{} frame_height:{} fps:{} ", m_dec_name, input_file, frame_width, frame_height, m_fps);

		for(int i = 0; i<5; i++){
			AVPacket* pkt = av_packet_alloc();
			av_init_packet( pkt );
			m_vec_pkt.push_back(pkt);
		}
		m_pktQueueptr->init(m_vec_pkt);

		return avctx;
	}while(0);

    LOG_ERROR("[{}]- init ffmpeg failed ! input:{} ", m_dec_name);

    return nullptr;
}

void FFReceiver::releaseFFmpeg(){
	m_dec_keyframe = false;
	if(h264bsfc){
		av_bsf_free(&h264bsfc);
		h264bsfc = nullptr;
	}
	if (fmt_ctx){
		avformat_close_input(&fmt_ctx);
		fmt_ctx = nullptr;
	}
    if(avctx){
        avcodec_free_context(&avctx);
        avctx = nullptr;
    }
}

void FFReceiver::read_thread(){

    int frame_count = 0;
	int ret = -1;
	while (m_bRunning)
	{
		if (!m_bReal)
		{
			if (m_bPause)
			{
				std::this_thread::sleep_for(std::chrono::milliseconds(3));
				continue;
			}
		}

		AVPacket* pkt = m_pktQueueptr->getTail();
		if(pkt == nullptr){
			std::this_thread::sleep_for(std::chrono::milliseconds(3));
			continue;
		}

		int result = av_read_frame(fmt_ctx, pkt);
		if (result == AVERROR_EOF || result < 0)
		{
            LOG_ERROR("[{}]- Failed to read frame!", m_dec_name);
			break;
		}

		if (m_dec_keyframe && !(pkt->flags & AV_PKT_FLAG_KEY)) {
			av_packet_unref(pkt);
			continue;
		}

		if (stream_index == pkt->stream_index){

            ret = av_bsf_send_packet(h264bsfc, pkt);
            if(ret < 0) {
                LOG_ERROR("[{}]- av_bsf_send_packet error!", m_dec_name);
            }

            while ((ret = av_bsf_receive_packet(h264bsfc, pkt)) == 0) {
                if(pkt->size > g_pkt_size){
                    LOG_ERROR("[{}]- pkt size 大于最大预设值!", m_dec_name);
					break;
                }

				if(!m_bRunning){
					break;
				}

                m_pktQueueptr->addTail();

                frame_count++;
            }
		}
	}

    LOG_INFO("[{}]- read thread exit.", m_dec_name);
	m_bFinished = true;

	receiver_finished_cbk(m_finishedReceiveArg);
}

bool FFReceiver::start(){
	m_bRunning = true;

	pthread_create(&m_read_thread,0,
        [](void* arg)
        {
            FFReceiver* a=(FFReceiver*)arg;
            a->read_thread();
            return (void*)0;
        }
    ,this);

	return true;
}

void FFReceiver::close(){
	m_bRunning=false;

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

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

bool FFReceiver::getResolution( int &width, int &height ){
	width = frame_width;
	height = frame_height;
	return true;
}

void FFReceiver::pause(){
	m_bPause = true;
}

void FFReceiver::resume(){
	m_bPause = false;
}

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

bool FFReceiver::isRunning(){
	return m_bRunning;
}

bool FFReceiver::isFinished(){
	return m_bFinished;
}

bool FFReceiver::isPausing(){
	return m_bPause;
}

void FFReceiver::setFinishCbkArg(const void* userPtr){
	m_finishedReceiveArg = userPtr;
}