/*
 *******************************************************************************
 *                                   INCLUDE
 *******************************************************************************
 */
#include "RTMPMuxer.h"
#include "../core/VideoEncoder.h"
#include "../core/AudioEncoder.h"
#include "../core/SampleFormat.h"
#include "../core/System.h"
#include <media-server/libflv/flv-proto.h>

namespace Video
{
namespace MediaServer
{
/*
 *******************************************************************************
 *                                   TYPEDEFS
 *******************************************************************************
 */
struct CodecIDItem
{
    Core::CodecID id;
    int           type;
};

/*
 *******************************************************************************
 *                                  LOCAL VARIABLES
 *******************************************************************************
 */
static const std::array<CodecIDItem, 7> codecMap = {{
    {Core::CodecID::H264, FLV_VIDEO_H264},
    {Core::CodecID::H265, FLV_VIDEO_H265},

    {Core::CodecID::MP3, FLV_AUDIO_MP3},
    {Core::CodecID::AAC, FLV_AUDIO_AAC},
    {Core::CodecID::OPUS, FLV_AUDIO_OPUS},
    {Core::CodecID::G711A, FLV_AUDIO_G711A},
    {Core::CodecID::G711U, FLV_AUDIO_G711U},
}};

/*
 *******************************************************************************
 *                                CLASS DEFINE
 *******************************************************************************
 */
RTMPMuxer::Stream::Stream(const std::string& name, Core::CodecID codecID, Core::PacketQueue::Ptr pQueue)
    : _name(name), _codecID(codecID), _pQueue(pQueue)
{
    switch (codecID)
    {
    case Core::CodecID::H264:
        _pMuxerMethod = flv_muxer_avc;
        break;

    case Core::CodecID::H265:
        _pMuxerMethod = flv_muxer_hevc;
        break;

    case Core::CodecID::MP3:
        _pMuxerMethod = flv_muxer_mp3;
        break;

    case Core::CodecID::AAC:
        _pMuxerMethod = flv_muxer_aac;
        break;

    case Core::CodecID::OPUS:
        _pMuxerMethod = flv_muxer_opus;
        break;

    case Core::CodecID::G711A:
        _pMuxerMethod = flv_muxer_g711a;
        break;

    case Core::CodecID::G711U:
        _pMuxerMethod = flv_muxer_g711u;
        break;

    default:
        _pMuxerMethod = nullptr;
        throw Poco::NotImplementedException(
            Poco::format("%s: Unsupported codec: %s", name, Core::codecIDToString(codecID)));
    }
}

RTMPMuxer::Stream::~Stream()
{
}

inline int Video::MediaServer::RTMPMuxer::Stream::muxer(flv_muxer_t* muxer, const void* data, size_t bytes,
                                                        uint32_t pts, uint32_t dts)
{
    if (!_pMuxerMethod)
        throw Poco::BugcheckException(Poco::format("%s: muxer method is NULL!", name()));

    return _pMuxerMethod(muxer, data, bytes, pts, dts);
}

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

RTMPMuxer::~RTMPMuxer()
{
}

void RTMPMuxer::connect()
{
    Poco::URI uri(url());

    std::string scheme = uri.getScheme();
    std::string host   = uri.getHost();
    short       port   = uri.getPort();
    if (port <= 0)
        port = 1935;

    std::string path = uri.getPath();

    if (!path.empty() && path[0] == '/')
        path = path.substr(1);

    size_t pos = path.find_first_of('/');
    if (pos == std::string::npos)
        throw Poco::InvalidArgumentException(Poco::format("%s: Invalid URL: %s", name(), url()));

    _socket = socket_connect_host(host.c_str(), port, 2000);
    if (socket_invalid == _socket)
    {
        throw Poco::InvalidAccessException(Poco::format("%s: Connect FAIL", name()));
    }

    socket_setnonblock(_socket, 0);

    std::string turl     = scheme + "://" + host + ":" + std::to_string(port) + "/" + path.substr(0, pos);
    std::string appname  = path.substr(0, pos);
    std::string playpath = path.substr(pos + 1);

    struct rtmp_client_handler_t handler = {};
    handler.send                         = RTMPMuxer::clientSend;
    _pClient = rtmp_client_create(appname.c_str(), playpath.c_str(), turl.c_str(), this, &handler);
    if (!_pClient)
        throw Poco::InvalidAccessException(Poco::format("%s: Create RTMP-Client FAIL!", name()));

    _pMuxer = flv_muxer_create(RTMPMuxer::muxerHandler, this);
    if (!_pMuxer)
        throw Poco::InvalidAccessException(Poco::format("%s: Create RTMP-Muxer FAIL!", name()));

    int rc = rtmp_client_start(_pClient, 0);
    if (rc != 0)
        throw Poco::InvalidAccessException(Poco::format("%s: Start RTMP-Client FAIL: %d", name(), rc));

    int     limit = 50;
    uint8_t buffer[4096];
    while (--limit > 0 && 4 != rtmp_client_getstate(_pClient) &&
           (rc = socket_recv(_socket, buffer, sizeof(buffer), 0)) > 0)
    {
        if (rtmp_client_input(_pClient, buffer, rc) != 0)
            throw Poco::InvalidAccessException(Poco::format("%s: RTMP-Client input handle FAIL!", name()));
    }

    struct flv_metadata_t metadata     = {};
    uint8_t               opusHead[19] = {0};
    bool                  hasOpus      = false;
    previousComponents().foreach (
        [&](Core::Component::Ptr pPrevious) -> void
        {
            auto pVideoEncoder = pPrevious.cast<Core::VideoEncoder>();
            if (!pVideoEncoder.isNull())
            {
                _streams.push_back(new Stream(Poco::format("%s/Stream(%s)", name(), pVideoEncoder->name()),
                                              pVideoEncoder->getCodecID(), packetQueues().find(pPrevious)));

                metadata.videocodecid  = getCodecID(pVideoEncoder->getCodecID());
                metadata.videodatarate = pVideoEncoder->getBitrate() / 1000;
                metadata.framerate     = pVideoEncoder->getFrameRate();
                metadata.width         = pVideoEncoder->getWidth();
                metadata.height        = pVideoEncoder->getHeight();
                return;
            }

            auto pAudioEncoder = pPrevious.cast<Core::AudioEncoder>();
            if (!pAudioEncoder.isNull())
            {
                _streams.push_back(new Stream(Poco::format("%s/Stream(%s)", name(), pAudioEncoder->name()),
                                              pAudioEncoder->getCodecID(), packetQueues().find(pPrevious)));

                metadata.audiocodecid = getCodecID(pAudioEncoder->getCodecID());
                metadata.audiocodecid >>= 4;
                metadata.audiosamplerate = pAudioEncoder->getSampleRate();
                metadata.audiodatarate   = pAudioEncoder->getBitrate() / 1000;
                metadata.audiosamplesize = Core::sampleFormatBitWidth(pAudioEncoder->getSampleFormat());
                metadata.stereo = pAudioEncoder->getChannelLayout() == Core::ChannelLayout::STEREO ? true : false;

                if (pAudioEncoder->getCodecID() == Core::CodecID::OPUS)
                {
                    memcpy(opusHead, "OpusHead", 8);
                    opusHead[8]  = 1;
                    opusHead[9]  = (uint8_t)(metadata.stereo ? 2 : 1);
                    uint32_t sr  = (uint32_t)metadata.audiosamplerate;
                    opusHead[12] = (uint8_t)(sr & 0xFF);
                    opusHead[13] = (uint8_t)((sr >> 8) & 0xFF);
                    opusHead[14] = (uint8_t)((sr >> 16) & 0xFF);
                    opusHead[15] = (uint8_t)((sr >> 24) & 0xFF);
                    hasOpus      = true;
                }
            }
        });

    flv_muxer_metadata(_pMuxer, &metadata);

    if (hasOpus)
    {
        for (auto& pStream : _streams)
        {
            if (pStream->codecID() == Core::CodecID::OPUS)
            {
                pStream->muxer(_pMuxer, opusHead, 19, 0, 0);
                break;
            }
        }
    }
}

void RTMPMuxer::mux()
{
    drainInput();

    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()));

