/*

  Compile with:
  g++ -Wall -O2 -g -D__STDC_CONSTANT_MACROS  -o video-sound-transcoding-test video-sound-transcoding-test.cpp -lm -lavdevice -lavformat -lavfilter -lavcodec -lswresample -lswscale -lavutil

*/

extern "C"
{
#include <libavutil/avutil.h>
#include <libavutil/parseutils.h>
#include <libavutil/mathematics.h>
#include <libavutil/opt.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libswscale/swscale.h>
#include <libswresample/swresample.h>
#include <libavfilter/avfilter.h>
#include <libavfilter/avfiltergraph.h>
#include <libavfilter/avcodec.h>
#include <libavfilter/buffersink.h>
#include <libavfilter/buffersrc.h>
#include <libavutil/imgutils.h>
}

#include <iostream>

using namespace std;

//
// Compilation error work around
//
#ifndef av_ts2timestr
#include <vector>
#define AV_TS_MAX_STRING_SIZE 32
static inline string av_ts_make_time_string(int64_t ts, AVRational *tb)
{
    vector<char> buf(AV_TS_MAX_STRING_SIZE);
    if (ts == AV_NOPTS_VALUE) snprintf(buf.data(), AV_TS_MAX_STRING_SIZE, "NOPTS");
    else                      snprintf(buf.data(), AV_TS_MAX_STRING_SIZE, "%.6g", av_q2d(*tb) * ts);
    return string(buf.data());
}
#define av_ts2timestr(ts, tb) av_ts_make_time_string(ts, tb)
#endif

struct MediaContext
{
    MediaContext()
        : formatCtx(0),
          videoCodec(0),
          audioCodec(0),
          videoCodecCtx(0),
          audioCodecCtx(0),
          videoStream(0),
          audioStream(0),
          swsCtx(0),
          videoStreamIndex(-1),
          audioStreamIndex(-1)
    {}

    // input or output file name or URI
    string resource;

    AVFormatContext *formatCtx;
    AVCodec         *videoCodec;
    AVCodec         *audioCodec;
    AVCodecContext  *videoCodecCtx;
    AVCodecContext  *audioCodecCtx;
    AVStream        *videoStream;
    AVStream        *audioStream;
    SwsContext      *swsCtx;

    int              videoStreamIndex;
    int              audioStreamIndex;
};


MediaContext inCtx;   // Media source
MediaContext outCtx;  // Media destination


static
void usage(const string& progName)
{
    cout << "Use: " << progName << " <input media file> <output media file>" << endl;
}



static
int openCodecContex(MediaContext &ctx, AVMediaType type)
{
    int ret = -1;

    AVStream        **stream;
    AVCodecContext  **codecCtx;
    AVCodec         **codec;
    AVFormatContext *formatCtx;
    int             *streamIndex;

    formatCtx = ctx.formatCtx;

    if (type == AVMEDIA_TYPE_VIDEO)
    {
        stream      = &ctx.videoStream;
        codecCtx    = &ctx.videoCodecCtx;
        streamIndex = &ctx.videoStreamIndex;
        codec       = &ctx.videoCodec;
    }
    else if (type == AVMEDIA_TYPE_AUDIO)
    {
        stream      = &ctx.audioStream;
        codecCtx    = &ctx.audioCodecCtx;
        streamIndex = &ctx.audioStreamIndex;
        codec       = &ctx.audioCodec;
    }
    else
    {
        cerr << "Unknown media type: " << av_get_media_type_string(type) << endl;
        return -1;
    }


    ret = av_find_best_stream(formatCtx, type, -1, -1, 0, 0);
    if (ret >= 0)
    {
        *stream   = formatCtx->streams[ret];
        *codecCtx    = (*stream)->codec;

        *codec = avcodec_find_decoder((*codecCtx)->codec_id);

        if (!(*codec))
        {
            cerr << "Failed to find " << (*codecCtx)->codec_name << " codec" << endl;
            return -1;
        }

        int result = avcodec_open2(*codecCtx, *codec, 0);
        if (result < 0)
        {
            cerr << "Failed to open " << (*codec)->name << " codec" << endl;
            return -1;
        }
    }

    *streamIndex = ret;

    return ret;
}



