#include "MediaRecoder.h"

CMediaRecoder::CMediaRecoder()
{
    Init();
}

CMediaRecoder::~CMediaRecoder()
{
    UnInit();
}

bool CMediaRecoder::IsRecoding()
{
    return m_bIsRecording;
}

bool CMediaRecoder::StartRecode(uint nFrameRate, uint nWidth, uint nHeight, const std::string &strFileName)
{
    m_nFrameRate = nFrameRate;
    m_nWidth = nWidth;
    m_nHeight = nHeight;
    m_strFileName = strFileName;

    bool bRet = StartRecodeAndWriteHeader();
    if(bRet)
    {
        nVideoNextPts = 0;
        m_bIsRecording = true;
    }

    return bRet;
}

void CMediaRecoder::Init()
{
    m_nFrameRate = 25;
    m_nWidth = 320;
    m_nHeight = 240;
    m_strFileName = "";

    pOutFmtCtx = NULL;
    pOutputVideoStream = NULL;

    nVideoNextPts = 0;

    m_bIsRecording = false;
}

void CMediaRecoder::UnInit()
{
    if(IsRecoding())
    {
        FinishRecode();
    }
}

int CMediaRecoder::OpenOutPutAndAddStreams(AVFormatContext **pFmt,const char *strOutFileName,AVOutputStream **pVideoOutput)
{
    int nRet;
    AVFormatContext *pOutFmtCtx;
    AVOutputStream *pOutputVideoStream;
    pOutFmtCtx = NULL;
    pOutputVideoStream = NULL;

    //1
    nRet = OpenOutput(&pOutFmtCtx,NULL,strOutFileName);
    if(0 != nRet)
    {
        return -1;
    }

    //2
    pOutputVideoStream = CreateOutputStreamByDefalutParam(AV_CODEC_ID_H264);
    if(NULL == pOutputVideoStream)
    {
        return -2;
    }

    pOutputVideoStream->pCodecContext->codec_type = AVMEDIA_TYPE_VIDEO;
    /* resolution must be a multiple of two */
    pOutputVideoStream->pCodecContext->width = m_nWidth;
    pOutputVideoStream->pCodecContext->height = m_nHeight;
    /* frames per second */
    pOutputVideoStream->pCodecContext->time_base.num = 1;
    pOutputVideoStream->pCodecContext->time_base.den = m_nFrameRate;
    pOutputVideoStream->pCodecContext->framerate.num = m_nFrameRate;
    pOutputVideoStream->pCodecContext->framerate.den = 1;
    pOutputVideoStream->pCodecContext->max_b_frames = 30;
    pOutputVideoStream->pCodecContext->pix_fmt = AV_PIX_FMT_YUV420P;
    pOutputVideoStream->pCodecContext->bit_rate = 1024*1000;

    nRet = AddStreamToOutput(pOutFmtCtx,pOutputVideoStream);
    if(0 != nRet)
    {
        return -3;
    }

    AVDictionary *pH265Dict = NULL;
    av_dict_set(&pH265Dict, "preset", "ultrafast", 0);
    av_dict_set(&pH265Dict, "tune", "zerolatency", 0);

    nRet = OpenOutputStream(pOutputVideoStream,pH265Dict);
    av_dict_free(&pH265Dict);
    if(0 != nRet)
    {
        return -4;
    }

    *pFmt = pOutFmtCtx;
    *pVideoOutput = pOutputVideoStream;

    return 0;
}

bool CMediaRecoder::StartRecodeAndWriteHeader()
{
    if(m_strFileName.empty())
        return false;

    if(IsRecoding())
        return false;

    int nRet;
    pOutFmtCtx = NULL;
    pOutputVideoStream = NULL;
    nVideoNextPts = 0;

    PrepareFFmpeg();

    nRet = OpenOutPutAndAddStreams(&pOutFmtCtx,m_strFileName.c_str(),&pOutputVideoStream);
    if(0 != nRet)
    {
        printf(u8"打开输出文件并向其中加入流失败！\n");
        return false;
    }

    nRet = WriteHeader(pOutFmtCtx, NULL);
    if (nRet < 0) {
        return false;
    }

    return true;
}

bool CMediaRecoder::RecoderPacket(AVPacket *pPacket,bool bKeyFrame)
{
    if(!IsRecoding())
        return false;

    if(nVideoNextPts == 0)
    {
        if(!bKeyFrame)
        {
            return false;
        }
    }

    nVideoNextPts++;
    pPacket->dts = nVideoNextPts;
    pPacket->pts = nVideoNextPts;

    int nRet;
    nRet = WritePacket(pOutputVideoStream,pPacket);
    if(0 != nRet)
    {
        printf(u8"add packet failed！\n");
        return false;
    }

    return true;
}

bool CMediaRecoder::FinishRecode()
{
    if(!IsRecoding())
        return false;

    WriteTrailer(pOutFmtCtx);

    FreeAVOutputStream(&pOutputVideoStream);
    CloseOutput(&pOutFmtCtx);

    pOutFmtCtx = NULL;
    pOutputVideoStream = NULL;

    nVideoNextPts = 0;
    m_bIsRecording = false;
}
