// 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 "PathToAttributesMap.h"

#include <carb/profiler/Profile.h>

namespace carb
{
namespace flatcache
{

const uint64_t kProfilerMask = 1;

// 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::enableGpuRead(MirroredArray& array,
                                               const size_t* elemToArraySize,
                                               MirroredArray* elemToArrayCpuCapacity,
                                               MirroredArray* elemToArrayGpuCapacity,
                                               MirroredArray* gpuPointerArray)
{
    CARB_PROFILE_ZONE(kProfilerMask, "enableGpuRead");
    using omni::gpucompute::MemcpyKind;

    log("begin enableGpuRead\n");
    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 (cpuValid)
    {
        std::lock_guard<MirroredArray::AttributeMutex> hostDeviceLock(array.attributeMutex);
        const TypeC type = array.type;
        const Typeinfo &typeInfo = getTypeInfo(type);

        if (typeInfo.isArray)
        {
            const size_t elemCount = array.count;

            uint8_t** cpuPointers = reinterpret_cast<uint8_t**>(cpuArray);
            uint8_t** gpuPointers = reinterpret_cast<uint8_t**>(gpuPointerArray->cpuData());

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

                if(gpuCapacity != desiredCapacity)
                {
                    destructiveResizeIfNecessaryGPU(*gpuPointerArray, elem, gpuCapacity, desiredCapacity, typeInfo.arrayElemSize, platform.gpuCuda, platform.gpuCudaCtx);
                }

                const size_t copyByteCount = std::min(desiredCapacity, cpuCapacity) * typeInfo.arrayElemSize;

                if (copyByteCount > 0)
                {
                    void* cpuPointer = cpuPointers[elem];
                    void* gpuPointer = gpuPointers[elem];
                    CARB_ASSERT(cpuPointer);
                    CARB_ASSERT(gpuPointer);

                    platform.gpuCuda->memcpyAsync( *platform.gpuCudaCtx, gpuPointer, cpuPointer, copyByteCount, MemcpyKind::hostToDevice);
                }
            }

            gpuPointerArray->cpuValid = true;
        }
        else
        {
            // Copy the outer array from CPU to GPU
            size_t byteCount = array.size();
            allocGpuMemIfNecessary(array, byteCount, typeInfo.size, platform.gpuCuda, platform.gpuCudaCtx);

            log("array values: to GPU\n");
            uint8_t* cpuArray = array.cpuData();
            carb::profiler::ZoneId zoneId = CARB_PROFILE_BEGIN(kProfilerMask, "outer array values");
            platform.gpuCuda->memcpyAsync(*platform.gpuCudaCtx, gpuArray, cpuArray, byteCount, MemcpyKind::hostToDevice);
            CARB_PROFILE_END(kProfilerMask, zoneId);
        }

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

inline void PathToAttributesMap::enableGpuWrite(PathToAttributesMap::MirroredArray& array,
                                                const size_t* elemToArraySize,
                                                PathToAttributesMap::MirroredArray* elemToArrayCpuCapacity,
                                                PathToAttributesMap::MirroredArray* elemToArrayGpuCapacity,
                                                PathToAttributesMap::MirroredArray* arrayGpuDataArray)
{
    CARB_PROFILE_ZONE(kProfilerMask, "enableGpuWrite");
    using omni::gpucompute::MemcpyKind;

    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 (!typeInfo.isArray && !gpuValid)
    {
        size_t byteCount = array.size();
        allocGpuMemIfNecessary(array, byteCount, typeInfo.size, platform.gpuCuda, platform.gpuCudaCtx);
    }
    else if (typeInfo.isArray)
    {
        // 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
        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.gpuCuda, platform.gpuCudaCtx);
        }

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

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

    if (elemToArrayCpuCapacity) elemToArrayCpuCapacity->usdValid = false;
    if (elemToArrayGpuCapacity) elemToArrayGpuCapacity->usdValid = false;
    if (arrayGpuDataArray) arrayGpuDataArray->usdValid = false;

}

inline ConstSpanC PathToAttributesMap::getArraySpanRdGpuC(const Bucket& bucket, const TokenC& attrName, NameSuffix suffix)
{
    APILOGGER("getArraySpanRdGpuC", apiLogEnabled, attrName);

    BucketId bucketId = findBucketId(bucket);
    if (bucketId == kInvalidBucketId)
        return { nullptr, 0, 0 };
    // We don't set dirty indices here because this method gives read-only access
    return getArraySpanC(bucketId, attrName, CudaReadConfig(), suffix).array;
}

inline const void* PathToAttributesMap::getArrayRdGpuC(const Bucket& bucket, const TokenC& attrName, NameSuffix suffix)
{
    APILOGGER("getArrayRdGpuC", apiLogEnabled, attrName);

    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, CudaReadConfig(), suffix).array.ptr;
}