static
int addStream(MediaContext &ctx, AVMediaType type)
{
    AVCodecID       codecId;
    AVStream        **stream;
    AVCodecContext  **codecCtx;
    AVCodec         **codec;
    AVFormatContext *formatCtx;
    int             *streamIndex;

    formatCtx = ctx.formatCtx;

    if (type == AVMEDIA_TYPE_VIDEO)
    {
        stream      = &ctx.videoStream;
        codecCtx    = &ctx.videoCodecCtx;
        streamIndex = &ctx.videoStreamIndex;
        codec       = &ctx.videoCodec;
        codecId     = formatCtx->oformat->video_codec;
    }
    else if (type == AVMEDIA_TYPE_AUDIO)
    {
        stream      = &ctx.audioStream;
        codecCtx    = &ctx.audioCodecCtx;
        streamIndex = &ctx.audioStreamIndex;
        codec       = &ctx.audioCodec;
        codecId     = formatCtx->oformat->audio_codec;
    }
    else
    {
        cerr << "Unknown media type: " << av_get_media_type_string(type) << endl;
        return -1;
    }


    *codec = avcodec_find_encoder(codecId);
    if (!(*codec))
    {
        cerr << "Can't found codec" << endl;
        return -1;
    }

    *stream = avformat_new_stream(ctx.formatCtx, *codec);
    if (!*stream)
    {
        cerr << "Can't alloc stream" << endl;
        return -1;
    }

    *codecCtx = (*stream)->codec;

    avcodec_get_context_defaults3(*codecCtx, *codec);

    (*codecCtx)->codec_id = codecId;

    if (formatCtx->oformat->flags & AVFMT_GLOBALHEADER)
        (*codecCtx)->flags |= CODEC_FLAG_GLOBAL_HEADER;

    *streamIndex = formatCtx->nb_streams - 1;
    return *streamIndex;
}


static
AVFrame *allocVideoFrame(PixelFormat pixFmt, int width, int height)
{
    if (pixFmt == PIX_FMT_NONE || width <= 0 || height <= 0)
    {
        cerr << "Invalid picture params: format:" << pixFmt << ", w:" << width << ", h:" << height << endl;
        return 0;
    }

    AVFrame *outFrame = avcodec_alloc_frame();
    if (!outFrame)
        return outFrame;

    outFrame->width  = width;
    outFrame->height = height;
    outFrame->format = pixFmt;

    int size = avpicture_get_size(pixFmt, width, height);
    uint8_t *pictBuf = new uint8_t[size];

    avpicture_fill((AVPicture*)outFrame,
                   pictBuf,
                   pixFmt, width, height);

    return outFrame;
}


static
void freeVideoFrame(AVFrame **frame)
{
    if (frame == 0 || *frame == 0)
    {
        return;
    }

    AVFrame *ptr = *frame;

    delete [] ptr->data[0];

    avcodec_free_frame(frame);
    *frame = 0;
}


