#include "FFmpegFunc.h"

void PrepareFFmpeg()
{
    av_register_all();
    avformat_network_init();
}

enum AVSampleFormat SelectSampleFormat(const AVCodec *codec)
{
    return codec->sample_fmts ? codec->sample_fmts[0] : AV_SAMPLE_FMT_S16;
}

int SelectSampleRate(const AVCodec *codec)
{
    const int *p;
    int best_samplerate = 0;
    if (!codec->supported_samplerates)
        return 44100;
    p = codec->supported_samplerates;
    while (*p) {
        if (!best_samplerate || abs(44100 - *p) < abs(44100 - best_samplerate))
            best_samplerate = *p;
        p++;
    }
    return best_samplerate;
}

int SelectChannelLayout(const AVCodec *codec)
{
    const uint64_t *p;
    uint64_t best_ch_layout = 0;
    int best_nb_channels   = 0;
    if (!codec->channel_layouts)
        return AV_CH_LAYOUT_STEREO;
    p = codec->channel_layouts;
    while (*p) {
        int nb_channels = av_get_channel_layout_nb_channels(*p);
        if (nb_channels > best_nb_channels) {
            best_ch_layout    = *p;
            best_nb_channels = nb_channels;
        }
        p++;
    }
    return best_ch_layout;
}


int OpenOutput(AVFormatContext **pFmtCtx,const char *strOutFmtName,const char *strOutputName)
{
    int nRet;
    avformat_alloc_output_context2(pFmtCtx, NULL, strOutFmtName, strOutputName);
    if (!(*pFmtCtx)) {
        printf("Could not deduce output format from file extension: %s\n",strOutputName);
        return -1;
    }

    if (!((*pFmtCtx)->flags & AVFMT_NOFILE)) {
        nRet = avio_open(&(*pFmtCtx)->pb, strOutputName, AVIO_FLAG_WRITE);
        if (nRet < 0) {
            printf("Could not open '%s'\n", strOutputName);
            return -2;
        }
    }

    return 0;
}

void CloseOutput(AVFormatContext **pFmtCtx)
{
    if (!((*pFmtCtx)->flags & AVFMT_NOFILE))
        avio_closep(&(*pFmtCtx)->pb);

    avformat_free_context(*pFmtCtx);
}



int WritePacket(AVOutputStream *pOutStream,AVPacket *pPacket)
{
    /* rescale output packet timestamp values from codec to stream timebase */
    av_packet_rescale_ts(pPacket, pOutStream->pCodecContext->time_base, pOutStream->pStream->time_base);
    pPacket->stream_index = pOutStream->pStream->index;
    return av_interleaved_write_frame(pOutStream->pFmtCtx, pPacket);
}

int WriteHeader(AVFormatContext *pFmtCtx,AVDictionary *pDict)
{
    int nRet;
    nRet = avformat_write_header(pFmtCtx, &pDict);
    if (nRet < 0) {
        return -1;
    }

    return 0;
}

int WriteTrailer(AVFormatContext *pFmtCtx)
{
    av_write_trailer(pFmtCtx);
}

