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

#include <array>
#include <unordered_map>
#include <vector>

namespace omni {
namespace graph {
namespace nodes {
// unnamed namespace to avoid multiple declaration when linking
namespace {
static constexpr size_t kMaxAttrNameLen{ 32 };

void formatAttrName(size_t n, std::array<char, kMaxAttrNameLen>& outName)
{
    snprintf(outName.data(), kMaxAttrNameLen, "inputs:input%zu", n);
}

template<typename BaseType>
bool tryComputeAssumingType(OgnConstructArrayDatabase& db)
{
    NodeObj nodeObj = db.abi_node();
    auto iNode = nodeObj.iNode;
    GraphObj graphObj = iNode->getGraph(nodeObj);
    GraphContextObj context = graphObj.iGraph->getDefaultGraphContext(graphObj);

    auto const outputArrayAttr = iNode->getAttributeByToken(nodeObj, outputs::array.token());
    auto const outputArrayType = outputArrayAttr.iAttribute->getResolvedType(outputArrayAttr);
    if (outputArrayType.baseType == BaseDataType::eUnknown)
    {
        return false;
    }

    auto outputArray = db.outputs.array().template get<BaseType[]>();
    if (!outputArray)
        return false;
    size_t const arraySize = static_cast<size_t>(std::max(0, db.inputs.arraySize()));
    (*outputArray).resize(arraySize);
    memset(outputArray->data(), 0, sizeof(BaseType) * arraySize);

    // read dynamic inputs
    size_t i = 0;
    std::array<char, kMaxAttrNameLen> outName;
    formatAttrName(i, outName);
    while (iNode->getAttributeExists(nodeObj, outName.data()) && i < arraySize) {
        auto inAttrib = iNode->getAttribute(nodeObj, outName.data());
        auto inputAttribType = inAttrib.iAttribute->getResolvedType(inAttrib);
        if (inputAttribType.baseType != BaseDataType::eUnknown)
        {
            ConstAttributeDataHandle handle = inAttrib.iAttribute->getConstAttributeDataHandle(inAttrib, db.getInstanceIndex());
            auto const dataPtr = getDataR<BaseType>(context, handle);
            if (dataPtr)
                (*outputArray)[i] = *dataPtr;
        }
        ++i;
        formatAttrName(i, outName);
    }
        
    // fill the rest of the array using the last input value if necessary
    if (0 < i && i < arraySize)
    {
        for (size_t j = i; j < arraySize; ++j)
        {
            (*outputArray)[j] = (*outputArray)[i - 1];
        }
    }

    return true;
}

template<typename BaseType, size_t TupleSize>
bool tryComputeAssumingType(OgnConstructArrayDatabase& db)
{
    NodeObj nodeObj = db.abi_node();
    auto iNode = nodeObj.iNode;
    GraphObj graphObj = iNode->getGraph(nodeObj);
    GraphContextObj context = graphObj.iGraph->getDefaultGraphContext(graphObj);

    auto const outputArrayAttr = iNode->getAttributeByToken(nodeObj, outputs::array.token());
    auto const outputArrayType = outputArrayAttr.iAttribute->getResolvedType(outputArrayAttr);
    if (outputArrayType.baseType == BaseDataType::eUnknown)
    {
        return false;
    }

    auto outputArray = db.outputs.array().template get<BaseType[][TupleSize]>();
    if (!outputArray)
        return false;
    size_t const arraySize = static_cast<size_t>(std::max(0, db.inputs.arraySize()));
    (*outputArray).resize(arraySize);
    memset(outputArray->data(), 0, sizeof(BaseType) * arraySize * TupleSize);

    // read dynamic inputs
    size_t i = 0;
    std::array<char, kMaxAttrNameLen> outName;
    formatAttrName(i, outName);
    while (iNode->getAttributeExists(nodeObj, outName.data()) && i < arraySize) {
        auto inAttrib = iNode->getAttribute(nodeObj, outName.data());
        auto inputAttribType = inAttrib.iAttribute->getResolvedType(inAttrib);
        if (inputAttribType.baseType != BaseDataType::eUnknown)
        {
            ConstAttributeDataHandle handle = inAttrib.iAttribute->getConstAttributeDataHandle(inAttrib, db.getInstanceIndex());
            auto const dataPtr = getDataR<BaseType[TupleSize]>(context, handle);
            if (dataPtr)
                memcpy(&((*outputArray)[i]), dataPtr, sizeof(BaseType) * TupleSize);
        }
        ++i;
        formatAttrName(i, outName);
    }
        
    // fill the rest of the array using the last input value if necessary
    if (0 < i && i < arraySize)
    {
        for (size_t j = i; j < arraySize; ++j)
        {
            memcpy(&((*outputArray)[j]), &((*outputArray)[i - 1]), sizeof(BaseType) * TupleSize);
        }
    }

    return true;
}
} // namespace

class OgnConstructArray
{
    static void tryResolveAttribute(omni::graph::core::NodeObj const& nodeObj,
                                    omni::graph::core::AttributeObj attrib,
                                    omni::graph::core::Type attribType)
    {
        Type valueType{ attrib.iAttribute->getResolvedType(attrib) };
        bool attribIsValid = (attribType.baseType != BaseDataType::eUnknown);

        if (valueType.baseType != BaseDataType::eUnknown)
        {
            // Resolved
            // Case 1: We didn't find a valid source attribute => unresolve
            // Case 2: We found an attribute but it is not compatible with our current resolution => unresolve
            // Else: All good
            if (!attribIsValid or (attribType != valueType))
            {
                // Not compatible! Request that the attribute be un-resolved. Note that this could fail if there are
                // connections that result in a contradiction during type propagation
                attrib.iAttribute->setResolvedType(attrib, Type(BaseDataType::eUnknown));
            }
        }

        // If it's unresolved (and we have a valid attribute) we can request a resolution
        if (attribIsValid and (attrib.iAttribute->getResolvedType(attrib).baseType == BaseDataType::eUnknown))
        {
            attrib.iAttribute->setResolvedType(attrib, attribType);
        }
    }