inline SpanC PathToAttributesMap::getArrayGpuC(BucketId bucketId, const TokenC& attrName, NameSuffix suffix)
{
    APILOGGER("getArrayGpuC", apiLogEnabled, attrName);

    ArrayAndDirtyIndices arrayAndDirtyIndices = getArraySpanC(bucketId, attrName, CudaReadWriteConfig(), suffix);
    setArrayDirty(arrayAndDirtyIndices);

    return arrayAndDirtyIndices.array;
}

inline ConstSpanC PathToAttributesMap::getArrayRdGpuC(BucketId bucketId, const TokenC& attrName, NameSuffix suffix)
{
    APILOGGER("getArrayRdGpuC", apiLogEnabled, attrName);

    const ArrayAndDirtyIndices arrayAndDirtyIndices = getArraySpanC(bucketId, attrName, CudaReadConfig(), suffix);
    // We don't set dirty indices here because this method gives read-only access
    return arrayAndDirtyIndices.array;
}

inline SpanC PathToAttributesMap::getArrayWrGpuC(BucketId bucketId, const TokenC& attrName, NameSuffix suffix)
{
    APILOGGER("getArrayWrC", apiLogEnabled, attrName);

    ArrayAndDirtyIndices arrayAndDirtyIndices = getArraySpanC(bucketId, attrName, CudaWriteConfig(), suffix);
    setArrayDirty(arrayAndDirtyIndices);

    return arrayAndDirtyIndices.array;
}

inline SpanC PathToAttributesMap::getArraySpanWrGpuC(const Bucket& bucket, const TokenC& attrName, NameSuffix suffix)
{
    APILOGGER("getArraySpanWrGpuC", apiLogEnabled, attrName);

    BucketId bucketId = findBucketId(bucket);
    if (bucketId == kInvalidBucketId)
        return { nullptr, 0, 0 };
    ArrayAndDirtyIndices array = getArraySpanC(bucketId, attrName, CudaWriteConfig(), suffix);
    setArrayDirty(array);
    return array.array;
}

inline void* PathToAttributesMap::getArrayWrGpuC(const Bucket& bucket, const TokenC& attrName, NameSuffix suffix)
{
    APILOGGER("getArrayWrGpuC", apiLogEnabled, attrName);

    // Get write-only GPU access
    BucketId bucketId = findBucketId(bucket);
    if (bucketId == kInvalidBucketId)
        return nullptr;
    ArrayAndDirtyIndices array = getArraySpanC(bucketId, attrName, CudaWriteConfig(), suffix);
    setArrayDirty(array);
    return array.array.ptr;
}

inline SpanC PathToAttributesMap::getArraySpanGpuC(const Bucket& bucket, const TokenC& attrName, NameSuffix suffix)
{
    APILOGGER("getArraySpanGpuC", apiLogEnabled, attrName);

    BucketId bucketId = findBucketId(bucket);
    if (bucketId == kInvalidBucketId)
        return { nullptr, 0, 0 };
    ArrayAndDirtyIndices array = getArraySpanC(bucketId, attrName, CudaReadWriteConfig(), suffix);
    setArrayDirty(array);
    return array.array;
}

inline void* PathToAttributesMap::getArrayGpuC(const Bucket& bucket, const TokenC& attrName, NameSuffix suffix)
{
    APILOGGER("getArrayGpuC", apiLogEnabled, attrName);

    BucketId bucketId = findBucketId(bucket);
    if (bucketId == kInvalidBucketId)
        return nullptr;
    ArrayAndDirtyIndices array = getArraySpanC(bucketId, attrName, CudaReadWriteConfig(), suffix);
    setArrayDirty(array);
    return array.array.ptr;
}