    pStream->setTimeline(pPacket->pts());
    int rc = pStream->muxer(_pMuxer, pPacket->data(), pPacket->length(), pPacket->pts() / 1000, pPacket->dts() / 1000);
    if (rc < 0)
        throw Poco::ProtocolException(Poco::format("%s: mux %s FAIL: %d (len:%u, pts:%Lu, dts:%Lu)", name(),
                                                   (pPacket->isVideo() ? std::string("video") : std::string("audio")),
                                                   rc, pPacket->length(), pPacket->pts(), pPacket->dts()));

#if (TRACE_VIDEO_FRAME_RATE)
    if (pPacket->isVideo())
    {
        count++;
        uint64_t now   = Core::System::timestamp();
        int64_t  delta = now - timeline;
        if (delta >= 1000000)
        {
            timeline = now;
            std::cout << " ==> " << name() << " fps: " << (double)count / (delta / 1000000) << std::endl;
            count = 0;
        }
    }
#endif
}

void RTMPMuxer::disconnect()
{
    if (_pClient)
    {
        rtmp_client_stop(_pClient);
        rtmp_client_destroy(_pClient);
        _pClient = nullptr;
    }

    Poco::Thread::sleep(100);
    if (_pMuxer)
    {
        flv_muxer_destroy(_pMuxer);
        _pMuxer = nullptr;
    }

    if (_socket != socket_invalid)
    {
        socket_close(_socket);
        _socket = socket_invalid;
    }

    _streams.clear();
}