    static void tryResolveArrayAttributes(omni::graph::core::NodeObj const& nodeObj, bool reconnectInputs)
    {
        auto& state = OgnConstructArrayDatabase::sInternalState<OgnConstructArray>(nodeObj);

        // Get the input attributes
        std::vector<AttributeObj> inputAttributes;
        size_t i = 0;
        std::array<char, kMaxAttrNameLen> outName;
        formatAttrName(i, outName);
        while (nodeObj.iNode->getAttributeExists(nodeObj, outName.data()))
        {
            inputAttributes.push_back(nodeObj.iNode->getAttribute(nodeObj, outName.data()));
            ++i;
            formatAttrName(i, outName);
        }

        // Determine the output array type from the specified array type and connected inputs
        Type outputType = state.m_inputArrayType;  // Initialize to the specified array type ('eUnknown' if unspecified, i.e. 'inputs:arrayType' is set to 'auto')
        for (auto const& inAttrib : inputAttributes)
        {
            // Check if input is an array
            if (inAttrib.iAttribute->getResolvedType(inAttrib).arrayDepth != 0)
            {
                OgnConstructArrayDatabase::logError(nodeObj, "Nested arrays not supported: The input array element at attribute '%s' is an array",
                    inAttrib.iAttribute->getName(inAttrib));

                // Unresolve output and return
                auto outputArrayAttrib = nodeObj.iNode->getAttributeByToken(nodeObj, outputs::array.token());
                tryResolveAttribute(nodeObj, outputArrayAttrib, Type(BaseDataType::eUnknown));
                return;
            }

            // Skip unconnected inputs
            size_t upstreamConnectionCount = inAttrib.iAttribute->getUpstreamConnectionCount(inAttrib);
            if (upstreamConnectionCount != 1)
                continue;

            // Get the type of the current input from the upstream connection
            ConnectionInfo upstreamConnection;
            inAttrib.iAttribute->getUpstreamConnectionsInfo(inAttrib, &upstreamConnection, 1);
            Type const upstreamType = upstreamConnection.attrObj.iAttribute->getResolvedType(upstreamConnection.attrObj);

            // The array type is not specified, so infer from the first connected and resolved input
            if (upstreamType.baseType != BaseDataType::eUnknown)
            {
                if (outputType.baseType == BaseDataType::eUnknown)
                {
                    outputType = upstreamType;
                }
                else
                {
                    // Check if the specified or inferred array type matches the type of the current input
                    if (!outputType.compatibleRawData(upstreamType))
                    {
                        OgnConstructArrayDatabase::logError(
                            nodeObj, "Mismatched array element type on input attribute '%s': expected '%s', got '%s'",
                            inAttrib.iAttribute->getName(inAttrib), getOgnTypeName(outputType).c_str(),
                            getOgnTypeName(upstreamType).c_str());
                        outputType = Type(BaseDataType::eUnknown);
                        break;
                    }
                }
            }
        }

        // Resolve inputs
        for (auto& inAttrib : inputAttributes)
        {
            size_t upstreamConnectionCount = inAttrib.iAttribute->getUpstreamConnectionCount(inAttrib);
            if (upstreamConnectionCount == 0)
            {
                // Resolve unconnected inputs to the specified array type, or the inferred array type if the array type is unspecified
                if (state.m_inputArrayType.baseType != BaseDataType::eUnknown)
                {
                    if (inAttrib.iAttribute->getResolvedType(inAttrib) != state.m_inputArrayType)
                    {
                        // Case: The current input is disconnected and the array type is specified
                        // Action: Resolve the current input to the specified array type
                        tryResolveAttribute(nodeObj, inAttrib, state.m_inputArrayType);
                    }
                }
                else
                {
                    if (inAttrib.iAttribute->getResolvedType(inAttrib) != outputType)
                    {
                        // Case: The current input is disconnected and the array type is unspecified (auto)
                        // Action: Resolve the current input to the inferred array type if any input is connected.
                        //         If no inputs are connected, then outputType will be 'eUnknown', unresolving the current input.
                        tryResolveAttribute(nodeObj, inAttrib, outputType);
                    }
                }
            }
            else if (upstreamConnectionCount == 1 && reconnectInputs)  // reconnectInputs avoids an infinite loop
            {
                // Resolve connected inputs to the specified array type, or their upstream types if the array type is unspecified
                if (state.m_inputArrayType.baseType != BaseDataType::eUnknown)
                {
                    // Check if already resolved to the specified array type
                    if (inAttrib.iAttribute->getResolvedType(inAttrib) != state.m_inputArrayType)
                    {
                        // Disconnect before re-resolving
                        ConnectionInfo upstreamConnection;
                        inAttrib.iAttribute->getUpstreamConnectionsInfo(inAttrib, &upstreamConnection, 1);
                        inAttrib.iAttribute->disconnectAttrs(upstreamConnection.attrObj, inAttrib, true);

                        // Case: The current input is connected and the array type is specified
                        // Action: Resolve the current input to the specified array type
                        tryResolveAttribute(nodeObj, inAttrib, state.m_inputArrayType);

                        // Reconnect
                        ConnectionInfo destConnection {inAttrib, upstreamConnection.connectionType};
                        inAttrib.iAttribute->connectAttrsEx(upstreamConnection.attrObj, destConnection, true);
                    }
                }
                else
                {
                    ConnectionInfo upstreamConnection;
                    inAttrib.iAttribute->getUpstreamConnectionsInfo(inAttrib, &upstreamConnection, 1);
                    Type const upstreamType = upstreamConnection.attrObj.iAttribute->getResolvedType(upstreamConnection.attrObj);
                    // Check if already resolved to the upstream type
                    if (inAttrib.iAttribute->getResolvedType(inAttrib) != upstreamType)
                    {
                        // Disconnect before re-resolving
                        inAttrib.iAttribute->disconnectAttrs(upstreamConnection.attrObj, inAttrib, true);

                        // Case: The current input is connected and the array type is unspecified (auto)
                        // Action: Resolve the current input to its upstream type (not the inferred array type 'outputType',
                        //         which could be a different but compatible type or 'eUnknown')
                        tryResolveAttribute(nodeObj, inAttrib, Type(BaseDataType::eUnknown));  // Must be resolved to 'eUnknown', not 'upstreamType'

                        // Reconnect
                        ConnectionInfo destConnection {inAttrib, upstreamConnection.connectionType};
                        inAttrib.iAttribute->connectAttrsEx(upstreamConnection.attrObj, destConnection, true);
                    }
                }
            }
        }

        // Resolve the output
        // If there is a type mismatch (regardless of whether the array type was specified), then 'outputType' gets set to 'eUnknown', so 'outputs:array' gets unresolved.
        // If the array type is unspecified (auto) and no inputs are connected (i.e. no type is inferred), then 'outputType' gets set to 'eUnknown' as well (but the node does not error).
        // Otherwise, 'outputs:array' gets resolved to the specified or inferred array type 'outputType'.
        outputType.arrayDepth = 1;
        auto outputArrayAttrib = nodeObj.iNode->getAttributeByToken(nodeObj, outputs::array.token());
        tryResolveAttribute(nodeObj, outputArrayAttrib, outputType);
    }

