/*
* Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
*
* This code is licensed under the MIT license (MIT) (http://opensource.org/licenses/MIT)
*/

#include "wfdSwapChain.h"
#include <assert.h>

VkResult SwapChainWfd::InitNvSciBufModule()
{
    NvSciError err = NvSciBufModuleOpen(&sciBufModule);
    if (err != NvSciError_Success)
        return VK_ERROR_UNKNOWN;
    return VK_SUCCESS;
}

VkResult SwapChainWfd::Init(VkInstance instance, VkPhysicalDevice physDevice, VkDevice dev, const uint32_t width, const uint32_t height)
{
    m_instance = instance;
    m_physDevice = physDevice;
    m_device = dev;

    bool ret = LoadSymbols();
    assert(ret);

    VkResult err = InitNvSciBufModule();
    assert(err == VK_SUCCESS);

    m_width = width;
    m_height = height;

    return err;
}

bool SwapChainWfd::LoadSymbols()
{
    fpGetDeviceProcAddr = (PFN_vkGetDeviceProcAddr)vkGetInstanceProcAddr(
                                m_instance, "vkGetDeviceProcAddr");
    if (!fpGetDeviceProcAddr) {
        return false;
    }
    if (fpGetDeviceProcAddr) {
        fpGetMemorySciBufNV = (PFN_vkGetMemorySciBufNV)fpGetDeviceProcAddr(
                                m_device, "vkGetMemorySciBufNV");
    }
    fpGetPhysicalDeviceSciBufAttributesNV =
            (PFN_vkGetPhysicalDeviceSciBufAttributesNV)vkGetInstanceProcAddr(
                    m_instance, "vkGetPhysicalDeviceSciBufAttributesNV");

    if (fpGetMemorySciBufNV && fpGetPhysicalDeviceSciBufAttributesNV) {
        return true;
    } else {
        return false;
    }
}

SwapChainWfd::~SwapChainWfd()
{
    NvSciBufObjFree(m_nvsciBufObjs[0]);
    NvSciBufObjFree(m_nvsciBufObjs[1]);
}

void SwapChainWfd::InitWFD()
{
    m_wfdRes.WfdInit(m_width, m_height);
    m_wfdRes.BindNvSciBufObj(m_nvsciBufObjs.data(), IMAGE_COUNT);
}

VkResult SwapChainWfd::getAttrListFromVkImage(const VkImageCreateInfo& imageInfo,
                            const VkMemoryRequirements& memReqs,
                            NvSciBufAttrList attrList)
{
    // Check the alignment has to be power of 2
    assert(__builtin_popcountl(memReqs.alignment) == 1);
    NvSciBufType bufType = NvSciBufType_Image;
    NvSciBufAttrValAccessPerm perm = NvSciBufAccessPerm_ReadWrite;  // Controlled by application
    NvSciBufAttrValImageLayoutType layout = imageInfo.tiling == VK_IMAGE_TILING_OPTIMAL ?
                            NvSciBufImage_BlockLinearType : NvSciBufImage_PitchLinearType;
    NvSciBufAttrValImageScanType planescantype = NvSciBufScan_ProgressiveType;
    uint32_t planeCount = 1;
    NvSciBufAttrValColorFmt colorFormat = NvSciColor_A8R8G8B8;
    uint32_t height = imageInfo.extent.height;
    uint32_t width = imageInfo.extent.width;
    bool needCpuAccess = memReqs.memoryTypeBits & VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT;

    NvSciBufAttrKeyValuePair pairArray[] = {
        { NvSciBufImageAttrKey_Layout, (void*) &layout, sizeof(layout) },
        { NvSciBufImageAttrKey_PlaneCount, (void*) &planeCount, sizeof(planeCount) },
        { NvSciBufImageAttrKey_PlaneColorFormat, (void*) &colorFormat, sizeof(colorFormat) },
        { NvSciBufImageAttrKey_PlaneHeight, (void*) &height, sizeof(height) },
        { NvSciBufImageAttrKey_PlaneWidth, (void*) &width, sizeof(width) },
        { NvSciBufImageAttrKey_ScanType, (void*) &planescantype, sizeof(planescantype) },
        { NvSciBufGeneralAttrKey_Types, (void*) &bufType, sizeof(bufType) },
        { NvSciBufGeneralAttrKey_NeedCpuAccess, (void*) &needCpuAccess, sizeof(needCpuAccess) },
        { NvSciBufGeneralAttrKey_RequiredPerm, (void*) &perm, sizeof(perm) },
    };

    // Application fills the public key-value pairs
    NvSciError err = NvSciBufAttrListSetAttrs(attrList, pairArray, sizeof(pairArray)/sizeof(NvSciBufAttrKeyValuePair));
    if (err != NvSciError_Success)
        return VK_ERROR_UNKNOWN;

    // Get GPU id
    VkResult res = fpGetPhysicalDeviceSciBufAttributesNV(m_physDevice, attrList);
    if (res != VK_SUCCESS)
        return VK_ERROR_INITIALIZATION_FAILED;

    return VK_SUCCESS;
}

