/*
 *******************************************************************************
 *                                   INCLUDE
 *******************************************************************************
 */
#include "Muxer.h"
#include "Util.h"
#include "Logger.h"
#include <Poco/Format.h>

namespace Video
{
namespace FFmpeg
{
/*
 *******************************************************************************
 *                                CLASS DEFINE
 *******************************************************************************
 */
Muxer::Stream::Stream(AVFormatContext& format, Core::VideoEncoder::Ptr pEncoder, AVRational tb,
                      Core::PacketQueue::Ptr pQueue)
    : _pQueue(pQueue), _timebase(tb), _timeline(0), _pEncoder(nullptr)
{
    open(format, pEncoder->getCodecID(),
         [&](AVCodecContext* pCodecCtx, AVStream* pStream) -> void
         {
             int fps                 = pEncoder->getFrameRate();
             pCodecCtx->width        = pEncoder->getWidth();
             pCodecCtx->height       = pEncoder->getHeight();
             pCodecCtx->time_base    = {1, fps};
             pCodecCtx->framerate    = {fps, 1};
             pCodecCtx->gop_size     = pEncoder->getGOP();
             pCodecCtx->max_b_frames = 0;
             pCodecCtx->pix_fmt      = AV_PIX_FMT_YUV420P;
             pCodecCtx->bit_rate     = pEncoder->getBitrate();

             pStream->time_base    = pCodecCtx->time_base;
             pStream->r_frame_rate = pCodecCtx->framerate;
         });
}

Muxer::Stream::Stream(AVFormatContext& format, Core::AudioEncoder::Ptr pEncoder, AVRational tb,
                      Core::PacketQueue::Ptr pQueue)
    : _pQueue(pQueue), _timebase(tb), _timeline(0), _pEncoder(nullptr)
{
    const Core::CodecID codecID = pEncoder->getCodecID();
    open(format, codecID,
         [&](AVCodecContext* pCodecCtx, AVStream* pStream) -> void
         {
             switch (codecID)
             {
             case Core::CodecID::G711A:
             case Core::CodecID::G711U:
                 pCodecCtx->sample_fmt = AV_SAMPLE_FMT_S16;
                 pCodecCtx->ch_layout  = AV_CHANNEL_LAYOUT_MONO;
                 break;

             case Core::CodecID::MP3:
             case Core::CodecID::AAC:
                 pCodecCtx->sample_fmt = AV_SAMPLE_FMT_FLTP;
                 pCodecCtx->ch_layout  = AV_CHANNEL_LAYOUT_STEREO;
                 break;

             case Core::CodecID::OPUS:
                 pCodecCtx->sample_fmt = AV_SAMPLE_FMT_S16;
                 pCodecCtx->ch_layout  = AV_CHANNEL_LAYOUT_STEREO;
                 break;

             default:
                 throw Poco::NotImplementedException(
                     Poco::format("Unsupported codec: %s", Core::codecIDToString(codecID)), -EINVAL);
             }

             pCodecCtx->sample_rate = pEncoder->getSampleRate();
             pCodecCtx->bit_rate    = pEncoder->getBitrate();
             pStream->time_base     = (AVRational){1, pCodecCtx->sample_rate};
         });
}

Muxer::Stream::~Stream()
{
    close();
}

void Muxer::Stream::open(AVFormatContext& format, const Core::CodecID codec,
                         std::function<void(AVCodecContext*, AVStream*)> init)
{
    _pEncoder = nullptr;

    const AVCodec* pCodec = Util::getEncoder(codec);
    if (!pCodec)
        throw Poco::NotImplementedException(Poco::format("Unsupported codec: %s", Core::codecIDToString(codec)),
                                            -EINVAL);

    _pStream = avformat_new_stream(&format, pCodec);
    if (!_pStream)
        throw Poco::OutOfMemoryException("Alloc stream object FAIL", -ENOMEM);

    _pStream->id = format.nb_streams - 1;
    _pEncoder    = avcodec_alloc_context3(pCodec);
    if (!_pEncoder)
        throw Poco::NotImplementedException(Poco::format("Unsupported codec: %s", Core::codecIDToString(codec)),
                                            -EINVAL);

    _pEncoder->codec_id = pCodec->id;
    init(_pEncoder, _pStream);

    if (format.oformat->flags & AVFMT_GLOBALHEADER)
        _pEncoder->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

    int rc = avcodec_open2(_pEncoder, pCodec, NULL);
    if (rc < 0)
        throw Poco::RuntimeException(
            Poco::format("Open codec(%s) FAIL: %d(%s)", Core::codecIDToString(codec), rc, Util::getErrorString(rc)),
            rc);

    avcodec_parameters_from_context(_pStream->codecpar, _pEncoder);
}

void Muxer::Stream::close()
{
    avcodec_free_context(&_pEncoder);
}

/*
 *******************************************************************************
 *                                CLASS DEFINE
 *******************************************************************************
 */
Muxer::Muxer(const std::string& url): Core::Muxer(Poco::format("FFmpeg::Muxer(%s)", url), url)
{
}

Muxer::~Muxer()
{
}

void Muxer::connect()
{
    const std::string format = getFormatName();
    int               rc =
        avformat_alloc_output_context2(&_pFormatCtx, nullptr, format.empty() ? nullptr : format.c_str(), url().c_str());
    if (rc < 0)
        throw Poco::NotFoundException(Poco::format("%s: Create output format FAIL: %s: %s", name(), scheme(), format),
                                      -EINVAL);

    AVRational vtb, atb;
    getSuitableTimebase(vtb, atb);

    previousComponents().foreach (
        [&](Core::Component::Ptr pPrevious) -> void
        {
            auto pVideoEncoder = pPrevious.cast<Core::VideoEncoder>();
            if (!pVideoEncoder.isNull())
            {
                _streams.push_back(new Stream(*_pFormatCtx, pVideoEncoder, vtb, packetQueues().find(pPrevious)));
                return;
            }

            auto pAudioEncoder = pPrevious.cast<Core::AudioEncoder>();
            if (!pAudioEncoder.isNull())
            {
                _streams.push_back(new Stream(*_pFormatCtx, pAudioEncoder, atb, packetQueues().find(pPrevious)));
            }
        });

    if (!(_pFormatCtx->oformat->flags & AVFMT_NOFILE))
    {
        AVDictionary* pOptions = nullptr;
        av_dict_set(&pOptions, "rw_timeout", "10000", 0);
        if ("udp" == scheme())
        {
            av_dict_set(&pOptions, "pkt_size", "1316", 0);
        }

        rc = avio_open2(&_pFormatCtx->pb, url().c_str(), AVIO_FLAG_WRITE, nullptr, &pOptions);
        if (rc < 0)
            throw Poco::RuntimeException(
                Poco::format("%s: Open: '%s' FAIL: %s", name(), url(), Util::getErrorString(rc)), rc);
    }

    AVDictionary* pOptions = NULL;
    if ("rtsp" == scheme())
    {
        av_dict_set(&pOptions, "rtsp_transport", "tcp", 0);
        av_dict_set(&pOptions, "buffer_size", "1024000", 0);
    }
    rc = avformat_write_header(_pFormatCtx, &pOptions);
    if (rc < 0)
        throw Poco::RuntimeException(
            Poco::format("%s: Open output: '%s' FAIL: %s", name(), url(), Util::getErrorString(rc)), rc);
}

void Muxer::mux()
{
    static const AVRational timebase = {1, 1000000};

    std::sort(_streams.begin(), _streams.end(), [](auto a, auto b) { return *a < *b; });
    auto& pStream = *_streams.begin();

    Core::Packet::Ptr pPacket = pStream->queue()->dequeue(4000);
    if (pStream->serialCodeWatcher().isChanged(pPacket->serialCode()))
        throw Poco::ApplicationException(Poco::format("%s: Stream changed, restart...", name()));

    AVPacket* pAVPacket;
    pPacket->moveTo(&pAVPacket);
    pAVPacket->stream_index = pStream->id();

    pStream->setTimeline(pAVPacket->pts);

    av_packet_rescale_ts(pAVPacket, timebase, pStream->timebase());

    int rc = av_interleaved_write_frame(_pFormatCtx, pAVPacket);
    av_packet_free(&pAVPacket);

    if (rc < 0)
    {
        logger().warning("%s: mux FAIL: %d (%s)", name(), rc, Util::getErrorString(rc));

        if (rc != -22)
            throw Poco::RuntimeException(Poco::format("%s: mux FAIL: %d (%s)", name(), rc, Util::getErrorString(rc)),
                                         rc);
    }
}

void Muxer::disconnect()
{
    if (_pFormatCtx && _pFormatCtx->pb)
    {
        av_write_trailer(_pFormatCtx);
        if (!(_pFormatCtx->oformat->flags & AVFMT_NOFILE))
            avio_close(_pFormatCtx->pb);
    }

    _streams.clear();
    avformat_free_context(_pFormatCtx);
}

const std::string Muxer::getFormatName()
{
    const std::string s = scheme();
    if (s == "rtmp")
        return "flv";
    if (s == "rtsp")
        return "rtsp";
    if (s == "udp")
        return "mpegts";
    return "";
}

void Muxer::getSuitableTimebase(AVRational& vtb, AVRational& atb)
{
    const std::string s = scheme();
    if ("rtmp" == s)
    {
        atb = {1, 1000};
        vtb = {1, 1000};
    }
    else if ("rtsp" == s)
    {
        previousComponents().foreach (
            [&](Core::Component::Ptr pPrevious) -> void
            {
                auto pEncoder = pPrevious.cast<Core::AudioEncoder>();
                if (!pEncoder.isNull())
                {
                    switch (pEncoder->getCodecID())
                    {
                    case Core::CodecID::AAC:
                    case Core::CodecID::OPUS:
                        atb = {1, pEncoder->getSampleRate()};
                        break;

                    case Core::CodecID::G711A:
                    case Core::CodecID::G711U:
                        atb = {1, 8000};
                        break;

                    default:
                        atb = {1, 90000};
                        break;
                    }
                }
            });

        vtb = {1, 90000};
    }
    else if ("udp" == s)
    {
        atb = {1, 90000};
        vtb = {1, 90000};
    }
}
}
}