#include "pch.h"
#include "mp4enc.h"

#include <codec_api.h>
#include <libyuv.h>

#define MP4V2_NO_STDINT_DEFS
#include <mp4v2/mp4v2.h>

static void encoder_trace_callback(void* context, int level, const char* message)
{
    LOG("ENCODER: %s", message);
}

MP4Encoder::~MP4Encoder()
{
    destroy();
}

bool MP4Encoder::init(int width, int height, int fps, int bitrate) noexcept
{
    if (WelsCreateSVCEncoder(&m_encoder) != 0)
        return false;

    m_width = width;
    m_height = height;
    m_framerate = fps;
    m_bitrate = bitrate;

    //Encoder params
    SEncParamExt param;
    m_encoder->GetDefaultParams(&param);
    param.iUsageType = CAMERA_VIDEO_REAL_TIME;
    param.fMaxFrameRate = m_framerate;
    param.iLtrMarkPeriod = 75;
    param.iPicWidth = m_width;
    param.iPicHeight = m_height;
    param.iTargetBitrate = m_bitrate;
    param.bEnableDenoise = false;
    param.iSpatialLayerNum = 1;
    param.bUseLoadBalancing = false;
    param.bEnableSceneChangeDetect = false;
    param.bEnableBackgroundDetection = false;
    param.bEnableAdaptiveQuant = false;
    param.bEnableFrameSkip = false;
    param.iMultipleThreadIdc = 0;
    //param.uiIntraPeriod = 10;

    for (int i = 0; i < param.iSpatialLayerNum; i++)
    {
        param.sSpatialLayers[i].iVideoWidth = m_width >> (param.iSpatialLayerNum - 1 - i);
        param.sSpatialLayers[i].iVideoHeight = m_height >> (param.iSpatialLayerNum - 1 - i);
        param.sSpatialLayers[i].fFrameRate = m_framerate;
        param.sSpatialLayers[i].iSpatialBitrate = m_bitrate;
        param.sSpatialLayers[i].uiProfileIdc = PRO_BASELINE;
        param.sSpatialLayers[i].uiLevelIdc = LEVEL_4_2;
        param.sSpatialLayers[i].iDLayerQp = 42;

        //SSliceArgument sliceArg;
        //sliceArg.uiSliceMode = SM_FIXEDSLCNUM_SLICE;
        //sliceArg.uiSliceNum = 0; 
        //param.sSpatialLayers[i].sSliceArgument = sliceArg;
    }

    param.uiMaxNalSize = 1500;
    param.iTargetBitrate *= param.iSpatialLayerNum;
    if (m_encoder->InitializeExt(&param) != 0)
    {
        LOG("Error initializing the encoder");
        return false;
    }

    int trace_level = WELS_LOG_ERROR;
    m_encoder->SetOption(ENCODER_OPTION_TRACE_LEVEL, &trace_level);
    m_encoder->SetOption(ENCODER_OPTION_TRACE_CALLBACK, (void*)encoder_trace_callback);
    m_encoder->SetOption(ENCODER_OPTION_TRACE_CALLBACK_CONTEXT, this);
    int videoFormat = videoFormatI420;
    m_encoder->SetOption(ENCODER_OPTION_DATAFORMAT, &videoFormat);

    return true;
}

bool MP4Encoder::init() noexcept
{
    return init(m_width, m_height, m_framerate, m_bitrate);
}

bool MP4Encoder::encode(const Image& rgba) noexcept
{
    // lazy allocation of YUV buffer
    if (m_yuv_buffer.empty())
        m_yuv_buffer.resize(m_width * m_height * 3 / 2);

    SFrameBSInfo info = { 0 };
    SSourcePicture pic = { 0 };
    pic.iPicWidth = m_width;
    pic.iPicHeight = m_height;
    pic.iColorFormat = videoFormatI420;
    pic.iStride[0] = pic.iPicWidth;
    pic.iStride[1] = pic.iStride[2] = pic.iPicWidth >> 1;
    pic.pData[0] = m_yuv_buffer.data();
    pic.pData[1] = pic.pData[0] + ((size_t)m_width * m_height);
    pic.pData[2] = pic.pData[1] + ((size_t)m_width * m_height / 4);

    if (rgba.width != m_width || rgba.height != m_height)
    {
        Image resized = rgba.resize(m_width, m_height);
        libyuv::ABGRToI420(resized.data(), resized.width * 4, pic.pData[0], m_width,
            pic.pData[1], m_width / 2, pic.pData[2], m_width / 2, m_width, m_height);
    }
    else
    {
        libyuv::ABGRToI420(rgba.data(), rgba.width * 4, pic.pData[0], m_width,
            pic.pData[1], m_width / 2, pic.pData[2], m_width / 2, m_width, m_height);
    }
    
    if (m_encoder->EncodeFrame(&pic, &info))
    {
        LOG("EcncodeFrame failed");
        return false;
    }

    // No frames generated at this time
    if (info.eFrameType == videoFrameTypeSkip)
        return false;

    for (int layer = 0; layer < info.iLayerNum; layer++)
    {
        size_t bs_size = 0;
        for (int nal = 0; nal < info.sLayerInfo[layer].iNalCount; nal++)
        {
            std::array<uint8_t, 5> nalu_bytes;
            for (int i = 0; i < nalu_bytes.size(); i++)
                nalu_bytes[i] = info.sLayerInfo[layer].pBsBuf[bs_size + i];
            if (nalu_bytes[4] == 0x67) // SPS
            {
                m_header.avc_profile = info.sLayerInfo[layer].pBsBuf[bs_size + 5];
                m_header.avc_compat = info.sLayerInfo[layer].pBsBuf[bs_size + 6];
                m_header.avc_level = info.sLayerInfo[layer].pBsBuf[bs_size + 7];
                uint8_t* ptr = info.sLayerInfo[layer].pBsBuf + bs_size + 4;
                size_t sz = info.sLayerInfo[layer].pNalLengthInByte[nal] - 4ull;
                m_header.SPS_data = std::vector<uint8_t>(ptr, ptr + sz);
            }
            else if (nalu_bytes[4] == 0x68) // PPS
            {
                uint8_t* ptr = info.sLayerInfo[layer].pBsBuf + bs_size + 4;
                size_t sz = info.sLayerInfo[layer].pNalLengthInByte[nal] - 4ull;
                m_header.PPS_data = std::vector<uint8_t>(ptr, ptr + sz);
            }
            else
            {
                uint32_t nalu_sz = info.sLayerInfo[layer].pNalLengthInByte[nal];
                uint8_t* data = info.sLayerInfo[layer].pBsBuf + bs_size;
                *(uint32_t*)data = BinaryStream::htonx(nalu_sz - 4);
                Frame::kType frame_type;
                if (nalu_bytes[4] == 0x65)      // 0x65 I-frame
                    frame_type = Frame::kType::IFrame;
                else if (nalu_bytes[4] == 0x61) // 0x61 P-frame
                    frame_type = Frame::kType::PFrame;
                else                            // something else, don't care
                    frame_type = Frame::kType::Unknown;
                m_frames.emplace_back();
                m_frames.back().type = frame_type;
                m_frames.back().data = std::vector<uint8_t>(data, data + nalu_sz);
            }
            //printf("nalu %x\n", nalu_bytes[4]);
            bs_size += info.sLayerInfo[layer].pNalLengthInByte[nal];
        }
    }

    return true;
}

