// Copyright (c) 2022-2023, 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 <OgnArraySetIndexDatabase.h>
#include <omni/graph/core/StringUtils.h>
#include <omni/graph/core/ogn/ComputeHelpers.h>
#include <carb/logging/Log.h>
#include <algorithm>

namespace omni {
namespace graph {
namespace nodes {
using core::ogn::array;

// unnamed namespace to avoid multiple declaration when linking
namespace {

// helper to wrap a given index from legal range [-arrayLength, arrayLength) to [0:arrayLength)
// This will throw if the wrapped index is greater than arraySize when resizeToFit is false,
// or if the wrapped index is negative.
size_t tryWrapIndex(int index, size_t arraySize, bool resizeToFit)
{
    int wrappedIndex = index;
    if (index < 0)
        wrappedIndex = static_cast<int>(arraySize) + index;
    if ((wrappedIndex >= static_cast<int>(arraySize) && !resizeToFit) || wrappedIndex < 0)
        throw ogn::compute::InputError(formatString("inputs:index %d is out of range for inputs:array of size %zu", wrappedIndex, arraySize));
    return static_cast<size_t>(wrappedIndex);
}

template<typename BaseType>
bool tryComputeAssumingType(OgnArraySetIndexDatabase& db)
{
    auto inputArray = db.inputs.array().template get<BaseType[]>();
    size_t const inputArraySize = db.inputs.array().size();
    size_t const index = tryWrapIndex(db.inputs.index(), inputArraySize, db.inputs.resizeToFit());
    auto const value = db.inputs.value().template get<BaseType>();
    auto outputArray = db.outputs.array().template get<BaseType[]>();

    if (!value || !inputArray)
        return false;

    // tryWrapIndex would have thrown already if index >= inputArraySize and resizeToFit is false
    size_t outputArraySize = std::max(inputArraySize, index + 1);
    (*outputArray).resize(outputArraySize);
    memcpy(outputArray->data(), inputArray->data(), sizeof(BaseType) * inputArraySize);
    if (outputArraySize > inputArraySize)
        memset(&((*outputArray)[inputArraySize]), 0, sizeof(BaseType) * (outputArraySize - inputArraySize));
    memcpy(&((*outputArray)[index]), &*value, sizeof(BaseType));

    return true;
}
} // namespace

class OgnArraySetIndex
{
public:
    static bool compute(OgnArraySetIndexDatabase& db)
    {
        auto& inputType = db.inputs.value().type();
        try
        {
            switch (inputType.baseType)
            {
            case BaseDataType::eBool:
                return tryComputeAssumingType<bool>(db);
            case BaseDataType::eToken:
                return tryComputeAssumingType<ogn::Token>(db);
            case BaseDataType::eDouble:
                switch (inputType.componentCount)
                {
                    case 1:  return tryComputeAssumingType<double>(db);
                    case 2:  return tryComputeAssumingType<double[2]>(db);
                    case 3:  return tryComputeAssumingType<double[3]>(db);
                    case 4:  return tryComputeAssumingType<double[4]>(db);
                    case 9:  return tryComputeAssumingType<double[9]>(db);
                    case 16: return tryComputeAssumingType<double[16]>(db);
                    default: break;
                }
            case BaseDataType::eFloat:
                switch (inputType.componentCount)
                {
                    case 1: return tryComputeAssumingType<float>(db);
                    case 2: return tryComputeAssumingType<float[2]>(db);
                    case 3: return tryComputeAssumingType<float[3]>(db);
                    case 4: return tryComputeAssumingType<float[4]>(db);
                    default: break;
                }
            case BaseDataType::eHalf:
                switch (inputType.componentCount)
                {
                    case 1: return tryComputeAssumingType<pxr::GfHalf>(db);
                    case 2: return tryComputeAssumingType<pxr::GfHalf[2]>(db);
                    case 3: return tryComputeAssumingType<pxr::GfHalf[3]>(db);
                    case 4: return tryComputeAssumingType<pxr::GfHalf[4]>(db);
                    default: break;
                }
            case BaseDataType::eInt:
                switch (inputType.componentCount)
                {
                    case 1: return tryComputeAssumingType<int32_t>(db);
                    case 2: return tryComputeAssumingType<int32_t[2]>(db);
                    case 3: return tryComputeAssumingType<int32_t[3]>(db);
                    case 4: return tryComputeAssumingType<int32_t[4]>(db);
                    default: break;
                }
                ;
            case BaseDataType::eInt64:
                return tryComputeAssumingType<int64_t>(db);
            case BaseDataType::eUChar:
                return tryComputeAssumingType<unsigned char>(db);
            case BaseDataType::eUInt:
                return tryComputeAssumingType<uint32_t>(db);
            case BaseDataType::eUInt64:
                return tryComputeAssumingType<uint64_t>(db);
            default: break;
            }

            db.logWarning("Failed to resolve input types");
        }
        catch (ogn::compute::InputError const &error)
        {
            db.logError(error.what());
        }
        return false;
    }

    static void onConnectionTypeResolve(const NodeObj& node){
        auto const inputArray = node.iNode->getAttributeByToken(node, inputs::array.token());
        auto const inputValue = node.iNode->getAttributeByToken(node, inputs::value.token());
        auto const outputArray = node.iNode->getAttributeByToken(node, outputs::array.token());
        auto const inputArrayType = inputArray.iAttribute->getResolvedType(inputArray);

        if (inputArrayType.baseType != BaseDataType::eUnknown)
        {
            std::array<AttributeObj, 3> attrs { inputArray, inputValue, outputArray };
            // all should have the same tuple count
            std::array<uint8_t, 3> tupleCounts {
                inputArrayType.componentCount,
                inputArrayType.componentCount,
                inputArrayType.componentCount
            };
            // value type can not be an array because we don't support arrays-of-arrays
            std::array<uint8_t, 3> arrayDepths {
                1,
                0,
                1
            };
            std::array<AttributeRole, 3> rolesBuf {
                inputArrayType.role,
                AttributeRole::eUnknown,
                AttributeRole::eUnknown
            };
            node.iNode->resolvePartiallyCoupledAttributes(node, attrs.data(), tupleCounts.data(),
                                                        arrayDepths.data(), rolesBuf.data(), attrs.size());
        }
    }
};

REGISTER_OGN_NODE()

} // namespace nodes
} // namespace graph
} // namespace omni
