/*
 * Copyright (c) 2018-2020, NVIDIA CORPORATION.  All rights reserved.
 *
 * NVIDIA CORPORATION and its licensors retain all intellectual property
 * and proprietary rights in and to this software, related documentation
 * and any modifications thereto.  Any use, reproduction, disclosure or
 * distribution of this software and related documentation without an express
 * license agreement from NVIDIA CORPORATION is strictly prohibited.
 */

#include <iostream>
#include <chrono>
#include <pthread.h>

#include "CComposite.hpp"
#include "CUtils.hpp"

using namespace std;

SIPLStatus CComposite::InitDisplay(uint32_t uDispID,
                                   uint32_t uWinID,
                                   uint32_t uWinDepth) {
    // Create IDP
    m_pNvmIDP.reset(NvMediaIDPCreate(uDispID, uWinID, NULL, NVMEDIA_FALSE));
    if (m_pNvmIDP == nullptr) {
        LOG_ERR("Compositor: NvMediaIDPCreate failed!.\n");
        return NVSIPL_STATUS_ERROR;
    }

    // Check properties of the display
    NvMediaIDPDeviceParams* dispOutputs = NULL;
    auto numDisplays = 0;
    auto nvmStatus = NvMediaIDPQuery(&numDisplays, NULL);
    if (nvmStatus != NVMEDIA_STATUS_OK) {
        LOG_ERR("Compositor: NvMediaIDPQuery failed! nvmStatus:%u\n", nvmStatus);
        return NVSIPL_STATUS_ERROR;
    }
    // Get available displays
    dispOutputs = new (std::nothrow) NvMediaIDPDeviceParams[numDisplays];
    if (dispOutputs == NULL) {
        LOG_ERR("Compositor: Allocating memory for NvMediaIDPDeviceParams failed!\n");
        return NVSIPL_STATUS_OUT_OF_MEMORY;
    }
    nvmStatus = NvMediaIDPQuery(&numDisplays, dispOutputs);
    if (nvmStatus != NVMEDIA_STATUS_OK) {
        LOG_ERR("Compositor: NvMediaIDPQuery failed! nvmStatus:%u\n", nvmStatus);
        delete [] dispOutputs;
        return NVSIPL_STATUS_ERROR;
    }
    for (auto i = 0; i < numDisplays; i++) {
        if (dispOutputs[i].displayId == uDispID) {
            m_oDispInfo = dispOutputs[i];
            break;
        }
    }
    delete [] dispOutputs;

    if (m_oDispInfo.refreshRate < 0.1) {
        m_oDispInfo.refreshRate = 60.0f; // Assume 60 fps display if the refresh rate is not available
    }

    // Set Depth
    NvMediaDispAttributes attrib { };
    attrib.depth = uWinDepth;
    nvmStatus = NvMediaIDPSetAttributes(m_pNvmIDP.get(),
                                        NVMEDIA_DISP_ATTR_SET_DEPTH,
                                        &attrib);
    if (nvmStatus != NVMEDIA_STATUS_OK) {
        LOG_ERR("Compositor: NvMediaIDPSetAttributes(Depth:%u) failed!\n", attrib.depth);
        return NVSIPL_STATUS_ERROR;
    }
    return NVSIPL_STATUS_OK;
}

