/*---------------------------------------------------------------------------*/
/* video_decoder.cpp                                                         */
/* copyright (c) innovaphone 2018                                            */
/*                                                                           */
/*---------------------------------------------------------------------------*/

#include "platform/platform.h"
#include "common/os/iomux.h"
#include "common/interface/socket.h"
#include "common/ilib/str.h"
#include "common/interface/media.h"
#include "video_transcoder.h"

#include "video_decoder.h"

class IMediaDecoder * IMediaDecoder::Create(class IIoMux * const iomux, class UMediaDecoder * const user, class IInstanceLog * log, enum VideoCoder coder)
{
    switch (coder) {
    case VideoCoder::VP8:
        return new VideoDecoderVPX(iomux, user, log, false);
    case VideoCoder::VP9:
        return new VideoDecoderVPX(iomux, user, log, true);
    default:
        return new VideoDecoderH264(iomux, user, log);
    }
}

VideoDecoderH264::VideoDecoderH264(class IIoMux * const iomux, class UMediaDecoder * const user, class IInstanceLog * log) :
    iomux(iomux),
    user(user),
    log(log)
{
    log->Log(LOG_VIDEO, "VideoDecoderH264(%p)::VideoDecoderH264", this);
    execCnt = 0;
    shutDown = false;
    videoFrame = (char *)malloc(640*480*2);
    dropFrame = false;
    sequenceGap = false;
    rtpMarker = false;
    frameLength = 0;
    frameLengthFua = 0;
    fragSb = 0;
    fuaTimestamp = 0;
    transcoder = IMediaTranscoder::Create(iomux, this, log, VideoCoder::H264, (enum VideoCoder) 0);
}

VideoDecoderH264::~VideoDecoderH264() 
{
    log->Log(LOG_VIDEO, "VideoDecoderH264(%p)::~VideoDecoderH264", this);
    if(videoFrame) free(videoFrame);
}

void VideoDecoderH264::Close()
{
    shutDown = true;
    transcoder->Close();
}

void VideoDecoderH264::MediaTranscoderCloseResult(class IMediaTranscoder * transcoder)
{
    delete transcoder;
    transcoder = NULL;
    user->MediaDecoderCloseResult(this);
}