    static void onValueChanged(AttributeObj const& attrObj, void const* userData)
    {
        NodeObj nodeObj{ attrObj.iAttribute->getNode(attrObj) };
        if (nodeObj.nodeHandle == kInvalidNodeHandle)
            return;

        GraphObj graphObj{ nodeObj.iNode->getGraph(nodeObj) };
        if (graphObj.graphHandle == kInvalidGraphHandle)
            return;

        GraphContextObj context{ graphObj.iGraph->getDefaultGraphContext(graphObj) };
        if (context.contextHandle == kInvalidGraphContextHandle)
            return;

        ConstAttributeDataHandle handle =
            attrObj.iAttribute->getConstAttributeDataHandle(attrObj, kAccordingToContextIndex);
        auto const dataPtr = getDataR<NameToken>(context, handle);
        if (dataPtr)
        {
            auto& state = OgnConstructArrayDatabase::sInternalState<OgnConstructArray>(nodeObj);
            state.m_inputArrayType = state.m_arrayTypes[*dataPtr];
        }

        tryResolveArrayAttributes(nodeObj, true);
    }

    static void onConnected(AttributeObj const& otherAttrib, AttributeObj const& inAttrib, void* userData)
    {
        NodeObj nodeObj{ inAttrib.iAttribute->getNode(inAttrib) };
        if (nodeObj.nodeHandle == kInvalidNodeHandle)
            return;

        NodeObj* thisNodeObj = reinterpret_cast<NodeObj*>(userData);
        if (nodeObj.nodeHandle != thisNodeObj->nodeHandle)
            return;

        auto& state = OgnConstructArrayDatabase::sInternalState<OgnConstructArray>(nodeObj);

        // Deregister onConnected callback to avoid infinite loop
        struct ConnectionCallback connectedCallback = {onConnected, &state.m_nodeObj};
        nodeObj.iNode->deregisterConnectedCallback(nodeObj, connectedCallback);

        tryResolveArrayAttributes(nodeObj, true);

        // Re-register onConnected callback
        nodeObj.iNode->registerConnectedCallback(nodeObj, connectedCallback);
    }

public:
    omni::graph::core::NodeObj m_nodeObj;