SIPLStatus CComposite::Init(uint32_t uDispID,
                            uint32_t uWinID,
                            uint32_t uWinDepth,
                            NvMediaRect *pRect,
                            NvSciBufModule bufModule,
                            NvSciSyncModule syncModule)
{
    // Initialize state
    m_bRunning = false;

    // Create an NvMedia Device to be used to allocated images.
    m_pDevice.reset(NvMediaDeviceCreate());
    if (!m_pDevice) {
        LOG_ERR("Compositor: NvMediaDeviceCreate failed\n");
        return NVSIPL_STATUS_ERROR;
    }

    // Create 2D for each input.
    for (auto row = 0U; row < NUM_OF_ROWS; row++) {
        m_p2DDevices[row].reset(NvMedia2DCreate(m_pDevice.get()));
        if (!m_p2DDevices[row]) {
            LOG_ERR("Compositor: NvMedia2DCreate failed for source:%u\n", row);
            return NVSIPL_STATUS_ERROR;
        }
    }

    // Initialize display
    auto status = InitDisplay(uDispID, uWinID, uWinDepth);
    if (status != NVSIPL_STATUS_OK) {
        LOG_ERR("Compositor: InitDisplay failed\n");
        return status;
    }
    m_outputWidth = m_oDispInfo.width;
    m_outputHeight = m_oDispInfo.height;
    if (pRect != nullptr) {
        m_outputRect = *pRect;
        m_outputWidth = (pRect->x1 - pRect->x0);
        m_outputHeight = (pRect->y1 - pRect->y0);
        if (m_outputWidth == 0 || m_outputHeight == 0) {
            LOG_ERR("Compositor: Invalid output resolution specified\n");
            return NVSIPL_STATUS_BAD_ARGUMENT;
        }
        m_bRectSet = true;
    }

    // Create output buffer pool
    vector<NvMediaSurfAllocAttr> vAllocAttr;
    vAllocAttr.push_back( { NVM_SURF_ATTR_WIDTH, m_outputWidth });
    vAllocAttr.push_back( { NVM_SURF_ATTR_HEIGHT, m_outputHeight });
    vAllocAttr.push_back( { NVM_SURF_ATTR_CPU_ACCESS, NVM_SURF_ATTR_CPU_ACCESS_UNCACHED });
    vAllocAttr.push_back( { NVM_SURF_ATTR_ALLOC_TYPE, NVM_SURF_ATTR_ALLOC_ISOCHRONOUS });

    NVM_SURF_FMT_DEFINE_ATTR(compositeFormatAttrs);
    NVM_SURF_FMT_SET_ATTR_RGBA(compositeFormatAttrs, RGBA, UINT, 8, PL);
    auto surfType = NvMediaSurfaceFormatGetType(compositeFormatAttrs, NVM_SURF_FMT_ATTR_MAX);

    // Create output images and add them to the queue.
    for (auto i = 0u; i < NUM_OUT_BUFFERS; i++) {
        m_outputImages[i].reset(NvMediaImageCreateNew(m_pDevice.get(),
                                surfType,
                                vAllocAttr.data(),
                                vAllocAttr.size(),
                                0));
        if (m_outputImages[i] == nullptr) {
            LOG_ERR("Compositor: Failed to create output image\n");
            return NVSIPL_STATUS_OUT_OF_MEMORY;
        }
        m_outputQueue.push(m_outputImages[i].get());
    }

    // Create a blank NvMediaImage to clear the frames during transitions
    m_pBlankFrame.reset(NvMediaImageCreateNew(m_pDevice.get(),
                                             surfType,
                                             vAllocAttr.data(),
                                             vAllocAttr.size(),
                                             0));
    if (m_pBlankFrame == nullptr) {
        LOG_ERR("Compositor: Failed to create clear frame\n");
        return NVSIPL_STATUS_OUT_OF_MEMORY;
    }

    // Memset the NvMediaImage (assumes the blank frame to be RGBA)
    uint8_t* buff = new (std::nothrow) uint8_t[m_outputWidth * m_outputHeight * 4];
    if (buff == nullptr) {
        LOG_ERR("Compositor: Failed to allocate buffer\n");
        return NVSIPL_STATUS_OUT_OF_MEMORY;
    }
    std::fill(buff, buff + m_outputWidth * m_outputHeight * 4, 0);
    uint8_t** buffers = &buff;
    auto buffPitches = static_cast<uint32_t>(m_outputWidth * 4);
    NvMediaImageSurfaceMap surfaceMap;
    NvMediaImageLock(m_pBlankFrame.get(), NVMEDIA_IMAGE_ACCESS_WRITE, &surfaceMap);
    auto nvmStatus = NvMediaImagePutBits(m_pBlankFrame.get(),
                                         nullptr,
                                         (void **)buffers,
                                         &buffPitches);
    if (nvmStatus != NVMEDIA_STATUS_OK) {
        LOG_ERR("Compositor: NvMediaImagePutBits failed\n");
    }
    NvMediaImageUnlock(m_pBlankFrame.get());
    delete [] buff;

    // Clear all output buffers
    ClearOutputBuffers();

    return NVSIPL_STATUS_OK;
}

