#include "ImageHandler.h"

CImageHandler::CImageHandler(CImageIO *piCVMatIO)
{
    m_psH264CodecCtx = NULL;
    m_psYUVFrame = NULL;
    av_init_packet(&m_iOutputPkt);
    m_iOutputPkt.data = NULL;
    m_iOutputPkt.size = 0;    
    m_piCVMatIO = piCVMatIO;

    Init();
}

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

int CImageHandler::Init()
{
    av_register_all();
    avformat_network_init();

    const AVCodec *pH264Codec = NULL;
    pH264Codec = avcodec_find_encoder_by_name("libx264");
    if (!pH264Codec) {
        fprintf(stderr, "Codec '%s' not found\n", "h264");
        return -1;
    }

    m_psH264CodecCtx = avcodec_alloc_context3(pH264Codec);
    if (!m_psH264CodecCtx) {
        fprintf(stderr, "Could not allocate video codec context\n");
        return -2;
    }

    m_psH264CodecCtx->width = m_piCVMatIO->GetWidth();
    m_psH264CodecCtx->height = m_piCVMatIO->GetHeight();
    m_psH264CodecCtx->time_base.num = 1;
    m_psH264CodecCtx->time_base.den = m_piCVMatIO->GetFrameRate();
    m_psH264CodecCtx->pix_fmt = AV_PIX_FMT_YUV420P;

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


    int nRet = avcodec_open2(m_psH264CodecCtx, pH264Codec, &pH264Dict);
    if (nRet < 0) {
        fprintf(stderr, "Could not open H.264\n");
        return -9;
    }

    av_init_packet(&m_iOutputPkt);
    m_iOutputPkt.data = NULL;
    m_iOutputPkt.size = 0;

    m_psYUVFrame = AllocAVFrame(AV_PIX_FMT_YUV420P,m_piCVMatIO->GetWidth(), m_piCVMatIO->GetHeight());
    return 0;
}

int CImageHandler::UnInit()
{
    avcodec_free_context(&m_psH264CodecCtx);
    av_packet_unref(&m_iOutputPkt);
    av_frame_free(&m_psYUVFrame);
    return 0;
}

bool CImageHandler::CVMatToAVFrame(const cv::Mat &iMat,AVFrame *pFrame)
{
    if (pFrame && !iMat.empty()) {

        cv::Mat yuv;
        cv::cvtColor(iMat, yuv, cv::COLOR_BGR2YUV_I420);
        // calc frame size
        int frame_size = iMat.cols * iMat.rows;
        unsigned char *pdata = yuv.data;

        memcpy(pFrame->data[0],pdata,frame_size);
        memcpy(pFrame->data[1],pdata + frame_size,frame_size*1/4);
        memcpy(pFrame->data[2],pdata + frame_size*5/4,frame_size*1/4);

        return true;
    }
    return false;
}

AVFrame *CImageHandler::AllocAVFrame(enum AVPixelFormat sPixFmt, int nWidth, int nHeight)
{
    AVFrame *pFrame = NULL;
    int nRet;

    pFrame = av_frame_alloc();
    if (!pFrame)
        return NULL;

    pFrame->format = sPixFmt;
    pFrame->width = nWidth;
    pFrame->height = nHeight;

    /* allocate the buffers for the frame data */
    nRet = av_frame_get_buffer(pFrame, 0);
    if (nRet < 0) {
        fprintf(stderr, "Could not allocate frame data.\n");
        return NULL;
    }

    return pFrame;
}

int CImageHandler::Encode(AVCodecContext *psCodecCtx,AVPacket *psPkt,AVFrame *frame)
{
    int nRet;
    nRet = avcodec_send_frame(psCodecCtx, frame);
    if (nRet < 0) {
        fprintf(stderr, "Error sending a frame for encoding\n");
        return -1;
    }

    while (nRet >= 0) {
        nRet = avcodec_receive_packet(psCodecCtx, psPkt);
        if (nRet == AVERROR(EAGAIN) || nRet == AVERROR_EOF)
            return -1;
        else if (nRet < 0)
        {
            fprintf(stderr, "Error during encoding\n");
            return -1;
        }
        return 0;
    }
    return -1;
}

int CImageHandler::GetFrameData(uint8_t **pBuffer, int *nBufferLen)
{
    int nMaxSize = *nBufferLen;
    int nRet = 0; 
    while (1){
        bool bKeyFrame = false;
        cv::Mat iMat = m_piCVMatIO->ReadCVMat(bKeyFrame);
        if (!iMat.empty())
        {
            if(CVMatToAVFrame(iMat,m_psYUVFrame))
            {
                static int nIndex = 0;
                if(nIndex == 0)
                {
                    m_psYUVFrame->pict_type = AV_PICTURE_TYPE_I;
                }
                else
                {
                    m_psYUVFrame->pict_type = AV_PICTURE_TYPE_P;
                }

                nIndex++;
                if(nIndex >= GOPSIZE)
                {
                    nIndex = 0;
                }

                if(bKeyFrame)
                {
                    m_psYUVFrame->pict_type = AV_PICTURE_TYPE_I;
                }

                av_packet_unref(&m_iOutputPkt);
                av_init_packet(&m_iOutputPkt);
                m_iOutputPkt.data = NULL;
                m_iOutputPkt.size = 0;

                nRet = Encode(m_psH264CodecCtx,&m_iOutputPkt,m_psYUVFrame);
                if(nRet == 0)
                {
                    int nCopyLen = m_iOutputPkt.size>nMaxSize ? nMaxSize:m_iOutputPkt.size;
                    memcpy(*pBuffer,m_iOutputPkt.data,nCopyLen);
                    *nBufferLen = nCopyLen;

                    m_piCVMatIO->HandleEncodeAVPacket(&m_iOutputPkt,bKeyFrame);
                }

                return 0;
            }

        }

        return -1;
    }
    return -1;
}