    std::unordered_map<NameToken, omni::graph::core::Type> m_arrayTypes;

    omni::graph::core::Type m_inputArrayType;

    static void initialize(GraphContextObj const& context, NodeObj const& nodeObj)
    {
        auto& state = OgnConstructArrayDatabase::sInternalState<OgnConstructArray>(nodeObj);

        state.m_nodeObj = nodeObj;

        state.m_arrayTypes = {
            {OgnConstructArrayDatabase::tokens.Bool, Type(BaseDataType::eBool)},
            {OgnConstructArrayDatabase::tokens.Double, Type(BaseDataType::eDouble)},
            {OgnConstructArrayDatabase::tokens.Float, Type(BaseDataType::eFloat)},
            {OgnConstructArrayDatabase::tokens.Half, Type(BaseDataType::eHalf)},
            {OgnConstructArrayDatabase::tokens.Int, Type(BaseDataType::eInt)},
            {OgnConstructArrayDatabase::tokens.Int64, Type(BaseDataType::eInt64)},
            {OgnConstructArrayDatabase::tokens.Token, Type(BaseDataType::eToken)},
            {OgnConstructArrayDatabase::tokens.UChar, Type(BaseDataType::eUChar)},
            {OgnConstructArrayDatabase::tokens.UInt, Type(BaseDataType::eUInt)},
            {OgnConstructArrayDatabase::tokens.UInt64, Type(BaseDataType::eUInt64)},
            {OgnConstructArrayDatabase::tokens.Double_2, Type(BaseDataType::eDouble, 2)},
            {OgnConstructArrayDatabase::tokens.Double_3, Type(BaseDataType::eDouble, 3)},
            {OgnConstructArrayDatabase::tokens.Double_4, Type(BaseDataType::eDouble, 4)},
            {OgnConstructArrayDatabase::tokens.Double_9, Type(BaseDataType::eDouble, 9, 0, AttributeRole::eMatrix)},
            {OgnConstructArrayDatabase::tokens.Double_16, Type(BaseDataType::eDouble, 16, 0, AttributeRole::eMatrix)},
            {OgnConstructArrayDatabase::tokens.Float_2, Type(BaseDataType::eFloat, 2)},
            {OgnConstructArrayDatabase::tokens.Float_3, Type(BaseDataType::eFloat, 3)},
            {OgnConstructArrayDatabase::tokens.Float_4, Type(BaseDataType::eFloat, 4)},
            {OgnConstructArrayDatabase::tokens.Half_2, Type(BaseDataType::eHalf, 2)},
            {OgnConstructArrayDatabase::tokens.Half_3, Type(BaseDataType::eHalf, 3)},
            {OgnConstructArrayDatabase::tokens.Half_4, Type(BaseDataType::eHalf, 4)},
            {OgnConstructArrayDatabase::tokens.Int_2, Type(BaseDataType::eInt, 2)},
            {OgnConstructArrayDatabase::tokens.Int_3, Type(BaseDataType::eInt, 3)},
            {OgnConstructArrayDatabase::tokens.Int_4, Type(BaseDataType::eInt, 4)},
        };

        AttributeObj arrayTypeAttrib = nodeObj.iNode->getAttributeByToken(nodeObj, inputs::arrayType.m_token);
        arrayTypeAttrib.iAttribute->registerValueChangedCallback(arrayTypeAttrib, onValueChanged, true);

        struct ConnectionCallback connectedCallback = {onConnected, &state.m_nodeObj};
        nodeObj.iNode->registerConnectedCallback(nodeObj, connectedCallback);

        onValueChanged(arrayTypeAttrib, nullptr);
    }

