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

#include <unordered_map>

namespace omni {
namespace graph {
namespace test {

// Helper that returns a tokenized string containing the comma-separated attribute name and resolved type
// e.g. if the output "value" resolves to integer the string returned is "outputs:value,int".
// Unresolved types will have the data type "unknown"
template <typename OutputToResolve>
NameToken getResolutionDescription(OutputToResolve& attribute, NameToken attributeName, OgnTestTypeResolutionDatabase& db)
{
    std::string attributeResolution{db.tokenToString(attributeName)};
    attributeResolution += ",";
    attributeResolution += getOgnTypeName(attribute.type());

    return db.stringToToken(attributeResolution.c_str());
}

class OgnTestTypeResolution
{
public:
    static bool compute(OgnTestTypeResolutionDatabase& db)
    {
        // Set up the output array to report all of the available resolved types
        auto& resolvedType = db.outputs.resolvedType();
        resolvedType.resize(6);

        // Populate the output with the individual output attributes resolution information
        resolvedType[0] = getResolutionDescription(db.outputs.value(), outputs::value.token(), db);
        resolvedType[1] = getResolutionDescription(db.outputs.arrayValue(), outputs::arrayValue.token(), db);
        resolvedType[2] = getResolutionDescription(db.outputs.tupleValue(), outputs::tupleValue.token(), db);
        resolvedType[3] = getResolutionDescription(db.outputs.tupleArrayValue(), outputs::tupleArrayValue.token(), db);
        resolvedType[4] = getResolutionDescription(db.outputs.mixedValue(), outputs::mixedValue.token(), db);
        resolvedType[5] = getResolutionDescription(db.outputs.anyValue(), outputs::anyValue.token(), db);

        // Copy the input 'any' attribute value to output 'any'.
        if (db.inputs.anyValueIn().resolved())
        {
            db.outputs.anyValue().copyData(db.inputs.anyValueIn());
        }

        return true;
    }

    static void onConnectionTypeResolve(const NodeObj& node)
    {
        // anyValue and anyValueIn are coupled in this example - set one and the other auto-sets
        // the other attribs are going to behave as unrelated inputs, so disconnecting an attrib will
        // flip it back to any state

        auto iNode = node.iNode;

        size_t nAttribs = iNode->getAttributeCount(node);
        std::vector<AttributeObj> allAttribs(nAttribs);
        if (!iNode->getAttributes(node, &allAttribs[0], nAttribs))
            return;

        std::unordered_map<std::string, AttributeObj> attribsByName;

        const std::string anyValue = outputs::anyValue.name(), anyValueIn = inputs::anyValueIn.name();

        for (auto& attribObj : allAttribs)
        {
            const char* name = attribObj.iAttribute->getName(attribObj);
            attribsByName.emplace(name, attribObj);
        }

        auto iAttribute = allAttribs[0].iAttribute;

        auto isDisconnected = [&](const std::string& name) {
            return iAttribute->getDownstreamConnectionCount(attribsByName[name]) == 0 &&
                   iAttribute->getUpstreamConnectionCount(attribsByName[name]) == 0;
        };

        Type anyValueType = iAttribute->getResolvedType(attribsByName[anyValue]);
        Type anyValueInType = iAttribute->getResolvedType(attribsByName[anyValueIn]);

        // if one attribute is resolved and not the other we can resolve the other
        // but if one is resolved and the other not, and the resolved one has no connections - we can unresolve the
        // other way
        if (anyValueType != anyValueInType)
        {
            if (anyValueInType.baseType != BaseDataType::eUnknown)
            {
                if (isDisconnected(anyValueIn))
                {
                    iAttribute->setResolvedType(attribsByName[anyValueIn], Type());
                }
                else
                    iAttribute->setResolvedType(attribsByName[anyValue], anyValueInType);
            }
            else if (anyValueType.baseType != BaseDataType::eUnknown)
            {
                if (isDisconnected(anyValue))
                {
                    iAttribute->setResolvedType(attribsByName[anyValue], Type());
                }
                else
                    iAttribute->setResolvedType(attribsByName[anyValueIn], anyValueType);
            }
        }
        // If both are disconnected and both resolved, we can unresolve them both
        else if (isDisconnected(anyValue) && isDisconnected(anyValueIn))
        {
            iAttribute->setResolvedType(attribsByName[anyValue], Type());
            iAttribute->setResolvedType(attribsByName[anyValueIn], Type());
            return;
        }
    }
};

REGISTER_OGN_NODE()

} // namespace test
} // namespace graph
} // namespace omni