SIPLStatus CComposite::RegisterSource(uint32_t groupIndex,
                                      uint32_t modIndex,
                                      uint32_t outIndex,
                                      bool isRaw,
                                      uint32_t &id,
                                      NvSciStreamBlock *consumer,
                                      NvSciStreamBlock **upstream,
                                      NvSciStreamBlock *queue,
                                      QueueType type)
{
    if ((groupIndex >= NUM_OF_GROUPS) or (modIndex >= NUM_OF_ROWS) or (outIndex >= NUM_OF_COLS)) {
        LOG_ERR("Compositor: RegisterSource: Invalid argument\n");
        return NVSIPL_STATUS_BAD_ARGUMENT;
    }

    // Compute global ID for the source
    id = groupIndex * (NUM_OF_COLS * NUM_OF_ROWS) + modIndex * NUM_OF_COLS + outIndex;

    // Add source to the group
    m_GroupInfos[groupIndex].push_back(id);

    LOG_INFO("Compositor: Registered output:%u from link:%u of quad: %u as id:%u\n", outIndex, modIndex, groupIndex, id);
    return NVSIPL_STATUS_OK;
}

SIPLStatus CComposite::Post(uint32_t id, INvSIPLClient::INvSIPLNvMBuffer *pBuffer)
{
    if (!m_bRunning) {
        // Composite is not ready to accept buffers
        LOG_WARN("Compositor is not ready to accept buffers\n");
        return NVSIPL_STATUS_OK;
    }

    // Check if ID belongs to currently active group
    bool found = false;
    for (const auto& i : m_GroupInfos[m_uGroupIndex]) {
        if (i == id) {
            found = true;
            break;
        }
    }

    if (found == false) {
        // Source does not belong to the active group
        return NVSIPL_STATUS_OK;
    }

    // Check if input image is RAW and allocate scratch RGB surface if necessary
    auto nvmImg = pBuffer->GetImage();
    bool isRaw = false;
    auto nvmStatus = CUtils::IsRawImage(nvmImg, isRaw);
    if (nvmStatus != NVSIPL_STATUS_OK) {
        LOG_ERR("Compositor: CUtils::IsRawImage failed for source:%u\n", id);
        return NVSIPL_STATUS_ERROR;
    }
    if (isRaw) {
        if (m_scratchRGBAImages[id] == nullptr) {
            // Allocate scratch RGBA image for raw to RGBA conversion
            m_scratchRGBAImages[id].reset(CUtils::CreateRgbaImage(m_pDevice.get(), nvmImg));
            if (m_scratchRGBAImages[id] == nullptr) {
                LOG_ERR("Compositor: Failed to allocate scratch RGBA image for source:%u\n", id);
                return NVSIPL_STATUS_ERROR;
            }
        }
        if (m_rawImageBufs[id] == nullptr) {
            // Allocate CPU buffer for raw image
            m_rawImageBufs[id].reset(CUtils::CreateRawBuffer(nvmImg));
            if (m_rawImageBufs[id] == nullptr) {
                LOG_ERR("Compositor: Failed to allocate raw image buffer for source:%u\n", id);
                return NVSIPL_STATUS_ERROR;
            }
        }
        if (m_rgbaImageBufs[id] == nullptr) {
            // Allocate CPU buffer for RGBA image
            m_rgbaImageBufs[id].reset(CUtils::CreateRgbaBuffer(m_scratchRGBAImages[id].get()));
            if (m_rgbaImageBufs[id] == nullptr) {
                LOG_ERR("Compositor: Failed to allocate RGBA image buffer for source:%u\n", id);
                return NVSIPL_STATUS_ERROR;
            }
        }
    }

    // Add buffer to queue
    pBuffer->AddRef();
    std::unique_lock<std::mutex> lck(m_vInputQueueMutex[id]);
    m_vInputQueue[id].push(pBuffer);

    return NVSIPL_STATUS_OK;
}