    static bool compute(OgnConstructArrayDatabase& db)
    {
        NodeObj nodeObj = db.abi_node();
        auto const outputArrayAttr = nodeObj.iNode->getAttributeByToken(nodeObj, outputs::array.token());
        auto const outputArrayType = outputArrayAttr.iAttribute->getResolvedType(outputArrayAttr);

        try
        {
            switch (outputArrayType.baseType)
            {
            case BaseDataType::eBool:
                if (outputArrayType.componentCount == 1)
                    if (tryComputeAssumingType<bool>(db)) return true;
                break;
            case BaseDataType::eDouble:
                switch (outputArrayType.componentCount)
                {
                case 1:
                    if (tryComputeAssumingType<double>(db)) return true;
                    break;
                case 2:
                    if (tryComputeAssumingType<double, 2>(db)) return true;
                    break;
                case 3:
                    if (tryComputeAssumingType<double, 3>(db)) return true;
                    break;
                case 4:
                    if (tryComputeAssumingType<double, 4>(db)) return true;
                    break;
                case 9:
                    if (tryComputeAssumingType<double, 9>(db)) return true;
                    break;
                case 16:
                    if (tryComputeAssumingType<double, 16>(db)) return true;
                    break;
                }
                break;
            case BaseDataType::eFloat:
                switch (outputArrayType.componentCount)
                {
                case 1:
                    if (tryComputeAssumingType<float>(db)) return true;
                    break;
                case 2:
                    if (tryComputeAssumingType<float, 2>(db)) return true;
                    break;
                case 3:
                    if (tryComputeAssumingType<float, 3>(db)) return true;
                    break;
                case 4:
                    if (tryComputeAssumingType<float, 4>(db)) return true;
                    break;
                }
                break;
            case BaseDataType::eHalf:
                switch (outputArrayType.componentCount)
                {
                case 1:
                    if (tryComputeAssumingType<pxr::GfHalf>(db)) return true;
                    break;
                case 2:
                    if (tryComputeAssumingType<pxr::GfHalf, 2>(db)) return true;
                    break;
                case 3:
                    if (tryComputeAssumingType<pxr::GfHalf, 3>(db)) return true;
                    break;
                case 4:
                    if (tryComputeAssumingType<pxr::GfHalf, 4>(db)) return true;
                    break;
                }
                break;
            case BaseDataType::eInt:
                switch (outputArrayType.componentCount)
                {
                case 1:
                    if (tryComputeAssumingType<int32_t>(db)) return true;
                    break;
                case 2:
                    if (tryComputeAssumingType<int32_t, 2>(db)) return true;
                    break;
                case 3:
                    if (tryComputeAssumingType<int32_t, 3>(db)) return true;
                    break;
                case 4:
                    if (tryComputeAssumingType<int32_t, 4>(db)) return true;
                    break;
                }
                break;
            case BaseDataType::eInt64:
                if (outputArrayType.componentCount == 1)
                    if (tryComputeAssumingType<int64_t>(db)) return true;
                break;
            case BaseDataType::eToken:
                if (outputArrayType.componentCount == 1)
                    if (tryComputeAssumingType<OgnToken>(db)) return true;
                break;
            case BaseDataType::eUChar:
                if (outputArrayType.componentCount == 1)
                    if (tryComputeAssumingType<uint8_t>(db)) return true;
                break;
            case BaseDataType::eUInt:
                if (outputArrayType.componentCount == 1)
                    if (tryComputeAssumingType<uint32_t>(db)) return true;
                break;
            case BaseDataType::eUInt64:
                if (outputArrayType.componentCount == 1)
                    if (tryComputeAssumingType<uint64_t>(db)) return true;
                break;
            default:
                break;
            }

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

        return false;
    }

    static void onConnectionTypeResolve(const NodeObj& nodeObj)
    {
        tryResolveArrayAttributes(nodeObj, false);
    }
};

REGISTER_OGN_NODE()

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