// 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 <OgnGetGatheredAttributeDatabase.h>

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

using namespace carb::flatcache;

namespace omni
{
namespace graph
{
namespace core
{

class OgnGetGatheredAttribute
{
public:
    static bool compute(OgnGetGatheredAttributeDatabase& 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>();

        NameToken attributeName = db.inputs.name();

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

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

        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;
        }

        Type outputType(elementType);
        ++outputType.arrayDepth;

        // Determine if the output attribute has already been resolved to an incompatible type
        if (db.outputs.value().resolved())
        {
            Type outType = db.outputs.value().type();
            if (!outputType.compatibleRawData(outType))
            {
                db.logWarning("Resolved type %s of outputs:value is not compatible with type %s",
                    outType.getOgnTypeName().c_str(), db.outputs.value().type().getOgnTypeName().c_str());
                return false;
            }
        }

        // If it's resolved, we already know that it is compatible from the above check
        if (!db.outputs.value().resolved())
        {
            // Not resolved, so we have to resolve it now. This node is strange in that the resolved output type
            // depends on external state instead of other attributes.
            AttributeObj out = nodeObj.iNode->getAttributeByToken(nodeObj, outputs::value.m_token);
            out.iAttribute->setResolvedType(out, outputType);
            db.outputs.value().reset(db.abi_context(), out.iAttribute->getAttributeDataHandle(out), out);
        }

        AttributeObj outputAttr = nodeObj.iNode->getAttributeByToken(
            nodeObj, OgnGetGatheredAttributeAttributes::outputs::value.m_token);
        AttributeDataHandle outputHandle = outputAttr.iAttribute->getAttributeDataHandle(outputAttr);

        // determine the length of the output
        size_t totalPrimCount{ 0 };

        for (size_t i = 0; i < numBuckets; ++i)
        {
            BucketId bucketId = buckets[i];
            size_t primCount = 0;
            void* ptr = db.abi_context().iContext->getBucketArray(db.abi_context(), bucketId, attributeName, primCount);
            if (!ptr)
            {
                CARB_LOG_WARN("Attribute %s not found in Gather %zd, bucket %zd", attributeNameStr, gatherId, size_t(bucketId));
                return true;
            }
            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;
        }

        //printf("Getting %zd prims worth of %s\n", totalPrimCount, attributeNameStr);
        
        // Set the required length of output array
        db.abi_context().iAttributeData->setElementCount(db.abi_context(), outputHandle, totalPrimCount + numRepeatedPaths);

        // Get pointer to target data
        uint8_t* destPtr = nullptr;
        {
            void** out = nullptr;
            void** outPtr = reinterpret_cast<void**>(&out);
            db.abi_context().iAttributeData->getDataW(outPtr, db.abi_context(), &outputHandle, 1);
            destPtr = (uint8_t*)(*out);
        }
        CARB_ASSERT(destPtr);


        // Finally, we copy the data into the output, bucket by bucket
        for (size_t i = 0; i < numBuckets; ++i)
        {
            BucketId bucketId = buckets[i];

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

            if (primCount == 0 || !srcPtr)
            {
                db.logWarning("Bucket %zd has no entries for the given attribute", bucketId);
                return false;
            }
            size_t byteCount = elementSize * primCount;
            {
                // Copy the data
                memcpy(destPtr, srcPtr, byteCount);
                // Move the write pointer
                destPtr += byteCount;
            }
        }

        // Copy the data for repeated paths
        for (size_t i = 0; i < numRepeatedPaths; ++i)
        {
            BucketId bucketId = std::get<1>(repeatedPaths[i]);

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

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

            ArrayIndex index =  std::get<2>(repeatedPaths[i]);
            if (index >= primCount)
            {
                db.logWarning("Bucket %zd has less entries than required", bucketId);
                return false;
            }

            size_t byteCount = elementSize * index;
            {
                // Copy the data
                memcpy(destPtr, srcPtr + byteCount, elementSize);
                // Move the write pointer
                destPtr += elementSize;
            }
        }
        return true;
    }
};

REGISTER_OGN_NODE()
}
}
}
