// 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.
//

//! @file Node.h
//!
//! @brief Defines @ref omni::graph::exec::unstable::Node.
#pragma once

#include <omni/core/ResultError.h>
#include <omni/graph/exec/unstable/Assert.h>
#include <omni/graph/exec/unstable/IGraph.h>
#include <omni/graph/exec/unstable/IGraphBuilderNode.h>
#include <omni/graph/exec/unstable/INode.h>
#include <omni/graph/exec/unstable/INodeDef.h>
#include <omni/graph/exec/unstable/INodeGraphDef.h>
#include <omni/graph/exec/unstable/ITopology.h>
#include <omni/graph/exec/unstable/SmallVector.h>
#include <omni/graph/exec/unstable/Types.h>

namespace omni
{
namespace graph
{
namespace exec
{
namespace unstable
{

//! @copydoc omni::graph::exec::unstable::INode
template <typename... Bases>
class NodeT : public Implements<Bases...>
{
public:
    //! Constructor of a node with an empty definition.
    //!
    //! May throw.
    static omni::core::ObjectPtr<NodeT> create(omni::core::ObjectParam<ITopology> topology, const char* idName)
    {
        OMNI_THROW_IF_ARG_NULL(topology);
        OMNI_THROW_IF_ARG_NULL(idName);

        return omni::core::steal(new NodeT(topology.get(), idName));
    }

    //! Constructor of a node with an empty definition.
    //!
    //! May throw.
    static omni::core::ObjectPtr<NodeT> create(omni::core::ObjectParam<IGraph> owner, const char* idName)
    {
        OMNI_THROW_IF_ARG_NULL(owner);
        OMNI_THROW_IF_ARG_NULL(idName);

        return omni::core::steal(new NodeT(owner->getTopology(), idName));
    }

    //! Constructor of a node with a node graph definition.
    //!
    //! May throw.
    static omni::core::ObjectPtr<NodeT> create(omni::core::ObjectParam<IGraph> owner,
                                               omni::core::ObjectParam<INodeGraphDef> nodeGraphDef,
                                               const char* idName)
    {
        OMNI_THROW_IF_ARG_NULL(owner);
        OMNI_THROW_IF_ARG_NULL(idName);

        return omni::core::steal(new NodeT(owner->getTopology(), nodeGraphDef.get(), idName));
    }

    //! Constructor of a node with a opaque node definition.
    //!
    //! May throw.
    static omni::core::ObjectPtr<NodeT> create(omni::core::ObjectParam<ITopology> owner,
                                               omni::core::ObjectParam<INodeGraphDef> nodeGraphDef,
                                               const char* idName)
    {
        OMNI_THROW_IF_ARG_NULL(owner);
        OMNI_THROW_IF_ARG_NULL(idName);

        return omni::core::steal(new NodeT(owner.get(), nodeGraphDef.get(), idName));
    }

    //! Constructor of a node with a opaque node definition.
    //!
    //! May throw.
    static omni::core::ObjectPtr<NodeT> create(omni::core::ObjectParam<IGraph> owner,
                                               omni::core::ObjectParam<INodeDef> nodeDef,
                                               const char* idName)
    {
        OMNI_THROW_IF_ARG_NULL(owner);
        OMNI_THROW_IF_ARG_NULL(idName);

        return omni::core::steal(new NodeT(owner->getTopology(), nodeDef.get(), idName));
    }

    //! Constructor of a node with a opaque node definition.
    //!
    //! May throw.
    static omni::core::ObjectPtr<NodeT> create(omni::core::ObjectParam<ITopology> topology,
                                               omni::core::ObjectParam<INodeDef> nodeDef,
                                               const char* idName)
    {
        OMNI_THROW_IF_ARG_NULL(topology);
        OMNI_THROW_IF_ARG_NULL(idName);

        return omni::core::steal(new NodeT(topology.get(), nodeDef.get(), idName));
    }