bool MP4Encoder::write_mp4(const std::string& filename) const noexcept
{
    MP4FileHandle mp4 = MP4Create(filename.c_str());
    MP4SetTimeScale(mp4, 90000);
    MP4SetVideoProfileLevel(mp4, 1);
    MP4TrackId mp4_track = MP4AddH264VideoTrack(mp4, 90000, 90000 / m_framerate, m_width, m_height, 
        m_header.avc_profile, m_header.avc_compat, m_header.avc_level, 3);
    MP4AddH264SequenceParameterSet(mp4, mp4_track, m_header.SPS_data.data(), m_header.SPS_data.size());
    MP4AddH264PictureParameterSet(mp4, mp4_track, m_header.PPS_data.data(), m_header.PPS_data.size());
    for (const auto& f : m_frames)
    {
        if (!MP4WriteSample(mp4, mp4_track, f.data.data(), f.data.size(), 
            MP4_INVALID_DURATION, 0, f.type == Frame::kType::IFrame))
        {
            LOG("failed to write mp4 sample");
            MP4Close(mp4);
            return false;
        }
    }
    MP4Close(mp4);
    return true;
}

void MP4Encoder::destroy() noexcept
{
    if (m_encoder)
    {
        m_encoder->Uninitialize();
        WelsDestroySVCEncoder(m_encoder);
    }
    m_encoder = nullptr;
    m_frames.clear();
    m_header.PPS_data.clear();
    m_header.SPS_data.clear();
}

bool MP4Encoder::read(BinaryStreamReader& r)
{
    Serializer::Descriptor d;
    r >> d;
    d.value<Serializer::Integer>("width", m_width);
    d.value<Serializer::Integer>("height", m_height);
    d.value<Serializer::Integer>("bitrate", m_bitrate);
    d.value<Serializer::Float>("framerate", m_framerate);
    d.value<Serializer::RawData>("SPS", m_header.SPS_data);
    d.value<Serializer::RawData>("PPS", m_header.PPS_data);
    d.value<Serializer::Integer>("avc_profile", m_header.avc_profile);
    d.value<Serializer::Integer>("avc_compat", m_header.avc_compat);
    d.value<Serializer::Integer>("avc_level", m_header.avc_level);
    auto frames = d.get<Serializer::List>("frames");
    for (const auto& f : frames->items)
    {
        if (auto fd = std::dynamic_pointer_cast<Serializer::Descriptor>(f))
        {
            m_frames.emplace_back();
            m_frames.back().type = (Frame::kType)fd->value<Serializer::Integer>("type");
            m_frames.back().data = fd->value<Serializer::RawData>("data");
        }
    }
    return true;
}

void MP4Encoder::write(BinaryStreamWriter& w) const
{
    Serializer::Descriptor d;
    d.class_id = "mp4enc";
    d.name = L"MP4 Encoder class";
    d.props["width"] = std::make_shared<Serializer::Integer>(m_width);
    d.props["height"] = std::make_shared<Serializer::Integer>(m_height);
    d.props["bitrate"] = std::make_shared<Serializer::Integer>(m_bitrate);
    d.props["framerate"] = std::make_shared<Serializer::Float>(m_framerate);
    d.props["SPS"] = std::make_shared<Serializer::RawData>(m_header.SPS_data);
    d.props["PPS"] = std::make_shared<Serializer::RawData>(m_header.PPS_data);
    d.props["avc_profile"] = std::make_shared<Serializer::Integer>(m_header.avc_profile);
    d.props["avc_compat"] = std::make_shared<Serializer::Integer>(m_header.avc_compat);
    d.props["avc_level"] = std::make_shared<Serializer::Integer>(m_header.avc_level);
    auto frames = std::make_shared<Serializer::List>();
    for (const auto& f : m_frames)
    {
        auto fd = frames->add<Serializer::Descriptor>();
        fd->props["type"] = std::make_shared<Serializer::Integer>((int)f.type);
        fd->props["data"] = std::make_shared<Serializer::RawData>(f.data);
    }
    d.props["frames"] = frames;
    w << d;
}