inline SpanC PathToAttributesMap::getAttributeGpuC(const PathC& path, const TokenC& attrName, PtrToPtrKind ptrToPtrKind, NameSuffix suffix)
{
    APILOGGER("getAttributeGpuC", apiLogEnabled, path, 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 };

    ArrayAndDirtyIndices array = getArraySpanC(bucketId, attrName, CudaReadWriteConfig().withPtrToPtrKind(ptrToPtrKind), suffix);
    setArrayElementDirty(array, element);
    return getArrayElementPtr(array.array, element);
}

inline ConstSpanC PathToAttributesMap::getAttributeRdGpuC(const PathC& path, const TokenC& attrName, PtrToPtrKind ptrToPtrKind, NameSuffix suffix)
{
    APILOGGER("getAttributeRdGpuC", apiLogEnabled, path, attrName);

    bool present;
    BucketId bucketId;
    size_t element;
    std::tie(present, bucketId, element) = getPresentAndBucketAndElement(path);
    if (!present)
        return { nullptr, 0, 0 };

    ConstSpanC array = getArraySpanC(bucketId, attrName, CudaReadConfig().withPtrToPtrKind(ptrToPtrKind), suffix).array;
    // We don't set dirty indices here because this method gives read-only access
    return getArrayElementPtr(array, element);
}

inline SpanC PathToAttributesMap::getAttributeWrGpuC(const PathC& path, const TokenC& attrName, PtrToPtrKind ptrToPtrKind, NameSuffix suffix)
{
    APILOGGER("getAttributeWrGpuC", apiLogEnabled, path, attrName);

    // Writing an element is a RMW on the whole array, so use getArrayGpu instead of getArrayGpuWr
    bool present;
    BucketId bucketId;
    size_t element;
    std::tie(present, bucketId, element) = getPresentAndBucketAndElement(path);
    if (!present)
        return { nullptr, 0, 0 };

    ArrayAndDirtyIndices array = getArraySpanC(bucketId, attrName, CudaReadWriteConfig().withPtrToPtrKind(ptrToPtrKind), suffix);
    setArrayElementDirty(array, element);
    return getArrayElementPtr(array.array, element);
}

// Typed accessors
template <typename T>
inline const T* PathToAttributesMap::getArrayRdGpu(const Bucket& bucket, const TokenC& attrName)
{
    APILOGGER("getArrayRdGpu", apiLogEnabled, attrName);

    // TODO: check that T is the correct type
    return reinterpret_cast<const T*>(getArrayRdGpuC(bucket, attrName));
}

template <typename T>
inline T* PathToAttributesMap::getArrayWrGpu(const Bucket& bucket, const TokenC& attrName)
{
    APILOGGER("getArrayWrGpu", apiLogEnabled, attrName);

    // TODO: check that T is the correct type
    return reinterpret_cast<T*>(getArrayWrGpuC(bucket, attrName));
}

template <typename T>
inline T* PathToAttributesMap::getArrayGpu(const Bucket& bucket, const TokenC& attrName)
{
    APILOGGER("getArrayGpu", apiLogEnabled, attrName);

    // TODO: check that T is the correct type
    return reinterpret_cast<T*>(getArrayGpuC(bucket, attrName));
}

template <typename T>
inline T* PathToAttributesMap::getAttributeGpu(const PathC& path, const TokenC& attrName, PtrToPtrKind ptrToPtrKind)
{
    APILOGGER("getAttributeGpu", apiLogEnabled, path, attrName);

    // TODO: check that T is the correct type
    return reinterpret_cast<T*>(getAttributeGpuC(path, attrName, ptrToPtrKind).ptr);
}

template <typename T>
inline const T* PathToAttributesMap::getAttributeRdGpu(const PathC& path, const TokenC& attrName, PtrToPtrKind ptrToPtrKind)
{
    APILOGGER("getAttributeRdGpu", apiLogEnabled, path, attrName);

    // TODO: check that T is the correct type
    return reinterpret_cast<const T*>(getAttributeRdGpuC(path, attrName, ptrToPtrKind).ptr);
}

template <typename T>
inline T* PathToAttributesMap::getAttributeWrGpu(const PathC& path, const TokenC& attrName, PtrToPtrKind ptrToPtrKind)
{
    APILOGGER("getAttributeWrGpu", apiLogEnabled, path, attrName);

    // TODO: check that T is the correct type
    return reinterpret_cast<T*>(getAttributeWrGpuC(path, attrName, ptrToPtrKind).ptr);
}

}
}