SIPLStatus CComposite::SetActiveGroup(uint32_t groupIndex)
{
    if (groupIndex >= NUM_OF_GROUPS) {
        LOG_ERR("Compositor: Invalid quad number: %u\n", groupIndex);
        return NVSIPL_STATUS_BAD_ARGUMENT;
    }

    if (m_GroupInfos[groupIndex].size() == 0) {
        LOG_ERR("Compositor: No display-able outputs in group: %u\n", groupIndex);
        return NVSIPL_STATUS_BAD_ARGUMENT;
    }

    m_uNewGroupIndex = groupIndex;
    return NVSIPL_STATUS_OK;
}

void CComposite::PrintDisplayableGroups() const
{
    string s[NUM_OF_GROUPS] = { "PORT-A", "PORT-B", "PORT-C", "PORT-D" };
    for (auto g = 0u; g < NUM_OF_GROUPS; g++) {
        if (m_GroupInfos[g].size() != 0) {
            cout << g << ":\t" << s[g] << endl;
        }
    }
}

SIPLStatus CComposite::Start()
{
    if (m_uGroupIndex == -1U) {
        // No valid group specified. Use the first one with non-zero sources
        for (auto g=0u; g < NUM_OF_GROUPS; g++) {
            if (m_GroupInfos[g].size() != 0) {
                m_uNewGroupIndex = g;
                m_uGroupIndex = g;
                break;
            }
        }
    }

    if (m_uGroupIndex == -1U) {
        LOG_ERR("CComposite: No input registered.\n");
        return NVSIPL_STATUS_INVALID_STATE;
    }

    // Set up the source rectangles
    uint16_t xStep = m_outputWidth / NUM_OF_COLS;
    uint16_t yStep = m_outputHeight / NUM_OF_ROWS;

    for (auto g = 0u; g < NUM_OF_GROUPS; g++) {
        auto& vOutIDs = m_GroupInfos[g];
        auto numInputs = vOutIDs.size();
        for (auto i = 0u; i < numInputs; i++) {
            auto id = vOutIDs[i] - g * (NUM_OF_COLS * NUM_OF_ROWS);
            auto modIndex = id / NUM_OF_COLS;
            auto outIndex = id % NUM_OF_COLS;
            uint16_t startx = outIndex * xStep;
            uint16_t starty = modIndex * yStep;
            uint16_t endx = startx + xStep;
            uint16_t endy = starty + yStep;

            if (numInputs == 1) {
                // If single input. use the full screen
                startx = 0;
                starty = 0;
                endx = m_outputWidth;
                endy = m_outputHeight;
            }

            NvMediaRect tmp = { startx, starty, endx, endy };
            m_oInputRects[vOutIDs[i]] = tmp;

            LOG_INFO("Rectangle for Group:%u Link:%u Out:%u is ", g, modIndex, outIndex);
            LOG_INFO("[(%u, %u) : (%u, %u)]\n", tmp.x0, tmp.y0, tmp.x1, tmp.y1);
        }
    }

    // Start the composite thread
    LOG_INFO("Starting compositor thread\n");
    m_pthread.reset(new std::thread(&CComposite::ThreadFunc, this));
    if (m_pthread == nullptr) {
        LOG_ERR("Failed to create compositor thread\n");
        return NVSIPL_STATUS_OUT_OF_MEMORY;
    }
    LOG_INFO("Created compositor thread: ID:%u\n", m_pthread->get_id());

    return NVSIPL_STATUS_OK;
}