    //! Constructor of a node with a base definition (can be null_ptr, NodeDef or NodeGraphDef).
    //!
    //! May throw.
    static omni::core::ObjectPtr<NodeT> createForDef(omni::core::ObjectParam<ITopology> topology,
                                                     omni::core::ObjectParam<IDef> def,
                                                     const char* idName)
    {
        OMNI_THROW_IF_ARG_NULL(topology);
        OMNI_THROW_IF_ARG_NULL(idName);

        if (!def)
            return omni::core::steal(new NodeT(topology.get(), idName));
        else if (auto* nodeDef = omni::graph::exec::unstable::cast<INodeDef>(def))
            return omni::core::steal(new NodeT(topology.get(), nodeDef, idName));
        else if (auto* nodeGraphDef = omni::graph::exec::unstable::cast<INodeGraphDef>(def))
            return omni::core::steal(new NodeT(topology.get(), nodeGraphDef, idName));
        else
            return nullptr;
    }

    //! Destructor
    virtual ~NodeT()
    {
        // in case we decide to implement move constructor
        if (m_indexInTopology != kInvalidNodeIndexInTopology)
        {
            m_topology->releaseNodeIndex(m_indexInTopology);

            if (isValidTopology_abi())
            {
                m_topology->invalidate();
            }
        }
    }

