/*
 * Copyright (c) 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 "CImageManager.hpp"

// Number of images (buffers) to be allocated and registered for the capture output
static constexpr size_t CAPTURE_IMAGE_POOL_SIZE {6U};
// Number of images (buffers) to be allocated and registered for the ISP0 and ISP1 outputs
static constexpr size_t ISP_IMAGE_POOL_SIZE {4U};

struct CloseNvSciBufAttrList {
    void operator()(NvSciBufAttrList *attrList) const {
        if (attrList != nullptr) {
            if ((*attrList) != nullptr) {
                NvSciBufAttrListFree(*attrList);
            }
            delete attrList;
        }
    }
};

SIPLStatus CImageManager::AllocateImageGroups(ImagePool &imagePool)
{
    NvMediaStatus nvmStatus = NVMEDIA_STATUS_OK;
    NvSciError err = NvSciError_Success;
    std::unique_ptr<NvSciBufAttrList, CloseNvSciBufAttrList> reconciledAttrList;
    std::unique_ptr<NvSciBufAttrList, CloseNvSciBufAttrList> unreconciledAttrList;
    std::unique_ptr<NvSciBufAttrList, CloseNvSciBufAttrList> conflictAttrList;

    reconciledAttrList.reset(new NvSciBufAttrList);
    unreconciledAttrList.reset(new NvSciBufAttrList);
    conflictAttrList.reset(new NvSciBufAttrList);

    err = NvSciBufAttrListCreate(m_sciBufModule, unreconciledAttrList.get());
    CHK_NVSCISTATUS_AND_RETURN(err, "NvSciBufAttrListCreate");

    NvSciBufAttrValAccessPerm accessPerm = NvSciBufAccessPerm_Readonly;
    NvSciBufAttrKeyValuePair attrKvp = {
        NvSciBufGeneralAttrKey_RequiredPerm,
        &accessPerm,
        sizeof(accessPerm)
    };

    err = NvSciBufAttrListSetAttrs(*unreconciledAttrList, &attrKvp, 1U);
    CHK_NVSCISTATUS_AND_RETURN(err, "NvSciBufAttrListSetAttrs");

    nvmStatus = NvMediaImageFillNvSciBufAttrs(m_nvmDevice,
                                              imagePool.imageAttrs.surfaceType,
                                              imagePool.imageAttrs.surfaceAllocAttr.data(),
                                              imagePool.imageAttrs.surfaceAllocAttr.size(),
                                              0U,
                                              *unreconciledAttrList);
    CHK_NVMSTATUS_AND_RETURN(nvmStatus, "NvMediaImageFillNvSciBufAttrs");

    err = NvSciBufAttrListReconcile(unreconciledAttrList.get(),
                                    1U,
                                    reconciledAttrList.get(),
                                    conflictAttrList.get());
    CHK_NVSCISTATUS_AND_RETURN(err, "NvSciBufAttrListReconcile");

    imagePool.sciBufObjs.resize(imagePool.size);
    imagePool.imageGroups.resize(imagePool.size);
    for (size_t i = 0U; i < imagePool.size; i++) {
        // Allocate image group
        imagePool.imageGroups[i] = new (std::nothrow) NvMediaImageGroup;
        CHK_PTR_AND_RETURN(imagePool.imageGroups[i], "NvMediaImageGroup allocation");

        imagePool.imageGroups[i]->numImages = 1U;
        err = NvSciBufObjAlloc(*reconciledAttrList, &(imagePool.sciBufObjs[i]));
        CHK_NVSCISTATUS_AND_RETURN(err, "NvSciBufObjAlloc");

        nvmStatus = NvMediaImageCreateFromNvSciBuf(m_nvmDevice,
                                                   imagePool.sciBufObjs[i],
                                                   &(imagePool.imageGroups[i]->imageList[0]));
        CHK_NVMSTATUS_AND_RETURN(nvmStatus, "NvMediaImageCreateFromNvSciBuf");
        CHK_PTR_AND_RETURN(imagePool.imageGroups[i]->imageList[0],
                           "NvMediaImageCreateFromNvSciBuf");
    }

    return NVSIPL_STATUS_OK;
}

SIPLStatus CImageManager::AllocateImages(ImagePool &imagePool)
{
    NvMediaStatus nvmStatus = NVMEDIA_STATUS_OK;
    NvSciError err = NvSciError_Success;
    std::unique_ptr<NvSciBufAttrList, CloseNvSciBufAttrList> reconciledAttrList;
    std::unique_ptr<NvSciBufAttrList, CloseNvSciBufAttrList> unreconciledAttrList;
    std::unique_ptr<NvSciBufAttrList, CloseNvSciBufAttrList> conflictAttrList;

    reconciledAttrList.reset(new NvSciBufAttrList);
    unreconciledAttrList.reset(new NvSciBufAttrList);
    conflictAttrList.reset(new NvSciBufAttrList);

    err = NvSciBufAttrListCreate(m_sciBufModule, unreconciledAttrList.get());
    CHK_NVSCISTATUS_AND_RETURN(err, "NvSciBufAttrListCreate");

    NvSciBufAttrValAccessPerm accessPerm = NvSciBufAccessPerm_Readonly;
    NvSciBufAttrKeyValuePair attrKvp = {
        NvSciBufGeneralAttrKey_RequiredPerm,
        &accessPerm,
        sizeof(accessPerm)
    };

    err = NvSciBufAttrListSetAttrs(*unreconciledAttrList, &attrKvp, 1U);
    CHK_NVSCISTATUS_AND_RETURN(err, "NvSciBufAttrListSetAttrs");

    nvmStatus = NvMediaImageFillNvSciBufAttrs(m_nvmDevice,
                                              imagePool.imageAttrs.surfaceType,
                                              imagePool.imageAttrs.surfaceAllocAttr.data(),
                                              imagePool.imageAttrs.surfaceAllocAttr.size(),
                                              0U,
                                              *unreconciledAttrList);
    CHK_NVMSTATUS_AND_RETURN(nvmStatus, "NvMediaImageFillNvSciBufAttrs");

    err = NvSciBufAttrListReconcile(unreconciledAttrList.get(),
                                    1U,
                                    reconciledAttrList.get(),
                                    conflictAttrList.get());
    CHK_NVSCISTATUS_AND_RETURN(err, "NvSciBufAttrListReconcile");

    imagePool.sciBufObjs.resize(imagePool.size);
    imagePool.images.resize(imagePool.size);
    for (size_t i = 0U; i < imagePool.size; i++) {
        err = NvSciBufObjAlloc(*reconciledAttrList, &(imagePool.sciBufObjs[i]));
        CHK_NVSCISTATUS_AND_RETURN(err, "NvSciBufObjAlloc");

        nvmStatus = NvMediaImageCreateFromNvSciBuf(m_nvmDevice,
                                                   imagePool.sciBufObjs[i],
                                                   &(imagePool.images[i]));
        CHK_NVMSTATUS_AND_RETURN(nvmStatus, "NvMediaImageCreateFromNvSciBuf");
        CHK_PTR_AND_RETURN(imagePool.images[i], "NvMediaImageCreateFromNvSciBuf");
    }

    return NVSIPL_STATUS_OK;
}

void CImageManager::DeleteImageGroups(ImagePool &imagePool)
{
    for (size_t i = 0U; i < imagePool.imageGroups.size(); i++) {
        if (imagePool.imageGroups[i] == nullptr) {
            LOG_WARN("Attempt to destroy null image group\n");
            continue;
        }
        for (uint32_t p = 0U; p < imagePool.imageGroups[i]->numImages; p++) {
            if (imagePool.imageGroups[i]->imageList[p] == nullptr) {
                LOG_WARN("Attempt to destroy null image\n");
                continue;
            }
            NvMediaImageDestroy(imagePool.imageGroups[i]->imageList[p]);
        }
        delete imagePool.imageGroups[i];
    }
    // Swap imageGroups vector with an equivalent empty vector to force deallocation
    std::vector<NvMediaImageGroup *>().swap(imagePool.imageGroups);
}

void CImageManager::DeleteImages(ImagePool &imagePool)
{
    for (size_t i = 0U; i < imagePool.images.size(); i++) {
        if (imagePool.images[i] == nullptr) {
            LOG_WARN("Attempt to destroy null image\n");
            continue;
        }
        NvMediaImageDestroy(imagePool.images[i]);
    }
    // Swap images vector with an equivalent empty vector to force deallocation
    std::vector<NvMediaImage *>().swap(imagePool.images);
}

SIPLStatus CImageManager::Allocate()
{
    SIPLStatus status = NVSIPL_STATUS_OK;
    for (uint32_t i = 0U; i < MAX_NUM_IMAGE_OUTPUTS; i++) {
        if (m_imagePools[i].enable) {
            status = m_siplCamera->GetImageAttributes(m_index,
                                                      m_imagePools[i].outputType,
                                                      m_imagePools[i].imageAttrs);
            CHK_STATUS_AND_RETURN(status, "INvSIPLCamera::GetImageAttributes()");
            switch (m_imagePools[i].outputType) {
                case INvSIPLClient::ConsumerDesc::OutputType::ICP:
                    status = AllocateImageGroups(m_imagePools[i]);
                    CHK_STATUS_AND_RETURN(status, "CImageManager::AllocateImageGroups()");
                    break;
                case INvSIPLClient::ConsumerDesc::OutputType::ISP0:
                case INvSIPLClient::ConsumerDesc::OutputType::ISP1:
                    status = AllocateImages(m_imagePools[i]);
                    CHK_STATUS_AND_RETURN(status, "CImageManager::AllocateImages()");
                    break;
                default:
                    LOG_ERR("Unexpected output type\n");
                    return NVSIPL_STATUS_ERROR;
            }
        }
    }

    return status;
}

SIPLStatus CImageManager::Register()
{
    SIPLStatus status = NVSIPL_STATUS_OK;
    for (uint32_t i = 0U; i < MAX_NUM_IMAGE_OUTPUTS; i++) {
        if (m_imagePools[i].enable) {
            switch (m_imagePools[i].outputType) {
                case INvSIPLClient::ConsumerDesc::OutputType::ICP:
                    status = m_siplCamera->RegisterImageGroups(m_index,
                                                               m_imagePools[i].imageGroups);
                    CHK_STATUS_AND_RETURN(status, "INvSIPLCamera::RegisterImageGroups()");
                    break;
                case INvSIPLClient::ConsumerDesc::OutputType::ISP0:
                case INvSIPLClient::ConsumerDesc::OutputType::ISP1:
                    status = m_siplCamera->RegisterImages(m_index,
                                                          m_imagePools[i].outputType,
                                                          m_imagePools[i].images);
                    CHK_STATUS_AND_RETURN(status, "INvSIPLCamera::RegisterImages()");
                    break;
                default:
                    LOG_ERR("Unexpected output type\n");
                    return NVSIPL_STATUS_ERROR;
            }
        }
    }

    return status;
}

SIPLStatus CImageManager::Init(INvSIPLCamera *siplCamera, uint32_t index)
{
    SIPLStatus status = NVSIPL_STATUS_OK;
    NvMediaStatus nvmStatus = NVMEDIA_STATUS_OK;
    NvSciError err = NvSciError_Success;

    m_siplCamera = siplCamera;
    m_index = index;

    // TODO: Expose these initialization parameters so that they are configurable
    m_imagePools[0].enable = true;
    m_imagePools[0].outputType = INvSIPLClient::ConsumerDesc::OutputType::ICP;
    m_imagePools[0].size = CAPTURE_IMAGE_POOL_SIZE;
    m_imagePools[1].enable = true;
    m_imagePools[1].outputType = INvSIPLClient::ConsumerDesc::OutputType::ISP0;
    m_imagePools[1].size = ISP_IMAGE_POOL_SIZE;
    m_imagePools[2].enable = true;
    m_imagePools[2].outputType = INvSIPLClient::ConsumerDesc::OutputType::ISP1;
    m_imagePools[2].size = ISP_IMAGE_POOL_SIZE;

    m_nvmDevice = NvMediaDeviceCreate();
    CHK_PTR_AND_RETURN(m_nvmDevice, "NvMediaDeviceCreate");

    err = NvSciBufModuleOpen(&m_sciBufModule);
    CHK_NVSCISTATUS_AND_RETURN(err, "NvSciBufModuleOpen");

    nvmStatus = NvMediaImageNvSciBufInit();
    CHK_NVMSTATUS_AND_RETURN(nvmStatus, "NvMediaImageNvSciBufInit");

    status = Allocate();
    CHK_STATUS_AND_RETURN(status, "CImageManager::Allocate()");

    status = Register();
    CHK_STATUS_AND_RETURN(status, "CImageManager::Register()");

    return status;
}

void CImageManager::Deinit()
{
    for (uint32_t i = 0U; i < MAX_NUM_IMAGE_OUTPUTS; i++) {
        if (m_imagePools[i].enable) {
            switch (m_imagePools[i].outputType) {
                case INvSIPLClient::ConsumerDesc::OutputType::ICP:
                    DeleteImageGroups(m_imagePools[i]);
                    break;
                case INvSIPLClient::ConsumerDesc::OutputType::ISP0:
                case INvSIPLClient::ConsumerDesc::OutputType::ISP1:
                    DeleteImages(m_imagePools[i]);
                    break;
                default:
                    LOG_ERR("Unexpected output type\n");
                    break;
            }
            // Swap surfAllocAttr vector with an equivalent empty vector to force deallocation
            std::vector<NvMediaSurfAllocAttr>().swap(m_imagePools[i].imageAttrs.surfaceAllocAttr);
            for (uint32_t j = 0U; j < m_imagePools[i].sciBufObjs.size(); j++) {
                if (m_imagePools[i].sciBufObjs[j] == nullptr) {
                    LOG_WARN("Attempt to free null NvSciBufObj\n");
                    continue;
                }
                NvSciBufObjFree(m_imagePools[i].sciBufObjs[j]);
            }
            // Swap sciBufObjs vector with an equivalent empty vector to force deallocation
            std::vector<NvSciBufObj>().swap(m_imagePools[i].sciBufObjs);
        }
    }

    NvMediaImageNvSciBufDeinit();

    if (m_sciBufModule != nullptr) {
        NvSciBufModuleClose(m_sciBufModule);
    }

    if (m_nvmDevice != nullptr) {
        NvMediaDeviceDestroy(m_nvmDevice);
    }
}

CImageManager::~CImageManager()
{
    Deinit();
}
