// Copyright (c) 2019-2021, 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.
//
#pragma once

#include <carb/Defines.h>
#include <carb/graphics/Graphics.h>

#include <omni/gpucompute/GpuCompute.h>

using namespace carb::graphics;
using std::unique_ptr;

namespace carb
{
namespace flatcache
{
// If this is an array-of-arrays:
// array.cpuData - array of CPU pointers on CPU
// gpuPointerArray->cpuData() - array of GPU pointers on CPU
inline void PathToAttributesMap::enableD3dGpuRead(MirroredArray& array,
                                                  const size_t* elemToArraySize,
                                                  MirroredArray* elemToArrayCpuCapacity,
                                                  MirroredArray* elemToArrayGpuCapacity,
                                                  MirroredArray* gpuPointerArray)
{
    using omni::gpucompute::MemcpyKind;

    bool& cpuValid = array.cpuValid;
    bool& gpuValid = array.gpuValid;
    bool& gpuAllocedWithCuda = array.gpuAllocedWithCuda;
    uint8_t* cpuArray = array.cpuData();
    uint8_t*& gpuArray = array.gpuArray;

    if (gpuValid)
    {
        // Nothing to do
    }
    else if (!gpuValid && cpuValid)
    {
        const TypeC type = array.type;
        const Typeinfo &typeInfo = getTypeInfo(type);

        std::lock_guard<MirroredArray::AttributeMutex> hostDeviceLock(array.attributeMutex);
        // If each element is an array, then they could be of different sizes
        // So alloc and memcpy each one
        if (typeInfo.isArray)
        {
            size_t elemCount = array.count;
            uint8_t** elemToArrayCpuData = (uint8_t**)cpuArray;
            uint8_t** elemToArrayGpuData = (uint8_t**)gpuPointerArray->cpuData();

            for (size_t elem = 0; elem != elemCount; elem++)
            {
                // Make sure that the dest (GPU) buffer is large enough
                const uint8_t* const& cpuData = elemToArrayCpuData[elem]; // src
                size_t& destCapacity = reinterpret_cast<size_t*>(elemToArrayGpuCapacity->cpuData())[elem];
                size_t desiredElemCount = elemToArraySize[elem];
                destructiveResizeIfNecessaryGPU(
                    *gpuPointerArray, elem, destCapacity, desiredElemCount, typeInfo.arrayElemSize, platform.gpuD3dVk, platform.gpuD3dVkCtx);

                // Copy from CPU to GPU
                if (desiredElemCount != 0 && cpuData)
                {
                    uint8_t*& gpuData = elemToArrayGpuData[elem]; // dest
                    size_t copyByteCount = desiredElemCount * typeInfo.arrayElemSize;
                    platform.gpuD3dVk->memcpy(*platform.gpuD3dVkCtx, gpuData, cpuData, copyByteCount, MemcpyKind::hostToDevice);
                }
                else if (desiredElemCount != 0 && !cpuData)
                {
                    printf("Warning: GPU read access requested, CPU is valid but not allocated\n");
                }
            }

            // We don't need to copy the outer array to GPU here.
            // In D3dVk, the outer array is currently a CPU array of descriptors that we copy to
            // a kernel calls descriptor set at dispatch time
        }
        else
        {
            // Copy the outer array from CPU to GPU
            size_t byteCount = array.size();
            allocGpuMemIfNecessary(array, byteCount, typeInfo.size, platform.gpuD3dVk, platform.gpuD3dVkCtx);

            uint8_t* cpuArray = array.cpuData();
            platform.gpuD3dVk->memcpy(*platform.gpuD3dVkCtx, gpuArray, cpuArray, byteCount, MemcpyKind::hostToDevice);
        }

        // New state
        cpuValid = true;
        gpuValid = true;
        gpuAllocedWithCuda = false;
    }
}

inline void PathToAttributesMap::enableD3dGpuWrite(PathToAttributesMap::MirroredArray& array,
                                                   const size_t* elemToArraySize,
                                                   PathToAttributesMap::MirroredArray* elemToArrayCpuCapacity,
                                                   PathToAttributesMap::MirroredArray* elemToArrayGpuCapacity,
                                                   PathToAttributesMap::MirroredArray* arrayGpuDataArray)
{
    log("begin enableGpuWrite\n");
    bool& usdValid = array.usdValid;
    bool& cpuValid = array.cpuValid;
    bool& gpuValid = array.gpuValid;
    bool& gpuAllocedWithCuda = array.gpuAllocedWithCuda;

    const TypeC type = array.type;
    const Typeinfo &typeInfo = getTypeInfo(type);

    std::lock_guard<MirroredArray::AttributeMutex> hostDeviceLock(array.attributeMutex);
    if (!gpuValid)
    {
        size_t byteCount = array.size();
        allocGpuMemIfNecessary(array, byteCount, typeInfo.size, platform.gpuD3dVk, platform.gpuD3dVkCtx);

        // Array-valued elements are lazily allocated, meaning they are only
        // resized when write access is requested.
        // Write access has been requested, so resize if necessary
        if (typeInfo.isArray)
        {
            size_t elemCount = array.count;

            for (size_t elem = 0; elem != elemCount; elem++)
            {
                size_t& gpuCapacity = reinterpret_cast<size_t*>(elemToArrayGpuCapacity->cpuData())[elem];
                size_t desiredElemCount = elemToArraySize[elem];

                resizeIfNecessaryGPU(
                    *arrayGpuDataArray, elem, gpuCapacity, desiredElemCount, typeInfo.arrayElemSize, platform.gpuD3dVk, platform.gpuD3dVkCtx);
            }

            // Upload of allocated pointers to GPU happens outside this function
        }
    }

    // New state
    usdValid = false;
    cpuValid = false;
    gpuValid = true;
    gpuAllocedWithCuda = false;

    log("end enableGpuWrite\n\n");
}

inline const void* PathToAttributesMap::getArrayRdD3d(const Bucket& bucket, const TokenC& attrName, NameSuffix suffix)
{
    BucketId bucketId = findBucketId(bucket);
    if (bucketId == kInvalidBucketId)
        return nullptr;
    // We don't set dirty indices here because this method gives read-only access
    return getArraySpanC(bucketId, attrName, D3dVkReadConfig(), suffix).array.ptr;
}

inline void* PathToAttributesMap::getArrayWrD3d(const Bucket& bucket, const TokenC& attrName, NameSuffix suffix)
{
    BucketId bucketId = findBucketId(bucket);
    if (bucketId == kInvalidBucketId)
        return nullptr;
    ArrayAndDirtyIndices array = getArraySpanC(bucketId, attrName, D3dVkReadWriteConfig(), suffix);
    setArrayDirty(array);
    return array.array.ptr;
}

inline void* PathToAttributesMap::getArrayD3d(const Bucket& bucket, const TokenC& attrName, NameSuffix suffix)
{
    BucketId bucketId = findBucketId(bucket);
    if (bucketId == kInvalidBucketId)
        return nullptr;
    ArrayAndDirtyIndices array = getArraySpanC(bucketId, attrName, D3dVkReadWriteConfig(), suffix);
    setArrayDirty(array);
    return array.array.ptr;
}

#if 0
// If array of values, return the Buffer* that was returned by malloc
// If array of arrays, return array of Buffer* for each element array
inline std::pair<void**, size_t> PathToAttributesMap::getArrayD3d(const Bucket& bucket, const TokenC& attrName)
{
    std::pair<void**, size_t> retval = { nullptr, 0 };
    BucketImpl& bucketImpl = buckets[bucket];
    auto iter = bucketImpl.arrays.find({ pxr::TfType(), TypeC(), attrName });
    bool found = (iter != bucketImpl.arrays.end());
    if (found)
    {
        bool isTag = (typeToInfo[iter->first.type].size == 0);

        if (!isTag)
        {
            pxr::TfType type = iter->first.type;
            size_t elemSize = typeToInfo[type].size;
            size_t arrayElemSize = typeToInfo[type].arrayElemSize;

            // Read enable must come before write enable
            enableD3dGpuRead(iter->second, elemSize, arrayElemSize);
            enableD3dGpuWrite(iter->second, elemSize, arrayElemSize);

            retval.first = iter->second.d3dArrays.data();
            retval.second = iter->second.d3dArrays.size();
        }
        else if (isTag)
        {
            // If is a tag, then array.data() will be zero, so set special value
            // to distinguish from tag absent case
            retval.first = (void**)-1;
        }
    }

    return retval;
}
#endif

inline omni::gpucompute::GpuPointer PathToAttributesMap::getAttributeD3d(const PathC& path, const TokenC& attrName)
{
    bool present; // Whether this path has a bucket
    BucketId bucketId; // Pointer to the bucket if it does
    size_t element; // Index corresponding to path in this bucket's arrays
    std::tie(present, bucketId, element) = getPresentAndBucketAndElement(path);
    if (!present)
        return { nullptr, 0, 0 };

    // TODO: Get rid of double hash lookup below (getArraySpanC + explicit call)

    ArrayAndDirtyIndices arrayAndDirtyIndices = getArraySpanC(bucketId, attrName, D3dVkReadWriteConfig());
    setArrayElementDirty(arrayAndDirtyIndices, element);
    void* array = arrayAndDirtyIndices.array.ptr;

    if (array != nullptr)
    {
        // Get elemSize
        const BucketImpl* bucketImplPtr = buckets.find(bucketId);
        if (!bucketImplPtr)
            return { nullptr, 0, 0 };

        const AttrName name{ attrName, NameSuffix::none };
        const MirroredArray *valuesArray;
        if (!bucketImplPtr->scalarAttributeArrays.find(name, &valuesArray))
        {
            const ArrayAttributeArray *arrayAttributeArray;
            if (bucketImplPtr->arrayAttributeArrays.find(name, &arrayAttributeArray))
            {
                valuesArray = &arrayAttributeArray->values;
            }
            else
            {
                return { nullptr, 0, 0 };
            }
        }
        assert(valuesArray);

        const Typeinfo &typeinfo = getTypeInfo(valuesArray->type);
        const bool isArrayOfArray = typeinfo.isArray;
        const size_t elemSize = typeinfo.size;

        if (!isArrayOfArray)
        {
            return { array, element * elemSize, elemSize };
        }
        else if (isArrayOfArray)
        {
            // For arrays of arrays we return the Buffer* of the inner array
            uint8_t* const* elemToArrayData = (uint8_t* const*)array;
            return { elemToArrayData[element], 0, 0 };
        }
    }

    return { nullptr, 0, 0 };
}
}
}