VkResult SwapChainWfd::SetupNvSciBufVkDeviceMemory(VkImageCreateInfo imageInfo, VkMemoryRequirements memReqs)
{
    printf("vkimage mem reqs, type: %u, align = %lu, size = %lu\n",
                    memReqs.memoryTypeBits,
                    memReqs.alignment,
                    memReqs.size);
    // setup NvSciBufAttrList
    NvSciBufAttrList unreconciledList[2], reconciledList, conflictList;
    NvSciBufAttrListCreate(sciBufModule, &unreconciledList[0]);
    NvSciBufAttrListCreate(sciBufModule, &unreconciledList[1]);
    VkResult err = getAttrListFromVkImage(imageInfo, memReqs, unreconciledList[0]);
    assert(err == VK_SUCCESS);
    bool ret = WFDResources::InitializeRGB(m_width, m_height, unreconciledList[1],
                    imageInfo.tiling == VK_IMAGE_TILING_LINEAR ?
                    NvSciBufImage_PitchLinearType : NvSciBufImage_BlockLinearType);
    assert(ret == true);
    NvSciError sciErr = NvSciBufAttrListReconcile(unreconciledList, 2, &reconciledList, &conflictList);
    assert(sciErr == NvSciError_Success);
    NvSciBufAttrListFree(unreconciledList[0]);
    NvSciBufAttrListFree(unreconciledList[1]);

    NvSciBufObjAlloc(reconciledList, &m_nvsciBufObjs[0]);
    NvSciBufObjAlloc(reconciledList, &m_nvsciBufObjs[1]);

    // Update allocation size
    NvSciBufAttrKeyValuePair pairArray[] = {
        { NvSciBufImageAttrKey_Size, nullptr, 0 },
        { NvSciBufImageAttrKey_Alignment, nullptr, 0 },
        { NvSciBufImageAttrKey_PlanePitch, nullptr, 0 }
    };
    NvSciBufAttrListGetAttrs(reconciledList, pairArray, sizeof(pairArray) / sizeof(NvSciBufAttrKeyValuePair));
    memReqs.size = *static_cast<const uint64_t*>(pairArray[0].value);
    memReqs.alignment = *static_cast<const uint64_t*>(pairArray[1].value);
    uint64_t pitch = *static_cast<const uint64_t*>(pairArray[2].value);
    printf("NvSciBuf image size = %lu, alignement = %lu, pitch = %lu\n", memReqs.size, memReqs.alignment, pitch);

    VkImportMemorySciBufInfoNV importSciBufInfo = {
        .sType = VK_STRUCTURE_TYPE_IMPORT_MEMORY_SCI_BUF_INFO_NV,
        .pNext = nullptr,
        .handleType = VK_EXTERNAL_MEMORY_HANDLE_TYPE_SCI_BUF_BIT_NV,
        .handle = m_nvsciBufObjs[0]
    };

    VkMemoryAllocateInfo memAllocInfo = {
        .sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
        .pNext = &importSciBufInfo,
        .allocationSize  = 0, // ignore the allocate size when importing NvSciBufObj
        .memoryTypeIndex = 0  // local bit
    };

    err = vkAllocateMemory(m_device, &memAllocInfo, nullptr, &m_imageMems[0]);
    assert(err == VK_SUCCESS);

    importSciBufInfo.handle = m_nvsciBufObjs[1];
    err = vkAllocateMemory(m_device, &memAllocInfo, nullptr, &m_imageMems[1]);
    assert(err == VK_SUCCESS);
    return err;
}

void SwapChainWfd::CreateImages(VkImage* images)
{
    VkResult err;
    VkExternalMemoryImageCreateInfo externalMemInfo = {
        .sType = VK_STRUCTURE_TYPE_EXTERNAL_MEMORY_IMAGE_CREATE_INFO,
        .pNext = nullptr,
        .handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_SCI_BUF_BIT_NV
    };

    VkImageCreateInfo imageInfo = {};
    imageInfo.sType             = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
    imageInfo.pNext             = &externalMemInfo;
    imageInfo.imageType         = VK_IMAGE_TYPE_2D;
    imageInfo.format            = VK_FORMAT_B8G8R8A8_UNORM;
    imageInfo.extent.width      = m_width;
    imageInfo.extent.height     = m_height;
    imageInfo.extent.depth      = 1;
    imageInfo.mipLevels         = 1;
    imageInfo.arrayLayers       = 1;
    imageInfo.samples           = VK_SAMPLE_COUNT_1_BIT;
    imageInfo.tiling            = VK_IMAGE_TILING_OPTIMAL;
    imageInfo.usage             = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;

    // Create 2 images as back/front buffer
    err = vkCreateImage(m_device, &imageInfo, nullptr, &images[0]);
    err = vkCreateImage(m_device, &imageInfo, nullptr, &images[1]);
    assert(err == VK_SUCCESS);

    VkMemoryRequirements memReqs;
    vkGetImageMemoryRequirements(m_device, images[0], &memReqs);

    // Allocate NvSciBufObj and import it to VkDeviceMemory.
    err = SetupNvSciBufVkDeviceMemory(imageInfo, memReqs);

    // Bind image to memory.
    err = vkBindImageMemory(m_device, images[0], m_imageMems[0], 0);
    err = vkBindImageMemory(m_device, images[1], m_imageMems[1], 0);
    (void) err;
}

void SwapChainWfd::GetSwapChainImages(VkImage* images, size_t imageCount)
{
    assert(imageCount >= IMAGE_COUNT);
    CreateImages(images);

    InitWFD();
}

void SwapChainWfd::GetNexImage(uint32_t* currentBufferIdx, VkSemaphore* pCompletePresentSemaphore)
{
    // Change between [0, 1].
    m_currentImageIdx = m_currentImageIdx ^ 1;
    *currentBufferIdx = m_currentImageIdx;
}

void SwapChainWfd::PresentImage(VkSemaphore* pCompleteRenderSemaphore)
{
    m_wfdRes.FlipSubmit(m_currentImageIdx);
}