// Copyright (c) 2020-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.
//
#include <OgnConcatenateFloat3ArraysDatabase.h>
#include <carb/Framework.h>
#include <carb/Types.h>

#include <omni/graph/core/Accessors.h>
#include <omni/graph/core/IGatherPrototype.h>

#include <omni/math/linalg/vec.h>

using omni::math::linalg::vec3f;

namespace omni
{
namespace graph
{
namespace nodes
{

class OgnConcatenateFloat3Arrays
{
public:
    static bool compute(OgnConcatenateFloat3ArraysDatabase& db)
    {
        auto iGatherPrototype = carb::getCachedInterface<IGatherPrototype>();
        if (!iGatherPrototype)
        {
            db.logError("Could not acquire IGatherPrototype and therefore could not compute");
            return false;
        }

        auto& outputArray = db.outputs.outputArray();
        auto& outputSizesArray = db.outputs.arraySizes();

        // TODO: This uses the legacy approach instead of the database until gather arrays are supported
        AttributeObj inputArrayObj = db.abi_node().iNode->getAttribute(db.abi_node(), "inputs:inputArrays");
        AttributeObj inputArraySrcObj = getNthUpstreamAttributeOrSelf<0>(inputArrayObj);
        auto inputArrayValues = reinterpret_cast<const vec3f* const* const>(
            iGatherPrototype->getGatherArray(db.abi_context(), inputArraySrcObj, kReadOnly)
        );
        if (!inputArrayValues)
        {
            return false;
        }
        auto inputArraySizes = iGatherPrototype->getGatherArrayAttributeSizes(db.abi_context(), inputArraySrcObj, kReadOnly);
        const size_t inputArrayCount = iGatherPrototype->getElementCount(db.abi_context(), inputArraySrcObj);

        // Find the total output size.
        size_t outputCount = 0;
        for (size_t arrayi = 0; arrayi < inputArrayCount; ++arrayi)
        {
            outputCount += inputArraySizes[arrayi];
        }

        outputArray.resize(outputCount);
        outputSizesArray.resize(inputArrayCount);

        for (size_t arrayi = 0, outputi = 0; arrayi < inputArrayCount; ++arrayi)
        {
            const vec3f* const currentInputArray = inputArrayValues[arrayi];
            size_t currentSize = inputArraySizes[arrayi];
            outputSizesArray[arrayi] = int(currentSize);
            for (size_t inputi = 0; inputi < currentSize; ++inputi, ++outputi)
            {
                outputArray[outputi] = currentInputArray[inputi];
            }
        }

        return true;
    }
};

REGISTER_OGN_NODE()

}
}
}
