/*
 *******************************************************************************
 *                                   INCLUDE
 *******************************************************************************
 */
#include "DemuxHandler.h"
#include "Util.h"
#include "Logger.h"
#include "../core/System.h"
#include "util/Math.h"
#include <list>

namespace Video
{
namespace FFmpeg
{
/*
 *******************************************************************************
 *                                CLASS DEFINE
 *******************************************************************************
 */
DemuxHandler::DemuxHandler(const std::string& uri)
    : Core::DemuxHandler(Poco::format("FFmpeg::DemuxHandler(%s)", uri), uri)
{
}

DemuxHandler::~DemuxHandler()
{
}

void DemuxHandler::run(Core::VideoEncoder::Ptr pVideoEncoder, Core::AudioEncoder::Ptr pAudioEncoder, const Post& post,
                       std::atomic<bool>& isCancelled)
{
    _pFormatCtx                              = avformat_alloc_context();
    _pFormatCtx->interrupt_callback.callback = DemuxHandler::interruptCallback;
    _pFormatCtx->interrupt_callback.opaque   = this;

    std::string url          = uri();
    bool        isFileScheme = (scheme() == "file");
    if (isFileScheme)
        url = uri().substr(std::string("file:").size());

    AVDictionary* pOptions = nullptr;
    if (scheme() == "rtsp" && getUsingTCPTransport())
        av_dict_set(&pOptions, "rtsp_transport", "tcp", 0);

    av_dict_set(&pOptions, "stimeout", "2000000", 0);

    setAnchorTime(Core::System::timestamp());
    int rc = avformat_open_input(&_pFormatCtx, url.c_str(), nullptr, &pOptions);
    if (rc < 0)
        throw Poco::InvalidAccessException(
            Poco::format("%s: Open input FAIL: %d(%s)", name(), rc, Util::getErrorString(rc)), rc);

    if (!isFileScheme)
    {
        _pFormatCtx->probesize            = 256 * 1024;
        _pFormatCtx->max_analyze_duration = 300 * 1000;
    }
    logger().debug("avformat_find_stream_info enter");
    rc = avformat_find_stream_info(_pFormatCtx, nullptr);
    logger().debug("avformat_find_stream_info  exit");
    if (rc < 0)
    {
        avformat_close_input(&_pFormatCtx);
        throw Poco::NotFoundException(
            Poco::format("%s: Find stream FAIL: %d(%s)", name(), rc, Util::getErrorString(rc)), rc);
    }

    av_dump_format(_pFormatCtx, 0, url.c_str(), 0);

    int             vIndex = -1;
    AVRational      vTimebase;
    AVCodecContext* pVidecCodecCtx = nullptr;
    if (openCodec(AVMEDIA_TYPE_VIDEO, vIndex, &pVidecCodecCtx) >= 0)
    {
        auto pStream = _pFormatCtx->streams[vIndex];
        int  fr      = av_q2d(pStream->r_frame_rate);

        vTimebase = pStream->time_base;
        pVideoEncoder->setCodecID(Util::getCodecID(pVidecCodecCtx->codec_id));
        pVideoEncoder->setWidth(pVidecCodecCtx->width);
        pVideoEncoder->setHeight(pVidecCodecCtx->height);
        pVideoEncoder->setFrameRate(fr);
        pVideoEncoder->setGOP(pVidecCodecCtx->gop_size);
        pVideoEncoder->setBitrate(pVidecCodecCtx->bit_rate);
    }
    pVideoEncoder->setSignalLost(vIndex < 0);

    int             aIndex = -1;
    AVRational      aTimebase;
    AVCodecContext* pAudioCodecCtx = nullptr;
    if (openCodec(AVMEDIA_TYPE_AUDIO, aIndex, &pAudioCodecCtx) >= 0)
    {
        auto pStream = _pFormatCtx->streams[aIndex];
        logger().information("%s: Audio found: %s, %dHz, frame-size: %d, %dKbps", name(),
                             Util::getCodecIDString(pAudioCodecCtx->codec_id), pAudioCodecCtx->sample_rate,
                             pAudioCodecCtx->frame_size, (int)((float)pAudioCodecCtx->bit_rate / 1000));

        aTimebase = pStream->time_base;
        pAudioEncoder->setCodecID(Util::getCodecID(pAudioCodecCtx->codec_id));
        pAudioEncoder->setSampleFormat(Util::getSampleFormat(pAudioCodecCtx->sample_fmt));
        pAudioEncoder->setBitrate(pAudioCodecCtx->bit_rate);
        pAudioEncoder->setFrameSize(pAudioCodecCtx->frame_size);
        pAudioEncoder->setChannelLayout(Util::getChannelLayout(pAudioCodecCtx->ch_layout));
        pAudioEncoder->setSampleRate(pAudioCodecCtx->sample_rate);
    }
    pAudioEncoder->setSignalLost(aIndex < 0);
    pAudioEncoder->commit();

    if (vIndex < 0 && aIndex < 0)
    {
        avformat_close_input(&_pFormatCtx);
        throw Poco::NotFoundException(Poco::format("%s: Could NOT find any stream", name()), -ENOSYS);
    }

    int                primaryStream      = vIndex >= 0 ? vIndex : aIndex;
    AVBSFContext*      pAVBSFContext      = nullptr;
    AVBitStreamFilter* pAVBitStreamFilter = nullptr;
    if (vIndex >= 0)
    {
        if (pVidecCodecCtx->codec_id == AV_CODEC_ID_H264)
            pAVBitStreamFilter = (AVBitStreamFilter*)av_bsf_get_by_name("h264_mp4toannexb");
        else if (pVidecCodecCtx->codec_id == AV_CODEC_ID_H265)
            pAVBitStreamFilter = (AVBitStreamFilter*)av_bsf_get_by_name("hevc_mp4toannexb");

        if (pAVBitStreamFilter)
        {
            av_bsf_alloc(pAVBitStreamFilter, &pAVBSFContext);
            avcodec_parameters_copy(pAVBSFContext->par_in, _pFormatCtx->streams[vIndex]->codecpar);
            av_bsf_init(pAVBSFContext);
        }

        if (!pAVBSFContext)
        {
            avcodec_free_context(&pVidecCodecCtx);
            avcodec_free_context(&pAudioCodecCtx);
            avformat_close_input(&_pFormatCtx);
            throw Poco::NullValueException(Poco::format("%s: Init 'AVBSFContext' FAIL", name()), -ENOSYS);
        }
    }

    if (_seekReqeust)
    {
        if (isFileScheme && _seekTo < _pFormatCtx->duration)
        {
            rc = avformat_seek_file(_pFormatCtx, -1, INT64_MIN, _seekTo, INT64_MAX, AVSEEK_FLAG_FRAME);
            if (rc < 0)
                logger().error("%s: Seek file FAIL: %d(%s)", name(), rc, Util::getErrorString(rc));
            else
                logger().debug("%s: Seek file SUCCESS", name());
        }

        _seekReqeust = false;
    }

    AVPacket* pPacket = nullptr;

    pVideoEncoder->commit();
    logger().information("%s: Video found: %s, %dx%d@%d, gop: %d, %dKbps", name(),
                         Util::getCodecIDString(pVidecCodecCtx->codec_id), pVidecCodecCtx->width,
                         pVidecCodecCtx->height, pVideoEncoder->getFrameRate(), pVidecCodecCtx->gop_size,
                         (int)((float)pVidecCodecCtx->bit_rate / 1000));

    auto postPacket = [&](AVPacket*& pAVPacket) -> void
    {
        if (pAVPacket->stream_index == primaryStream)
            setPosition(
                av_rescale_q(pAVPacket->pts, _pFormatCtx->streams[pAVPacket->stream_index]->time_base, {1, 1000000}));

        if (pAVPacket->stream_index == vIndex)
        {
            rc = av_bsf_send_packet(pAVBSFContext, pAVPacket);
            av_packet_unref(pAVPacket);
            if (rc < 0)
            {
                logger().error("%s: Send BSF packet FAIL: %d(%s)", name(), rc, Util::getErrorString(rc));
                return;
            }

            while (av_bsf_receive_packet(pAVBSFContext, pAVPacket) >= 0)
            {
                adjustTimestamp(pAVPacket, vTimebase);
                updateTimeline(pAVPacket->pts);
                post(new Core::Packet(pVideoEncoder->code(), true, &pAVPacket), Core::MediaType::VIDEO);
            }
        }
        else if (pAVPacket->stream_index == aIndex)
        {
            adjustTimestamp(pAVPacket, aTimebase);
            updateTimeline(pAVPacket->pts);
            post(new Core::Packet(pAudioEncoder->code(), false, &pAVPacket), Core::MediaType::AUDIO);
        }
        else
            av_packet_unref(pAVPacket);
    };

    while (!isCancelled)
    {
        setAnchorTime(Core::System::timestamp());

        if (!pPacket)
            pPacket = av_packet_alloc();

        if (_seekReqeust)
        {
            if (isFileScheme && _seekTo < _pFormatCtx->duration)
            {
                rc = avformat_seek_file(_pFormatCtx, -1, INT64_MIN, _seekTo, INT64_MAX, AVSEEK_FLAG_FRAME);
                if (rc < 0)
                    logger().error("%s: Seek file FAIL: %d(%s)", name(), rc, Util::getErrorString(rc));
                else
                {
                    logger().debug("%s: Seek file SUCCESS", name());
                    setBaseline(getTimeline());
                    resetTimeDrift();
                }
            }

            _seekReqeust = false;
        }

        rc = av_read_frame(_pFormatCtx, pPacket);
        if (rc < 0)
        {
            logger().error("%s: Read FAIL: %d(%s)", name(), rc, Util::getErrorString(rc));
            break;
        }

        postPacket(pPacket);
    }

    av_bsf_free(&pAVBSFContext);
    av_packet_free(&pPacket);

    avcodec_free_context(&pVidecCodecCtx);
    avcodec_free_context(&pAudioCodecCtx);
    avformat_close_input(&_pFormatCtx);
}

int DemuxHandler::openCodec(AVMediaType mediaType, int& index, AVCodecContext** pCodecCtx)
{
    index = -1;

    int rc = av_find_best_stream(_pFormatCtx, mediaType, -1, -1, NULL, 0);
    if (rc < 0)
    {
        logger().warning("%s: Could not find %s stream in input: %d(%s)", name(), Util::getMediaTypeString(mediaType),
                         rc, Util::getErrorString(rc));
        return rc;
    }

    index                  = rc;
    AVStream*      pStream = _pFormatCtx->streams[index];
    const AVCodec* pCodec  = avcodec_find_decoder(pStream->codecpar->codec_id);
    if (!pCodec)
    {
        logger().error("%s: Failed to find %s codec", name(), Util::getMediaTypeString(mediaType));
        return AVERROR(EINVAL);
    }

    *pCodecCtx = avcodec_alloc_context3(pCodec);
    if (!*pCodecCtx)
    {
        logger().error("%s: Failed to allocate the %s codec context", name(), Util::getMediaTypeString(mediaType));
        return AVERROR(ENOMEM);
    }

    if ((rc = avcodec_parameters_to_context(*pCodecCtx, pStream->codecpar)) < 0)
    {
        logger().error("%s: Failed to copy %s codec parameters to decoder context: %d(%s)", name(),
                       Util::getMediaTypeString(mediaType), rc, Util::getErrorString(rc));
        return rc;
    }

    if ((rc = avcodec_open2(*pCodecCtx, pCodec, NULL)) < 0)
    {
        logger().error("%s: Failed to open %s codec: %d(%s)", name(), Util::getMediaTypeString(mediaType), rc,
                       Util::getErrorString(rc));
        return rc;
    }

    return 0;
}

void DemuxHandler::adjustTimestamp(AVPacket* pPacket, AVRational& timebase)
{
    av_packet_rescale_ts(pPacket, timebase, {1, 1000000});
    auto drift = getTimeDrift(pPacket->pts);
    pPacket->pts += drift;
    pPacket->dts += drift;
}

void DemuxHandler::updateTimeline(uint64_t ts)
{
    if (getTimeline() < ts)
        setTimeline(ts);
}

void DemuxHandler::seek(uint64_t ts)
{
    _seekReqeust = true;
    _seekTo      = ts;
}

int DemuxHandler::interruptCallback(void* pCtx)
{
    return 0;
}
}
}