SIPLStatus CComposite::Stop()
{
    // Signal thread to stop
    m_bRunning = false;

    // Wait for the thread
    if (m_pthread != nullptr) {
        LOG_INFO("Waiting to join compositor thread: ID:%u\n", m_pthread->get_id());
        m_pthread->join();
    }

    return NVSIPL_STATUS_OK;
}

SIPLStatus CComposite::Deinit()
{
    // Clear all output buffers
    ClearOutputBuffers();

    // Dequeue and release all input buffers
    for (auto id = 0u; id < NUM_OF_INPUTS; id++) {
        std::unique_lock<std::mutex> lck(m_vInputQueueMutex[id]);
        while(!m_vInputQueue[id].empty()) {
            auto pBuffer = m_vInputQueue[id].front();
            pBuffer->Release();
            m_vInputQueue[id].pop();
        }
    }

    return NVSIPL_STATUS_OK;
}

bool CComposite::CheckInputQueues(void)
{
    bool anyReady = false;
    const auto& vInput = m_GroupInfos[m_uGroupIndex];
    for (const auto& id : vInput) {
        std::unique_lock<std::mutex> lck(m_vInputQueueMutex[id]);
        if (!m_vInputQueue[id].empty()) {
            anyReady = true;
            break;
        }
    }
    return anyReady;
}

void CComposite::ThreadFunc(void)
{
    pthread_setname_np(pthread_self(), "CComposite");

    m_bRunning = true;

    uint32_t sleepTimeMs = uint32_t(1000 / (m_oDispInfo.refreshRate));

    while (m_bRunning) {
        // Check if there is a pending group change.
        if (m_uNewGroupIndex != m_uGroupIndex) {
            // Clear all output buffers
            ClearOutputBuffers();

            // Release buffers from input queues for current group
            const auto &vInput = m_GroupInfos[m_uGroupIndex];
            for (const auto& id : vInput) {
                std::unique_lock<std::mutex> lck(m_vInputQueueMutex[id]);
                if (!m_vInputQueue[id].empty()) {
                    INvSIPLClient::INvSIPLNvMBuffer *pBuffer = m_vInputQueue[id].front();
                    if (pBuffer != nullptr) {
                        pBuffer->Release();
                    }
                    m_vInputQueue[id].pop();
                }
            }
            // Update current group index
            m_uGroupIndex = m_uNewGroupIndex;
        }

        // Check for input readiness
        bool anyReady = CheckInputQueues();
        if (!anyReady) {
            LOG_INFO("Compositor does not have any inputs available yet\n");

            // Sleep for refresh rate
            std::this_thread::sleep_for(std::chrono::milliseconds(sleepTimeMs));
            continue;
        }

        // Get output image
        if (m_outputQueue.empty()) {
            LOG_ERR("Compositor output queue empty\n");
            return;
        }

        NvMediaImage* pDstImg = m_outputQueue.front();
        m_outputQueue.pop();

        // Get full buffer from input queue and composite on to output image
        const auto& vInput = m_GroupInfos[m_uGroupIndex];
        for (const auto&id : vInput) {
            INvSIPLClient::INvSIPLNvMBuffer *pBuffer = nullptr;
            {
                std::unique_lock<std::mutex> lck(m_vInputQueueMutex[id]);
                if (m_vInputQueue[id].empty()) {
                    LOG_INFO("Compositor: Input queue for source %u is empty\n", id);
                    continue;
                }
                pBuffer = m_vInputQueue[id].front();
                m_vInputQueue[id].pop();
            }

            auto pSrcImg = pBuffer->GetImage();
            NvMediaRect& dstRect = m_oInputRects[id];

            // Convert RAW to RGBA if necessary
            bool isRaw = false;
            auto utilStatus = CUtils::IsRawImage(pSrcImg, isRaw);
            if (utilStatus != NVSIPL_STATUS_OK) {
                LOG_ERR("Compositor: CUtils::IsRawImage failed for source %u\n", id);
                pBuffer->Release();
                return;
            }
            if (isRaw) {
                auto status = CUtils::ConvertRawToRgba(pSrcImg,
                                                       m_rawImageBufs[id].get(),
                                                       m_scratchRGBAImages[id].get(),
                                                       m_rgbaImageBufs[id].get());
                if (status != NVSIPL_STATUS_OK) {
                    LOG_ERR("Compositor: CUtils::IsRawConvertRawToRgbaImage failed for source %u\n", id);
                    pBuffer->Release();
                    return;
                }
                pSrcImg = m_scratchRGBAImages[id].get();
            }

            // Blit
            auto row = (id - m_uGroupIndex * (NUM_OF_COLS * NUM_OF_ROWS))/NUM_OF_COLS;
            auto nvmStatus = NvMedia2DBlitEx(m_p2DDevices[row].get(),
                                             pDstImg,
                                             &dstRect,
                                             pSrcImg,
                                             NULL,
                                             NULL,
                                             NULL);
            if (nvmStatus != NVMEDIA_STATUS_OK) {
                LOG_ERR("Compositor: NvMedia2DBlitEx failed for source %u\n", id);
                pBuffer->Release();
                return;
            }

            // Release the input buffer
            pBuffer->Release();
        }

        // Send composited output image to display
        NvMediaImage *releaseList[2] = { NULL }; // IDPFlip needs at least two elements
        NvMediaRect* pRect = nullptr;
        if (m_bRectSet) {
            pRect = &m_outputRect;
        }
        auto nvmStatus = NvMediaIDPFlip(m_pNvmIDP.get(), pDstImg, NULL, pRect, releaseList, NULL);
        if (nvmStatus != NVMEDIA_STATUS_OK) {
            LOG_ERR("Compositor: NvMediaIDPFlip failed!. nvmStatus: %u\n", nvmStatus);
            return;
        }
        auto pPrevImage = releaseList[0];
        if (pPrevImage != NULL) {
            // Put the previous image back on the queue.
            m_outputQueue.push(pPrevImage);
        }
    } // while ()

    return;
}