void RTMPMuxer::drainInput()
{
    uint8_t buffer[2048];

    while (true)
    {
        int rc = socket_recv_by_time(_socket, buffer, sizeof(buffer), 0, 0);

        if (SOCKET_TIMEDOUT == rc)
            break;

        if (0 == rc)
            throw Poco::IOException(Poco::format("%s: Connection closed by peer", name()));

        if (rc < 0)
            throw Poco::IOException(Poco::format("%s: socket_recv FAIL: %d", name(), rc));

        if (rtmp_client_input(_pClient, buffer, (size_t)rc) != 0)
            throw Poco::ProtocolException(Poco::format("%s: RTMP-Client input handle FAIL!", name()));
    }
}

int RTMPMuxer::clientSend(void* param, const void* header, size_t len, const void* data, size_t bytes)
{
    RTMPMuxer* pPusher = (RTMPMuxer*)param;

    socket_bufvec_t vec[2];
    socket_setbufvec(vec, 0, (void*)header, len);
    socket_setbufvec(vec, 1, (void*)data, bytes);

#if defined(CORRUPT_RTMP_CHUNK_DATA)

    rtmp_corrupt_data(header, len);
    rtmp_corrupt_data(data, bytes);
#endif

    return socket_send_v_all_by_time(pPusher->_socket, vec, bytes > 0 ? 2 : 1, 0, 5000);
}

int RTMPMuxer::muxerHandler(void* param, int type, const void* data, size_t bytes, uint32_t timestamp)
{
    RTMPMuxer* pPusher = (RTMPMuxer*)param;
    int        rc;
    switch (type)
    {
    case FLV_TYPE_SCRIPT:
        rc = rtmp_client_push_script(pPusher->_pClient, data, bytes, timestamp);
        break;

    case FLV_TYPE_VIDEO:
        rc = rtmp_client_push_video(pPusher->_pClient, data, bytes, timestamp);
        break;

    case FLV_TYPE_AUDIO:
        rc = rtmp_client_push_audio(pPusher->_pClient, data, bytes, timestamp);
        break;

    default:
        rc = -EINVAL;
        break;
    }

    return rc;
}

const int RTMPMuxer::getCodecID(Core::CodecID id)
{
    for (const auto& m : codecMap)
    {
        if (id == m.id)
            return m.type;
    }

    throw Poco::NotFoundException(
        Poco::format("Video::MediaServer: Can't find %s in codec-map", Core::codecIDToString(id)));
}
}
}