AVOutputStream *CreateOutputStreamByDefalutParam(enum AVCodecID iCodecID)
{
    AVCodec *pCodec;
    AVCodecContext *pCodecCxt;

    if(AV_CODEC_ID_NONE == iCodecID)
        return NULL;

    pCodec = avcodec_find_encoder(iCodecID);
    if(!pCodec)
    {
        printf("Could not find encoder for '%s'\n",avcodec_get_name(iCodecID));
        return NULL;
    }

    pCodecCxt = avcodec_alloc_context3(pCodec);
    if (!pCodecCxt) {
        fprintf(stderr, "Could not alloc an encoding context\n");
        return NULL;
    }

    switch (pCodec->type) {
    case AVMEDIA_TYPE_AUDIO:
        pCodecCxt->codec_id = iCodecID;
        pCodecCxt->sample_fmt  = SelectSampleFormat(pCodec);
        pCodecCxt->bit_rate    = 64000;
        pCodecCxt->sample_rate = SelectSampleRate(pCodec);
        pCodecCxt->channel_layout = SelectChannelLayout(pCodec);
        pCodecCxt->channels = av_get_channel_layout_nb_channels(pCodecCxt->channel_layout);
        break;
    case AVMEDIA_TYPE_VIDEO:
        pCodecCxt->codec_id     = iCodecID;
        pCodecCxt->bit_rate     = 400000;
        pCodecCxt->width        = 800;
        pCodecCxt->height       = 600;
        pCodecCxt->time_base    = AVRational{ 1, 25 };
        pCodecCxt->framerate    = AVRational{25, 1};
        pCodecCxt->gop_size     = 10;
        pCodecCxt->max_b_frames = 1;
        pCodecCxt->pix_fmt      = AV_PIX_FMT_YUV420P;
        break;
    default:
        break;
    }

    AVOutputStream *pAVOutputStream = (AVOutputStream *)av_mallocz(sizeof(AVOutputStream));

    pAVOutputStream->pCodec = pCodec;
    pAVOutputStream->pCodecContext = pCodecCxt;

    return pAVOutputStream;
}

int OpenOutputStream(AVOutputStream *pAVOutputStream, AVDictionary *pOpt)
{
    int nRet;
    AVDictionary    *pOutDict;
    pOutDict = NULL;

    if(NULL == pAVOutputStream)
        return -1;

    av_dict_copy(&pOutDict, pOpt, 0);

    nRet = avcodec_open2(pAVOutputStream->pCodecContext, pAVOutputStream->pCodec, &pOutDict);
    av_dict_free(&pOutDict);
    if (nRet < 0) {
        printf("Could not open video codec: %s\n",avcodec_get_name(pAVOutputStream->pCodec->id));
        return -2;
    }

    //如果Stream加入了输出文件
    if(pAVOutputStream->pFmtCtx && pAVOutputStream->pStream)
    {
        nRet = avcodec_parameters_from_context(pAVOutputStream->pStream->codecpar,pAVOutputStream->pCodecContext);
        if (nRet < 0) {
            fprintf(stderr, "Could not copy the stream parameters\n");
            return -2;
        }
    }

    return 0;
}

void FreeAVOutputStream(AVOutputStream **pOutStream)
{
    if(NULL == pOutStream)
        return;

    if(NULL == (*pOutStream))
        return;

    //empty
    AVOutputStream *pTmpOutputStream = *pOutStream;
    if(pTmpOutputStream->pCodecContext)
    {
        avcodec_free_context(&(pTmpOutputStream->pCodecContext));
    }
    memset(pTmpOutputStream,0,sizeof(AVOutputStream));

    //free
    av_freep(pOutStream);
    *pOutStream = NULL;

    return;
}

int AddStreamToOutput(AVFormatContext *pFmtCtx,AVOutputStream *pOutStream)
{
    int nRet;
    AVStream *pStream;

    pStream = avformat_new_stream(pFmtCtx,NULL);
    if(!pStream)
    {
        printf("ould not allocate stream\n");
        return -1;
    }
    pStream->id = pFmtCtx->nb_streams -1;

    switch (pOutStream->pCodec->type) {
    case AVMEDIA_TYPE_AUDIO:
        pStream->time_base = AVRational{1, pOutStream->pCodecContext->sample_rate};
        break;
    case AVMEDIA_TYPE_VIDEO:
        pStream->time_base = pOutStream->pCodecContext->time_base;
        break;
    }

    if (pFmtCtx->oformat->flags & AVFMT_GLOBALHEADER)
        pOutStream->pCodecContext->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

    pOutStream->pFmtCtx = pFmtCtx;
    pOutStream->pStream = pStream;
    return 0;
}