int main(int argc, char **argv)
{
    // register all formats, protocols and codecs
    av_register_all();

    int stat;

    if (argc < 3)
    {
        usage(argv[0]);
        return 1;
    }

    inCtx.resource  = argv[1];
    outCtx.resource = argv[2];

    ////
    //// Prepare input file
    ////

    // Open input file
    if (avformat_open_input(&inCtx.formatCtx, inCtx.resource.c_str(), 0, 0) < 0)
    {
        cerr << "Can't open source file: " << inCtx.resource << endl;
        return 1;
    }

    // Retrieve stream information
    if (avformat_find_stream_info(inCtx.formatCtx, 0) < 0)
    {
        cerr << "Can't find stream information" << endl;
        return 1;
    }

    // Walk via streams
    if (openCodecContex(inCtx, AVMEDIA_TYPE_VIDEO) >= 0)
    {
        // TODO
    }

    if (openCodecContex(inCtx, AVMEDIA_TYPE_AUDIO) >= 0)
    {
        // TODO
    }

    av_dump_format(inCtx.formatCtx, 0, inCtx.resource.c_str(), 0);




    ////
    //// Prepare output file
    ////

    // alloc context
    outCtx.formatCtx = avformat_alloc_context();

    outCtx.formatCtx->iformat = 0;
    outCtx.formatCtx->oformat = av_guess_format(0, outCtx.resource.c_str(), 0);

    if (!outCtx.formatCtx->oformat)
    {
        cout << "Fall back to MPEGTS format" << endl;
        outCtx.formatCtx->oformat = av_guess_format("mpegts", 0, 0);

        if (!outCtx.formatCtx->oformat)
        {
            cerr << "Can't found output format for file: " << outCtx.resource << endl;
            return 1;
        }
    }


    // Setup video stream
    if (inCtx.videoStreamIndex >= 0 && addStream(outCtx, AVMEDIA_TYPE_VIDEO) >= 0)
    {
        outCtx.videoCodecCtx->bit_rate = 400000;

        outCtx.videoCodecCtx->width    = inCtx.videoCodecCtx->width;
        outCtx.videoCodecCtx->height   = inCtx.videoCodecCtx->height;

        outCtx.videoCodecCtx->time_base.den = 25;
        outCtx.videoCodecCtx->time_base.num = 1;

        outCtx.videoCodecCtx->gop_size      = 12;
        outCtx.videoCodecCtx->pix_fmt       = PIX_FMT_YUV420P;

        if (outCtx.videoCodec->pix_fmts)
        {
            outCtx.videoCodecCtx->pix_fmt = *outCtx.videoCodec->pix_fmts;
        }

        if (outCtx.videoCodecCtx->codec_id == AV_CODEC_ID_MPEG2VIDEO)
        {
            outCtx.videoCodecCtx->max_b_frames = 2;
        }

        if (outCtx.videoCodecCtx->codec_id == AV_CODEC_ID_MPEG1VIDEO)
        {
            outCtx.videoCodecCtx->mb_decision = 2;
        }

        // Open codec
        if (avcodec_open2(outCtx.videoCodecCtx, outCtx.videoCodec, 0) < 0)
        {
            cerr << "Can't open output video codec" << endl;
            return 1;
        }

        if (outCtx.videoCodecCtx->pix_fmt != inCtx.videoCodecCtx->pix_fmt)
        {
            outCtx.swsCtx = sws_getContext(inCtx.videoCodecCtx->width,
                                           inCtx.videoCodecCtx->height,
                                           inCtx.videoCodecCtx->pix_fmt,
                                           outCtx.videoCodecCtx->width,
                                           outCtx.videoCodecCtx->height,
                                           outCtx.videoCodecCtx->pix_fmt,
                                           SWS_BICUBIC, 0, 0, 0);
            if (!outCtx.swsCtx)
            {
                cerr << "Can't initialize video conversion context" << endl;
                return 1;
            }
        }
    }


    // Setup audio stream
    if (inCtx.audioStreamIndex >= 0 && addStream(outCtx, AVMEDIA_TYPE_AUDIO) >= 0)
    {
        outCtx.audioStream->id = 1; // WHAT IS IT???

        outCtx.audioCodecCtx->sample_fmt     = inCtx.audioCodecCtx->sample_fmt;
        outCtx.audioCodecCtx->bit_rate       = inCtx.audioCodecCtx->bit_rate;
        outCtx.audioCodecCtx->sample_rate    = inCtx.audioCodecCtx->sample_rate;
        outCtx.audioCodecCtx->channels       = inCtx.audioCodecCtx->channels;
        outCtx.audioCodecCtx->channel_layout = inCtx.audioCodecCtx->channel_layout;

        // Open codec
        if (avcodec_open2(outCtx.audioCodecCtx, outCtx.audioCodec, 0) < 0)
        {
            cerr << "Can't open output audio codec" << endl;
            return 1;
        }
    }

    av_dump_format(outCtx.formatCtx, 0, outCtx.resource.c_str(), 1);

    stat = avio_open2(&outCtx.formatCtx->pb, outCtx.resource.c_str(), AVIO_FLAG_WRITE, 0, 0);
    if (stat < 0)
    {
        cerr << "Can't open output file: " << outCtx.resource << endl;
        return 1;
    }

    // write header
    if (avformat_write_header(outCtx.formatCtx, 0) < 0)
    {
        cerr << "Can't write header to output file" << endl;
        return 1;
    }


    ////
    //// Transcode
    ////

    cout << "streams: " << inCtx.videoStreamIndex << "/" << inCtx.audioStreamIndex << endl;
    cout << "streams: " << outCtx.videoStreamIndex << "/" << outCtx.audioStreamIndex << endl;

    AVFrame *frame      = avcodec_alloc_frame();
    AVFrame *videoFrame = 0;

    if (outCtx.swsCtx)
    {
        videoFrame = allocVideoFrame(outCtx.videoCodecCtx->pix_fmt,
                                     outCtx.videoCodecCtx->width,
                                     outCtx.videoCodecCtx->height);
    }


    AVPacket pkt;
    av_init_packet(&pkt);
    pkt.data = 0;
    pkt.size = 0;

    int gotFrame;
    uint64_t inputAudioSamples = 0;

    while (av_read_frame(inCtx.formatCtx, &pkt) >= 0)
    {
        stat = -1;
        AVRational outPktTimeBase;
        AVPacket outPkt;
        int      gotPacket = 0;
        av_init_packet(&outPkt);
        outPkt.data = 0;
        outPkt.size = 0;

        if (pkt.stream_index == inCtx.videoStreamIndex)
        {
            stat = avcodec_decode_video2(inCtx.videoCodecCtx, frame, &gotFrame, &pkt);
            if (stat < 0)
            {
                cerr << "Can't decode video frame" << endl;
                return 1;
            }

            if (gotFrame)
            {
                cout << "video_frame coded_n:" << frame->coded_picture_number
                     << " pts:" << av_ts2timestr(frame->pts, &inCtx.videoCodecCtx->time_base)
                     << endl;


                if (outCtx.swsCtx)
                {
                    sws_scale(outCtx.swsCtx,
                              (const uint8_t * const *)frame->data, frame->linesize, 0, inCtx.videoCodecCtx->height,
                              videoFrame->data, videoFrame->linesize);
                }
                else
                {
                    videoFrame = frame;
                }

                // Encode
                stat = avcodec_encode_video2(outCtx.videoCodecCtx, &outPkt, videoFrame, &gotPacket);

                if (stat >= 0 && gotPacket)
                {
                    if (outCtx.videoCodecCtx->coded_frame->pts != AV_NOPTS_VALUE)
                    {
                        outPkt.pts = av_rescale_q(outCtx.videoCodecCtx->coded_frame->pts,
                                                  outCtx.videoCodecCtx->time_base,
                                                  outCtx.videoStream->time_base);
                    }

                    if (outCtx.videoCodecCtx->coded_frame->key_frame)
                    {
                        outPkt.flags |= AV_PKT_FLAG_KEY;
                    }

                    outPkt.stream_index = outCtx.videoStreamIndex;
                    outPktTimeBase = outCtx.videoStream->time_base;
                }
            }
        }
        else if (pkt.stream_index == inCtx.audioStreamIndex)
        {
            stat = avcodec_decode_audio4(inCtx.audioCodecCtx, frame, &gotFrame, &pkt);
            if (stat < 0)
            {
                cerr << "Can't decode audio frame" << endl;
                return 1;
            }

            if (gotFrame)
            {
                if (frame->pts == AV_NOPTS_VALUE)
                {
                    AVRational samplesRateInv = {1, outCtx.audioCodecCtx->sample_rate};
                    int64_t pts = inputAudioSamples;
                    frame->pts = av_rescale_q(pts, samplesRateInv, inCtx.audioCodecCtx->time_base);
                }

                inputAudioSamples += frame->nb_samples;

                cout << "audio_frame nb_samples:" << frame->nb_samples
                     << " pts:" << av_ts2timestr(frame->pts, &inCtx.audioCodecCtx->time_base)
                     << " timebase: " << inCtx.audioCodecCtx->time_base.num << "/" << inCtx.audioCodecCtx->time_base.den
                     << endl;

                // Encode
                stat = avcodec_encode_audio2(outCtx.audioCodecCtx, &outPkt, frame, &gotPacket);
                if (stat >= 0 && gotPacket)
                {
                    outPkt.stream_index = outCtx.audioStreamIndex;
                    outPktTimeBase = outCtx.audioStream->time_base;
                }
            }

        }

        // Muxing
        if (stat >= 0 && gotPacket && outPkt.data && outPkt.size > 0)
        {
            cout << "write frame st:" << outPkt.stream_index
                 << " pts:" << av_ts2timestr(outPkt.pts, &outPktTimeBase)
                 << endl;
            stat = av_interleaved_write_frame(outCtx.formatCtx, &outPkt);
            if (stat < 0)
            {
                cerr << "Error while writing frame for stream: " << outPkt.stream_index << endl;
                return 1;
            }
        }
    }

    // write trailer
    av_write_trailer(outCtx.formatCtx);

    //
    // Free resources
    //

    // input
    if (inCtx.videoStream)
        avcodec_close(inCtx.videoStream->codec);
    if (inCtx.audioStream)
        avcodec_close(inCtx.audioStream->codec);

    avformat_close_input(&inCtx.formatCtx);

    // output
    if (outCtx.videoStream)
        avcodec_close(outCtx.videoStream->codec);
    if (outCtx.audioStream)
        avcodec_close(outCtx.audioStream->codec);

    for (uint i = 0; i < outCtx.formatCtx->nb_streams; ++i)
    {
        av_freep(&outCtx.formatCtx->streams[i]->codec);
        av_freep(&outCtx.formatCtx->streams[i]);
    }

    if (outCtx.swsCtx)
    {
        sws_freeContext(outCtx.swsCtx);
        freeVideoFrame(&videoFrame);
    }

    avio_close(outCtx.formatCtx->pb);

    av_free(outCtx.formatCtx);


    return 0;
}