// Copyright (c) 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 <OgnWritePrimRelationshipDatabase.h>
#include "PrimCommon.h"

namespace omni {
namespace graph {
namespace nodes {

class OgnWritePrimRelationship {

public:
    // ----------------------------------------------------------------------------
    // Prefetch prim to fabric
    static void setup(NodeObj const& nodeObj, GraphObj const& graphObj, OgnWritePrimRelationshipDatabase& db, size_t offset)
    {
        db.state.correctlySetup(offset) = false;

        TargetPath prim = db.inputs.prim.firstOrDefault(offset);
        NameToken name = db.inputs.name(offset);
        std::set<omni::fabric::TokenC> attributes = {name.token};
        prefetchPrim(graphObj, nodeObj, prim, attributes);

        db.state.correctlySetup(offset) = true;
        db.state.prim(offset).resize(1);
        db.state.prim(offset)[0] = prim;
        db.state.name(offset) = name;
    }

    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 contextObj = graphObj.iGraph->getDefaultGraphContext(graphObj);

        OgnWritePrimRelationshipDatabase db(contextObj, nodeObj);
        setup(nodeObj, graphObj, db, 0);
    }

    // ----------------------------------------------------------------------------
    static void initialize(GraphContextObj const& context, NodeObj const& nodeObj)
    {
        std::array<NameToken, 2> attribNames{inputs::name.m_token, inputs::prim.m_token};
        for (auto const& attribName : attribNames)
        {
            AttributeObj attribObj = nodeObj.iNode->getAttributeByToken(nodeObj, attribName);
            attribObj.iAttribute->registerValueChangedCallback(attribObj, onValueChanged, true);
        }
    }

    // ----------------------------------------------------------------------------
    static size_t computeVectorized(OgnWritePrimRelationshipDatabase& db, size_t count)
    {
        NodeObj nodeObj = db.abi_node();
        GraphContextObj context = db.abi_context();
        GraphObj graphObj = nodeObj.iNode->getGraph(nodeObj);

        auto name = db.inputs.name.vectorized(count);
        auto usdWriteBack = db.inputs.usdWriteBack.vectorized(count);
        auto correctlySetup = db.state.correctlySetup.vectorized(count);
        auto nameState = db.state.name.vectorized(count);
        auto execOut = db.outputs.execOut.vectorized(count);

        for (size_t idx = 0; idx < count; ++idx)
        {
            TargetPath curPrim = db.inputs.prim.firstOrDefault(idx);
            if(!correctlySetup[idx] ||
               db.state.prim(idx).size() == 0 ||
               db.state.prim(idx)[0] != curPrim ||
               nameState[idx] != name[idx])
            {
                setup(nodeObj, graphObj, db, idx);
            }

            if(correctlySetup[idx])
            {
                setRelationshipTargets(context, nodeObj, curPrim, name[idx], db.inputs.value(idx), usdWriteBack[idx]);
                execOut[idx] = kExecutionAttributeStateEnabled;
            }
        }
        return count;
    }

};

REGISTER_OGN_NODE()

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