void VideoDecoderH264::DecodeStream(const void *stream, int len, dword timestamp, short sequenceNumberDiff, bool marker)
{
    const char * buf = (const char *)stream;
    int nalUnitType = (buf[0] & 0x1F);
    static const byte startCode[] ={0, 0, 0, 1};
    static const unsigned int startCodeSize = sizeof(startCode);

    log->Log(LOG_VIDEO, "VideoDecoderH264(%p)::Recv h264 len=%d timestamp=%x sequenceNumberDiff=%d marker=%d,%d sequenceGap=%d", this, len, timestamp, sequenceNumberDiff, rtpMarker, marker, sequenceGap);

    if(lTimestamp != timestamp) {
        if(frameLength && rtpMarker && !dropFrame) {
            DecodeFrame((const byte *)videoFrame, frameLength, lTimestamp/*, sequenceGap ? 1 : 0*/);
            sequenceGap = false;
        }
        if(sequenceNumberDiff) {
            dropFrame = true;
            sequenceGap = true;
        }
        else dropFrame = false;
        frameLength = 0;
        frameLengthFua = 0;
        fragSb = 0;
        fuaTimestamp = 0;
    }
    else if(sequenceNumberDiff) {
        dropFrame = true;
        sequenceGap = true;
    }
    rtpMarker = marker;

    if(nalUnitType < RTP_STAP_A) {
        memcpy(&videoFrame[frameLength], startCode, startCodeSize);
        memcpy(&videoFrame[frameLength + startCodeSize], buf, len);
        frameLength += startCodeSize + len;
    }
    else if(nalUnitType == RTP_STAP_A) {
        int indPkt = 1;

        while(indPkt < len) {
            int nalLength = READ_WORD_NET(&buf[indPkt]);

            if((2+nalLength) > len) {
                break;
            }
            indPkt += 2;
            nalUnitType = (buf[indPkt] & 0x1F);
            if ((frameLength == 0) && ((nalUnitType == NAL_UNIT_IDR) || (nalUnitType == NAL_UNIT_SPS) || (nalUnitType == NAL_UNIT_PPS))) {
                log->Log(LOG_VIDEO, "VideoDecoderH264(%p)::recv h264 IDR", this);
            }
            memcpy(&videoFrame[frameLength], startCode, startCodeSize);
            memcpy(&videoFrame[frameLength + startCodeSize], &buf[indPkt], nalLength);
            frameLength += startCodeSize + nalLength;

            indPkt += nalLength;
        }
    }
    else if(nalUnitType == RTP_FU_A) {
        if(buf[1] & 0x80) {  // Start Bit (not allowed start and end bit inside the same packet)
            int nalRefIdc;

            if(fragSb) {
                log->Log(LOG_VIDEO, "VideoDecoderH264(%p)::recv Error: start bit already set", this);
                frameLengthFua = 0;
                fragSb = 0;
                fuaTimestamp = 0;
            }
            nalRefIdc = buf[0] & 0xE0;
            nalUnitType = buf[1] & 0x1F;

            if ((frameLength == 0) && ((nalUnitType == NAL_UNIT_IDR) || (nalUnitType == NAL_UNIT_SPS) || (nalUnitType == NAL_UNIT_PPS))) {
                log->Log(LOG_VIDEO, "VideoDecoderH264(%p)::recv h264 IDR", this);
            }
            memcpy(&videoFrame[frameLength], startCode, startCodeSize);
            videoFrame[frameLength + startCodeSize] = (nalRefIdc | nalUnitType);
            memcpy(&videoFrame[frameLength + startCodeSize + 1], &buf[2], (len-2));
            frameLengthFua = startCodeSize + (len-1);

            fragSb = 1;
            fuaTimestamp = timestamp;
        }
        else {
            if(!frameLengthFua || !fragSb || (fuaTimestamp != timestamp) || sequenceNumberDiff) {
                log->Log(LOG_VIDEO, "VideoDecoderH264(%p)::recv No start Bit available, no data, different timestamp or frame lost (%d,%d,%d,%d) ts=%d",
                    this, !frameLengthFua, !fragSb, (fuaTimestamp != timestamp), sequenceNumberDiff, lTimestamp);
                fragSb = 0;
                fuaTimestamp = 0;
                frameLengthFua = 0;
            }
            else {
                memcpy(&videoFrame[frameLength+frameLengthFua], &buf[2], len-2);
                frameLengthFua += len-2;
                if(buf[1] & 0x40) {  // End Bit
                    fragSb = 0;
                    fuaTimestamp = 0;
                    frameLength += frameLengthFua;
                }
            }
        }
    }
    else if(nalUnitType <= RTP_MTAP24) {
        log->Log(LOG_VIDEO, "VideoDecoderH264(%p)::recv aggregation type not implemented yet", this);
    }
    else if(nalUnitType == RTP_FU_B) {
        log->Log(LOG_VIDEO, "VideoDecoderH264(%p)::recv fragmentation type FU_B not implemented yet", this);
    }
    else {
        log->Log(LOG_VIDEO, "VideoDecoderH264(%p)::recv Unknown NAL type received(%d)=%d", this, nalUnitType);
    }
    lTimestamp = timestamp;
}

void VideoDecoderH264::DecodeFrame(const void *buf, int len, dword timestamp)
{
    if(shutDown) return;
    transcoder->Transcode(buf, len, timestamp, 0, 0, 0, 0, 0, false);
}

void VideoDecoderH264::TranscodeResult(class IMediaTranscoder * transcoder, const void * buf, int len, dword timestamp, dword duration, short sequenceDiff, int width, int height, int stride, dword flags)
{
    struct VideoFrameFormat format;

    format.captureWidth = width;
    format.captureHeight = height;
    format.width = width;
    format.height = height;
    format.stride = stride;
    format.sampleSize = stride * height + (stride / 2) * height;
    format.interlace = VideoInterlaceMode::VIDEO_INTERLACE_Unknown; 
    format.format = VideoFormat::VIDEO_FORMAT_IYUV;
    format.frameRate.numerator = 15;
    format.frameRate.denominator = 1;
    if (width * 9 == height * 16) {
        format.aspectRatio.numerator = 16;
        format.aspectRatio.denominator = 9;
    }
    else {
        format.aspectRatio.numerator = 4;
        format.aspectRatio.denominator = 3;
    }
    format.hdVideo = ((width >= 1080) || (height >= 960));
    format.targetBitrate = format.hdVideo ? 700 : 350;
    format.averageBitrate = format.targetBitrate * 1000;
    format.appSharing = false;

    user->DecodeResult(this, buf, len, timestamp, duration, &format);
}


VideoDecoderVPX::VideoDecoderVPX(class IIoMux * const iomux, class UMediaDecoder * const user, class IInstanceLog * log, bool vp9) :
    iomux(iomux),
    user(user),
    log(log)
{
    log->Log(LOG_VIDEO, "VideoDecoderVPX(%p)::VideoDecoderVPX", this);
    this->vp9 = vp9;
    execCnt = 0;
    shutDown = false;
    videoFrame = (byte *)(malloc(640 * 480 * 2));
    frameLength = 0;
    dropFrame = false;
    sequenceGap = false;
    rtpMarker = false;
    lTimestamp = 0;
    fragSb = 0;
    vp8PicId = 0;
    rxVp8PicId = 0;
    transcoder = IMediaTranscoder::Create(iomux, this, log, vp9 ? VideoCoder::VP9 : VideoCoder::VP8, (enum VideoCoder) 0);
}

