// Copyright (c) 2021-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.
//

// clang-format off
#include "UsdPCH.h"
// clang-format on

#include <OgnSetGatheredAttributeDatabase.h>

#include <omni/graph/core/IGatherPrototype.h>
#include <carb/flatcache/FlatCache.h>

using namespace carb::flatcache;

namespace omni
{
namespace graph
{
namespace core
{

// WARNING!
// The following code uses low-level ABI functionality and should not be copied for other purposes when such
// low level access is not required. Please use the OGN-generated API whenever possible.

#define RETURN_TRUE_EXEC \
{\
db.outputs.execOut() = kExecutionAttributeStateEnabled;\
return true;\
}

class OgnSetGatheredAttribute
{
public:
    static bool compute(OgnSetGatheredAttributeDatabase& db)
    {
        auto& nodeObj = db.abi_node();
        const INode& iNode = *nodeObj.iNode;
        const IGraphContext& iContext = *db.abi_context().iContext;
        const omni::graph::core::IGatherPrototype* iGatherPrototype =
            carb::getCachedInterface<omni::graph::core::IGatherPrototype>();

        const auto& value = db.inputs.value();
        const auto& mask = db.inputs.mask();

        NameToken attributeName = db.inputs.name();

        const char* attributeNameStr = db.tokenToString(attributeName);
        if (!attributeNameStr || strlen(attributeNameStr) == 0)
            RETURN_TRUE_EXEC

        GatherId gatherId = static_cast<GatherId>(db.inputs.gatherId());

        if (gatherId == kInvalidGatherId)
            RETURN_TRUE_EXEC

        if (!mask.empty() && value.size() > 1 && mask.size() != value.size())
        {
            db.logError("The length of the write mask (%zd) does not match the length of the value (%zd)",
                mask.size(), value.size());
            return false;
        }

        BucketId const* buckets{ nullptr };
        size_t numBuckets{ 0 };

        if (!iGatherPrototype->getGatheredBuckets(db.abi_context(), gatherId, buckets, numBuckets))
        {
            db.logError("Could not get gathered bucket list for Gather %zd", gatherId);
            return false;
        }

        if (numBuckets == 0)
        {
            db.logError("Gathered bucket list is empty for Gather %zd", gatherId);
            return false;
        }

        Type elementType;
        size_t elementSize{ 0 };
        if (!iGatherPrototype->getGatheredType(db.abi_context(), gatherId, attributeName, elementType, elementSize))
        {
            db.logError("Could not determine gathered type");
            return false;
        }

        if (elementType.arrayDepth > 0)
        {
            db.logError("Gathering Array Type %s is not yet supported", elementType.getOgnTypeName().c_str());
            return false;
        }

        if (elementSize == 0)
        {
            db.logError("The element type %s has zero size", elementType.getOgnTypeName().c_str());
            return false;
        }

        Type srcDataType = db.inputs.value().type();

        bool srcIsScaler = srcDataType.arrayDepth == 0;

        if (!srcIsScaler)
        {
            // A scaler value will be broadcast
            --srcDataType.arrayDepth;
        }

        if (!srcDataType.compatibleRawData(elementType))
        {
            db.logWarning("Attribute %s is not compatible with type '%s'", elementType.getTypeName().c_str(),
                          srcDataType.getTypeName().c_str());
            RETURN_TRUE_EXEC
        }

        // determine the length of the content
        size_t totalPrimCount{ 0 };
        size_t inputArraySize = value.size();

        for (size_t i = 0; i < numBuckets; ++i)
        {
            BucketId bucketId = buckets[i];
            size_t primCount = 0;
            (void)db.abi_context().iContext->getBucketArray(db.abi_context(), bucketId, attributeName, primCount);

            totalPrimCount += primCount;
        }

        PathBucketIndex const* repeatedPaths{ nullptr };
        size_t numRepeatedPaths{ 0 };
        if (!iGatherPrototype->getGatheredRepeatedPaths(db.abi_context(), gatherId, repeatedPaths, numRepeatedPaths))
        {
            db.logError("Could not get repeated paths list for Gather %zd", gatherId);
            return false;
        }

        if (inputArraySize > 1 && totalPrimCount + numRepeatedPaths != inputArraySize)
        {
            db.logError(
                "Given value of length %zd is not equal to the number of gathered attributes (%zd)", inputArraySize, totalPrimCount + numRepeatedPaths);
            return false;
        }

        //printf("Setting %zd prims worth (%zd bytes) to %s\n", totalPrimCount, totalPrimCount * elementSize,
        //       attributeNameStr);

        // Finally, we copy the data from the attribute to the buckets
        const IAttributeData& iAttributeData = *db.abi_context().iAttributeData;
        AttributeObj inputAttr =
            nodeObj.iNode->getAttributeByToken(nodeObj, OgnSetGatheredAttributeAttributes::inputs::value.m_token);
        ConstAttributeDataHandle inputHandle = inputAttr.iAttribute->getConstAttributeDataHandle(inputAttr);

        ConstRawPtr srcPtr{ nullptr };
        {
            const void** out = nullptr;
            void** outPtr = reinterpret_cast<void**>(&out);
            iAttributeData.getDataR((const void**)outPtr, db.abi_context(), &inputHandle, 1);

            if (srcIsScaler)
                srcPtr = reinterpret_cast<ConstRawPtr>(out);
            else
                srcPtr = reinterpret_cast<ConstRawPtr>(*out);
        }

        if (!srcPtr)
        {
            // No data to read
            RETURN_TRUE_EXEC
        }

        size_t maskIndex = 0;

        for (size_t i = 0; i < numBuckets; ++i)
        {
            BucketId bucketId = buckets[i];

            size_t primCount{ 0 };
            uint8_t* destPtr = (uint8_t*)db.abi_context().iContext->getBucketArray(
                db.abi_context(), bucketId, attributeName, primCount);

            if (primCount == 0 || !destPtr)
            {
                db.logWarning("Bucket %zd has no entries for the given attribute", bucketId);
                continue;
            }

            if (mask.empty())
            {
                size_t totalBytes = elementSize * primCount;
                memcpy(destPtr, srcPtr, totalBytes);
                if (!srcIsScaler)
                    srcPtr += totalBytes;
            }
            else
            {
                for (size_t j = 0; j < primCount; ++j)
                {
                    if (mask[maskIndex++])
                        memcpy(destPtr, srcPtr, elementSize);
                    destPtr += elementSize;
                    if (!srcIsScaler)
                        srcPtr += elementSize;
                }
            }
        }

        // Set attribute for repeated paths
        if (numRepeatedPaths > 0)
        {
            // If there are repeated paths, the set behaviour might be unexpected. Warn the user.
            db.logWarning("Trying to set attribute when there are repeated prims in the gather. The first %zd values might be overwritten", totalPrimCount);
        }

        for (size_t i = 0; i < numRepeatedPaths; ++i)
        {
            PathBucketIndex pathBucketIndex = repeatedPaths[i];
            BucketId bucketId = std::get<1>(pathBucketIndex);
            ArrayIndex index = std::get<2>(pathBucketIndex);

            size_t primCount{ 0 };
            uint8_t* destPtr = (uint8_t*)db.abi_context().iContext->getBucketArray(
                db.abi_context(), bucketId, attributeName, primCount);

            if (primCount == 0 || !destPtr)
            {
                db.logWarning("Bucket %zd has no entries for the given attribute", bucketId);
                continue;
            }

            if (index >= primCount)
            {
                db.logWarning("Bucket %zd has less entries than required", bucketId);
                return false;
            }

            if (mask.empty() || mask[maskIndex++])
            {
                size_t byteCount = elementSize * index;
                memcpy(destPtr + byteCount, srcPtr, elementSize);
                if (!srcIsScaler)
                    srcPtr += elementSize;
            }
        }
        RETURN_TRUE_EXEC
    }
};

REGISTER_OGN_NODE()
}
}
}