SIPLStatus CComposite::ClearOutputBuffers()
{
    // Get back any pending image from the display
    NvMediaImage *releaseList[2] = { NULL }; // IDPFlip needs at least two elements
    auto nvmStatus = NvMediaIDPFlip(m_pNvmIDP.get(), NULL, NULL, NULL, releaseList, NULL);
    if (nvmStatus != NVMEDIA_STATUS_OK) {
        LOG_ERR("Compositor: NvMediaIDPFlip(NULL) failed!. nvmStatus: %u\n", nvmStatus);
        return NVSIPL_STATUS_ERROR;
    }
    auto pPrevImage = releaseList[0];
    if (pPrevImage != NULL) {
        // Put the previous image back on the queue.
        m_outputQueue.push(pPrevImage);
    }

    // Check output queue
    for (auto i = 0u; i < NUM_OUT_BUFFERS; i++) {
        if (m_outputQueue.empty()) {
            return NVSIPL_STATUS_OK;
        }
        auto pImage = m_outputQueue.front();
        m_outputQueue.pop();

        // Blit clear frame
        nvmStatus = NvMedia2DBlitEx(m_p2DDevices[0].get(),
                                    pImage,
                                    NULL,
                                    m_pBlankFrame.get(),
                                    NULL,
                                    NULL,
                                    NULL);
        if (nvmStatus != NVMEDIA_STATUS_OK) {
            LOG_ERR("Compositor: NvMedia2DBlitEx(ClearFrame) failed\n");
            return NVSIPL_STATUS_ERROR;
        }

        // Put back the cleared image
        m_outputQueue.push(pImage);
    };

    return NVSIPL_STATUS_OK;
}