    // disambiguate between INode and IGraphBuilderNode
    using INode::getChildren;
    using INode::getParents;
    using INode::getTopology;

protected:
    //! Core implementation of @ref omni::graph::exec::unstable::INode::getTopology_abi
    ITopology* getTopology_abi() noexcept override
    {
        return m_topology;
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INode::getName_abi
    const ConstName* getName_abi() noexcept override
    {
        return &m_name;
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INode::getIndexInTopology_abi
    NodeIndexInTopology getIndexInTopology_abi() noexcept override
    {
        return m_indexInTopology;
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INode::getParents_abi
    Span<INode* const> getParents_abi() noexcept override
    {
        return isValidTopology_abi() ? Span<INode* const>{ m_parents.begin(), m_parents.size() } :
                                       Span<INode* const>{ nullptr, 0 };
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INode::getChildren_abi
    Span<INode* const> getChildren_abi() noexcept override
    {
        return isValidTopology_abi() ? Span<INode* const>{ m_children.begin(), m_children.size() } :
                                       Span<INode* const>{ nullptr, 0 };
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INode::getCycleParentCount_abi
    uint32_t getCycleParentCount_abi() noexcept override
    {
        return isValidTopology_abi() ? m_cycleParentCount : 0;
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INode::isValidTopology_abi
    //!
    //! @note This method is called in the destructor and therefore must be marked as final
    bool isValidTopology_abi() noexcept final override
    {
        return m_topologyStamp.inSync(m_topology->getStamp());
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INode::validateOrResetTopology_abi
    virtual void validateOrResetTopology_abi() noexcept
    {
        if (m_topologyStamp.makeSync(m_topology->getStamp()))
        {
            // topology changed, let's clear the old one
            m_parents.clear();
            m_children.clear();
            m_cycleParentCount = 0;
        }
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INode::getDef_abi
    IDef* getDef_abi() noexcept override
    {
        if (m_nodeDef.get())
        {
            return m_nodeDef.get();
        }
        else
        {
            return m_nodeGraphDef.get();
        }
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INode::getNodeDef_abi
    INodeDef* getNodeDef_abi() noexcept override
    {
        return m_nodeDef.get();
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INode::getNodeGraphDef_abi
    INodeGraphDef* getNodeGraphDef_abi() noexcept override
    {
        return m_nodeGraphDef.get();
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::_addParent_abi
    omni::core::Result _addParent_abi(IGraphBuilderNode* parent) noexcept override
    {
        try
        {
            OMNI_GRAPH_EXEC_CAST_OR_RETURN(asNode, INode, parent);
            m_parents.push_back(asNode);
            return omni::core::kResultSuccess;
        }
        OMNI_GRAPH_EXEC_CATCH_ABI_EXCEPTION();
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::_removeParent_abi
    omni::core::Result _removeParent_abi(IGraphBuilderNode* parent) noexcept override
    {
        try
        {
            OMNI_GRAPH_EXEC_CAST_OR_RETURN(asNode, INode, parent);
            _eraseRemove(m_parents, asNode);
            return omni::core::kResultSuccess;
        }
        OMNI_GRAPH_EXEC_CATCH_ABI_EXCEPTION();
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::_addChild_abi
    omni::core::Result _addChild_abi(IGraphBuilderNode* child) noexcept override
    {
        try
        {
            OMNI_GRAPH_EXEC_CAST_OR_RETURN(asNode, INode, child);
            m_children.push_back(asNode);
            return omni::core::kResultSuccess;
        }
        OMNI_GRAPH_EXEC_CATCH_ABI_EXCEPTION();
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::_removeChild_abi
    omni::core::Result _removeChild_abi(IGraphBuilderNode* child) noexcept override
    {
        try
        {
            OMNI_GRAPH_EXEC_CAST_OR_RETURN(asNode, INode, child);
            _eraseRemove(m_children, asNode);
            return omni::core::kResultSuccess;
        }
        OMNI_GRAPH_EXEC_CATCH_ABI_EXCEPTION();
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::_removeInvalidParents_abi
    void _removeInvalidParents_abi() noexcept override
    {
        if (isValidTopology_abi())
        {
            m_parents.erase(
                std::remove_if(m_parents.begin(), m_parents.end(), [](INode* n) { return !n->isValidTopology(); }),
                m_parents.end());
        }
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::_removeInvalidChildren_abi
    void _removeInvalidChildren_abi() noexcept override
    {
        if (isValidTopology_abi())
        {
            m_children.erase(
                std::remove_if(m_children.begin(), m_children.end(), [](INode* n) { return !n->isValidTopology(); }),
                m_children.end());
        }
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::_invalidateConnections_abi
    //!
    //! @warning This only removes connections on a single node. The topology has bi-directional connections
    //! for every node with the exception of the connection with the root node.
    void _invalidateConnections_abi() noexcept override
    {
        m_topologyStamp.invalidate();
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::setCycleParentCount_abi
    void setCycleParentCount_abi(uint32_t count) noexcept override
    {
        m_cycleParentCount = count;
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::_setNodeDef_abi
    void _setNodeDef_abi(INodeDef* nodeDef) noexcept override
    {
        m_nodeDef.borrow(nodeDef);
        m_nodeGraphDef.release();
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::_setNodeGraphDef_abi
    void _setNodeGraphDef_abi(INodeGraphDef* nodeGraphDef) noexcept override
    {
        m_nodeGraphDef.borrow(nodeGraphDef);
        m_nodeDef.release();
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::_clearDef_abi
    void _clearDef_abi() noexcept override
    {
        m_nodeDef.release();
        m_nodeGraphDef.release();
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::getParentAt_abi
    omni::core::Result getParentAt_abi(uint64_t index, IGraphBuilderNode** out) noexcept override
    {
        *out = nullptr;

        if (!isValidTopology_abi() || index >= m_parents.size())
        {
            OMNI_GRAPH_EXEC_RETURN_ERROR(omni::core::kResultInvalidIndex);
        }
        else
        {
            OMNI_GRAPH_EXEC_CAST_OR_RETURN(
                asGraphBuilderNode, IGraphBuilderNode, m_parents[static_cast<uint32_t>(index)]);
            *out = asGraphBuilderNode; // explicitly does not acquire
            return omni::core::kResultSuccess;
        }
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::getParentCount_abi
    uint64_t getParentCount_abi() noexcept override
    {
        return isValidTopology_abi() ? m_parents.size() : 0;
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::getChildAt_abi
    omni::core::Result getChildAt_abi(uint64_t index, IGraphBuilderNode** out) noexcept override
    {
        *out = nullptr;
        if (!isValidTopology_abi() || index >= m_children.size())
        {
            OMNI_GRAPH_EXEC_RETURN_ERROR(omni::core::kResultInvalidIndex);
        }
        else
        {
            OMNI_GRAPH_EXEC_CAST_OR_RETURN(
                asGraphBuilderNode, IGraphBuilderNode, m_children[static_cast<uint32_t>(index)]);
            *out = asGraphBuilderNode; // explicitly does not acquire
            return omni::core::kResultSuccess;
        }
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::getChildCount_abi
    uint64_t getChildCount_abi() noexcept override
    {
        return isValidTopology_abi() ? m_children.size() : 0;
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::hasChild_abi
    bool hasChild_abi(IGraphBuilderNode* node) noexcept override
    {
        if (!isValidTopology_abi())
            return false;

        auto asNode = omni::graph::exec::unstable::cast<INode>(node);
        if (!asNode)
        {
            return false;
        }

        return std::find(m_children.begin(), m_children.end(), asNode) != m_children.end();
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::isRoot_abi
    bool isRoot_abi() noexcept override
    {
        return (m_topology->getRoot() == static_cast<INode*>(this));
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilderNode::getRoot_abi
    omni::core::Result getRoot_abi(IGraphBuilderNode** out) noexcept override
    {
        *out = nullptr;
        OMNI_GRAPH_EXEC_CAST_OR_RETURN(asGraphBuilderNode, IGraphBuilderNode, m_topology->getRoot());
        *out = asGraphBuilderNode; // explicitly does not acquire
        return omni::core::kResultSuccess;
    }

    //! Constructor
    NodeT(ITopology* topology, const char* idName) // may throw
        : m_topology{ topology }, m_indexInTopology{ m_topology->acquireNodeIndex() }, m_name{ idName }
    {
    }

    //! Constructor
    NodeT(ITopology* topology,
          INodeGraphDef* nodeGraphDef,
          const char* idName) // may throw
        : m_topology{ topology },
          m_indexInTopology{ m_topology->acquireNodeIndex() },
          m_nodeGraphDef{ nodeGraphDef, omni::core::kBorrow },
          m_name{ idName }
    {
    }

    //! Constructor
    NodeT(ITopology* topology,
          INodeDef* nodeDef,
          const char* idName) // may throw
        : m_topology{ topology },
          m_indexInTopology{ m_topology->acquireNodeIndex() },
          m_nodeDef{ nodeDef, omni::core::kBorrow },
          m_name{ idName }
    {
    }

private:
    //! Container for connections.
    //!
    //! Using @ref omni::graph::exec::unstable::SmallVector with local storage space for two nodes.
    //! The local space storage was hand picked, following the experience that most of graph nodes
    //! have very few downstream nodes.
    using NodeArray = SmallVector<INode*, 2>;

    //! Helper erase-remove idiom to remove and eliminate a node from the container
    template <typename T>
    void _eraseRemove(T& v, INode* n) // may throw
    {
        v.erase(std::remove(v.begin(), v.end(), n), v.end());
    };

    ITopology* m_topology; //!< Topology owning this node

    //! Acquired local index
    NodeIndexInTopology m_indexInTopology{ kInvalidNodeIndexInTopology };

    NodeArray m_parents; //!< Edges to parents
    NodeArray m_children; //!< Edges to children

    uint32_t m_cycleParentCount{ 0 }; //!< Cycling parents (used by the graph traversal)
    SyncStamp m_topologyStamp; //!< Validity check for edges

    omni::core::ObjectPtr<INodeDef> m_nodeDef; //!< Node definition
    omni::core::ObjectPtr<INodeGraphDef> m_nodeGraphDef; //!< Node graph definition

    ConstName m_name; //!< Identifier name
};

//! Core Node implementation for @ref omni::graph::exec::unstable::INode
using Node = NodeT<INode, IGraphBuilderNode>;

} // namespace unstable
} // namespace exec
} // namespace graph
} // namespace omni