VideoDecoderVPX::~VideoDecoderVPX() 
{
    log->Log(LOG_VIDEO, "VideoDecoderVPX(%p)::~VideoDecoderVPX", this);
    if(videoFrame) free(videoFrame);
}

void VideoDecoderVPX::Close()
{
    shutDown = true;
    transcoder->Close();
}

void VideoDecoderVPX::MediaTranscoderCloseResult(class IMediaTranscoder * transcoder)
{
    delete transcoder;
    transcoder = NULL;
    user->MediaDecoderCloseResult(this);
}

void VideoDecoderVPX::DecodeStream(const void *stream, int len, dword timestamp, short sequenceNumberDiff, bool marker)
{
    const char * buf = (const char *)stream;
    if(lTimestamp != timestamp) {
        if(frameLength && rtpMarker && !dropFrame) {
            DecodeFrame(videoFrame, frameLength, lTimestamp/*, sequenceGap ? 1 : 0*/);
            sequenceGap = false;
        }
        if(sequenceNumberDiff) {
            dropFrame = true;
            sequenceGap = true;
        }
        else dropFrame = false;
        frameLength = 0;
        fragSb = 0;
        rxVp8PicId = 0;
    }
    else if(sequenceNumberDiff) {
        dropFrame = true;
        sequenceGap = true;
    }
    rtpMarker = marker;

    int rxPid = 0;
    int offHdr = 1;                        // Payload descriptor
    if(buf[0] & 0x80) {                    // Check payload descriptor for Extension header
        offHdr++;
        if(buf[1] & 0x80) {                // Picture ID present
            if(buf[2] & 0x80) {            // Long Picture ID
                rxPid = (((word)(buf[2] & 0x7f)) << 8) + (byte)buf[3];
                offHdr += 2;
            }
            else {
                rxPid = buf[2];
                offHdr++;
            }
        }
        if(buf[1] & 0x40) offHdr++;         // TL0PICIDX byte present
        if(buf[1] & 0x20) offHdr++;         // TID byte present, TID/KEYIDX octet
        else if(buf[1] & 0x10) offHdr++;    // KEYIDX byte present, TID/KEYIDX octet
    }
    else rxPid = buf[0] & 0x7;

    if(!frameLength) {
        fragSb = (buf[0] & 0x10) ? 1 : 0;
        if(fragSb) {
            rxVp8PicId = rxPid;
            if(!(buf[0] & 0x07) && !(buf[offHdr] & 0x01)) {
                log->Log(LOG_VIDEO, "VideoDecoderVPX(%p)::recv vp8 keyframe", this);
            }
        }
    }

    if(fragSb && (rxVp8PicId == rxPid)) {
        memcpy(&videoFrame[frameLength], &buf[offHdr], len - offHdr);
        frameLength += len - offHdr;
    }
    //else debug->printf("discard frame(%d) %d %d,%d buf %x,%x len=%d,%d,%d", (buf[0] & 0x80), fragSb, rxVp8PicId, rxPid, buf[2], buf[3], frameLength, len, offHdr);

    lTimestamp = timestamp;
}

void VideoDecoderVPX::DecodeFrame(const void *buf, int len, dword timestamp)
{
    if(shutDown) return;
    transcoder->Transcode(buf, len, timestamp, 0, 0, 0, 0, 0, false);
}

void VideoDecoderVPX::TranscodeResult(class IMediaTranscoder * transcoder, const void * buf, int len, dword timestamp, dword duration, short sequenceDiff, int width, int height, int stride, dword flags)
{
    struct VideoFrameFormat format;

    format.captureWidth = width;
    format.captureHeight = height;
    format.width = width;
    format.height = height;
    format.stride = stride;
    format.sampleSize = stride * height + (stride / 2) * height;
    format.interlace = VideoInterlaceMode::VIDEO_INTERLACE_Unknown; 
    format.format = VideoFormat::VIDEO_FORMAT_IYUV;
    format.frameRate.numerator = 15;
    format.frameRate.denominator = 1;
    if (width * 9 == height * 16) {
        format.aspectRatio.numerator = 16;
        format.aspectRatio.denominator = 9;
    }
    else {
        format.aspectRatio.numerator = 4;
        format.aspectRatio.denominator = 3;
    }
    format.hdVideo = ((width >= 1080) || (height >= 960));
    format.targetBitrate = format.hdVideo ? 700 : 350;
    format.averageBitrate = format.targetBitrate * 1000;
    format.appSharing = false;

    user->DecodeResult(this, buf, len, timestamp, duration, &format);
}
