// 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 GraphBuilder.h
//!
//! @brief Defines @ref omni::graph::exec::unstable::IGraphBuilder.
#pragma once

#include <carb/Format.h>

#include <omni/graph/exec/unstable/Assert.h>
#include <omni/graph/exec/unstable/IGraph.h>
#include <omni/graph/exec/unstable/IGraphBuilder.h>
#include <omni/graph/exec/unstable/IGraphBuilderContext.h>
#include <omni/graph/exec/unstable/INode.h>
#include <omni/graph/exec/unstable/INodeDef.h>
#include <omni/graph/exec/unstable/INodeFactory.h>
#include <omni/graph/exec/unstable/INodeGraphDef.h>
#include <omni/graph/exec/unstable/INodeGraphDefDebug.h>

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

//! @copydoc omni::graph::exec::unstable::IGraphBuilder
template <typename... Bases>
class GraphBuilderT : public Implements<Bases...>
{
public:
    //! Construct graph builder for a root @ref INodeGraphDef.
    //!
    //! Construction of a graph builder has a side effect on underlying @c topology causing its invalidation.
    //!
    //! May throw.
    static omni::core::ObjectPtr<GraphBuilderT> create(omni::core::ObjectParam<IGraphBuilderContext> context)
    {
        OMNI_THROW_IF_ARG_NULL(context);
        OMNI_GRAPH_EXEC_ASSERT(
            !omni::graph::exec::unstable::cast<INodeGraphDefDebug>(context->getGraph()->getNodeGraphDef()) ||
            !omni::graph::exec::unstable::cast<INodeGraphDefDebug>(context->getGraph()->getNodeGraphDef())->isExecuting());

        auto builder = omni::core::steal(new GraphBuilderT(context.get(), context->getGraph()->getNodeGraphDef()));
        auto topology = builder->getTopology();
        topology->invalidate();
        builder->_modifiedTopology(topology);

        return builder;
    }

    //! Construct graph builder for a given @ref INodeGraphDef.
    //!
    //! Construction of a graph builder has a side effect on underlying @c topology causing its invalidation.
    //!
    //! May throw.
    static omni::core::ObjectPtr<GraphBuilderT> create(omni::core::ObjectParam<IGraphBuilderContext> context,
                                                       omni::core::ObjectParam<INodeGraphDef> nodeGraphDef)
    {
        OMNI_THROW_IF_ARG_NULL(context);
        OMNI_THROW_IF_ARG_NULL(nodeGraphDef);
        OMNI_GRAPH_EXEC_ASSERT(!omni::graph::exec::unstable::cast<INodeGraphDefDebug>(nodeGraphDef.get()) ||
                               !omni::graph::exec::unstable::cast<INodeGraphDefDebug>(nodeGraphDef.get())->isExecuting());

        auto builder = omni::core::steal(new GraphBuilderT(context.get(), nodeGraphDef.get()));
        auto topology = builder->getTopology();
        topology->invalidate();
        builder->_modifiedTopology(topology);

        return builder;
    }

    //! Construct graph builder for a given @ref INodeGraphDef without causing topology invalidation.
    //!
    //! This builder is used by the pass pipeline when operations to the graph will alter existing topology.
    //!
    //! May throw.
    static omni::core::ObjectPtr<GraphBuilderT> createForPass(omni::core::ObjectParam<IGraphBuilderContext> context,
                                                              omni::core::ObjectParam<INodeGraphDef> nodeGraphDef)
    {
        OMNI_THROW_IF_ARG_NULL(context);
        OMNI_THROW_IF_ARG_NULL(nodeGraphDef);
        OMNI_GRAPH_EXEC_ASSERT(!omni::graph::exec::unstable::cast<INodeGraphDefDebug>(nodeGraphDef.get()) ||
                               !omni::graph::exec::unstable::cast<INodeGraphDefDebug>(nodeGraphDef.get())->isExecuting());

        auto builder = omni::core::steal(new GraphBuilderT(context.get(), nodeGraphDef.get()));

        // Detect when node graph was constructed outside of the pass pipeline. Tag these defs are created
        // during current construction stamp.
        //
        // This usage pattern we only have in tests currently.
        auto topology = nodeGraphDef->getTopology();
        if (!topology->getConstructionStamp().isValid())
        {
            builder->_modifiedTopology(topology);
        }

        return builder;
    }

protected:
    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilder::getGraph_abi
    IGraph* getGraph_abi() noexcept override
    {
        return m_context->getGraph();
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilder::getTopology_abi
    ITopology* getTopology_abi() noexcept override
    {
        return m_nodeGraphDef->getTopology();
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilder::getContext_abi
    IGraphBuilderContext* getContext_abi() noexcept override
    {
        return m_context;
    }

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

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilder::connect_abi
    omni::core::Result connect_abi(INode* upstreamNode, INode* downstreamNode) noexcept override
    {
        try
        {
            _modifiedTopology(upstreamNode->getTopology());

            IGraphBuilderNode* upstream = exec::unstable::cast<exec::unstable::IGraphBuilderNode>(upstreamNode);
            IGraphBuilderNode* downstream = exec::unstable::cast<exec::unstable::IGraphBuilderNode>(downstreamNode);
            if (_connect(upstream, downstream))
            {
                return omni::core::kResultSuccess;
            }

            OMNI_GRAPH_EXEC_RETURN_ERROR(omni::core::kResultFail);
        }
        OMNI_GRAPH_EXEC_CATCH_ABI_EXCEPTION()
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilder::disconnect_abi
    omni::core::Result disconnect_abi(INode* upstreamNode, INode* downstreamNode) noexcept override
    {
        try
        {
            _modifiedTopology(upstreamNode->getTopology());

            IGraphBuilderNode* upstream = exec::unstable::cast<exec::unstable::IGraphBuilderNode>(upstreamNode);
            IGraphBuilderNode* downstream = exec::unstable::cast<exec::unstable::IGraphBuilderNode>(downstreamNode);
            if (_disconnect(upstream, downstream))
            {
                return omni::core::kResultSuccess;
            }

            OMNI_GRAPH_EXEC_RETURN_ERROR(omni::core::kResultFail);
        }
        OMNI_GRAPH_EXEC_CATCH_ABI_EXCEPTION()
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilder::remove_abi
    omni::core::Result remove_abi(INode* node) noexcept override
    {
        _modifiedTopology(node->getTopology());

        try
        {
            IGraphBuilderNode* nodeToRemove = exec::unstable::cast<exec::unstable::IGraphBuilderNode>(node);
            nodeToRemove->validateOrResetTopology();

            // Nodes don't have an edge back to the root. Check if the removed node is a child of the root.
            auto graphRoot = nodeToRemove->getRoot();
            graphRoot->_removeChild(nodeToRemove); // Silently fails if node is not a root child.

            // Cache these pointers to avoid virtual method overhead.
            auto children = nodeToRemove->getChildren();
            auto parents = nodeToRemove->getParents();

            // Disconnect all parents from the node to be removed.
            for (auto parent : parents)
            {
                parent->_removeChild(nodeToRemove);
            }

            // Disconnect all children from the node to be removed.
            for (auto child : children)
            {
                child->_removeParent(nodeToRemove);
            }

            // Invalidate all remaining connections of the node to be removed.
            nodeToRemove->_invalidateConnections();

            return omni::core::kResultSuccess;
        }
        OMNI_GRAPH_EXEC_CATCH_ABI_EXCEPTION()
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilder::setNodeDef_abi
    void setNodeDef_abi(INode* node, INodeDef* nodeDef) noexcept override
    {
        _modifiedTopology(node->getTopology());

        exec::unstable::cast<exec::unstable::IGraphBuilderNode>(node)->_setNodeDef(nodeDef);
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilder::setNodeGraphDef_abi
    void setNodeGraphDef_abi(INode* node, INodeGraphDef* nodeGraphDef) noexcept override
    {
        _modifiedTopology(node->getTopology());

        exec::unstable::cast<exec::unstable::IGraphBuilderNode>(node)->_setNodeGraphDef(nodeGraphDef);
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilder::clearDef_abi
    void clearDef_abi(INode* node) noexcept override
    {
        _modifiedTopology(node->getTopology());

        exec::unstable::cast<exec::unstable::IGraphBuilderNode>(node)->_clearDef();
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilder::replacePartition_abi
    void replacePartition_abi(const NodePartition* partition, IDef* definition) noexcept override
    {
        if (partition->size() == 0)
            return;

        // validate the partition
        INode* rootNode = partition->front()->getRoot();
        for (auto nodeInPartition : *partition)
        {
            if (!nodeInPartition->isValidTopology() || nodeInPartition->getRoot() != rootNode)
            {
                return;
            }
        }

        // mutate the graph
        _commitPartition(m_nodeGraphDef, partition, definition);
        _modifiedTopology(m_nodeGraphDef->getTopology());
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IGraphBuilder::createNode_abi
    INode* createNode_abi(const char* name, IDef* def) noexcept override
    {
        if (auto factory = m_nodeGraphDef->getNodeFactory())
        {
            auto newNode = factory->createNode(name, def);
            m_createdNodes.push_back(newNode.get());
            return newNode.detach();
        }
        else
        {
            return nullptr;
        }
    }

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

    //! Constructor
    GraphBuilderT(IGraphBuilderContext* context, INodeGraphDef* nodeGraphDef)
        : m_context{ context }, m_nodeGraphDef{ nodeGraphDef }
    {
        m_context->getGraph()->_setInBuild(true);
    }

    ~GraphBuilderT()
    {
        m_context->getGraph()->_setInBuild(false);
    }

private:
    //! This builder modified topology of a graph. Currently it is possible it is not modifying topology belonging to
    //! NodeGraphDef it refers to.
    //!
    void _modifiedTopology(ITopology* modifiedTopology)
    {
        modifiedTopology->_setConstructionInSync(m_context->getConstructionStamp());
    }

    bool _connect(IGraphBuilderNode* upstream, IGraphBuilderNode* downstream)
    {
        if (upstream->getTopology() == downstream->getTopology())
        {
            upstream->validateOrResetTopology();
            downstream->validateOrResetTopology();

            if (!upstream->hasChild(downstream))
            {
                upstream->_addChild(downstream);

                if (!upstream->isRoot())
                {
                    downstream->_addParent(upstream);
                }
            }

            return true;
        }

        return false;
    }

    bool _disconnect(IGraphBuilderNode* upstream, IGraphBuilderNode* downstream)
    {
        if (upstream->getTopology() == downstream->getTopology())
        {
            upstream->validateOrResetTopology();
            downstream->validateOrResetTopology();

            if (upstream->hasChild(downstream))
            {
                upstream->_removeChild(downstream);

                if (!upstream->isRoot())
                {
                    downstream->_removeParent(upstream);
                }
            }

            return true;
        }

        return false;
    }

    //! Make changes to the topology with already validated partition and definition.
    void _commitPartition(INodeGraphDef* nodeGraphDef, const NodePartition* partition, IDef* definition) noexcept
    {
        OMNI_GRAPH_EXEC_ASSERT(nodeGraphDef && definition && (partition->size() > 0));
        OMNI_GRAPH_EXEC_ASSERT(nodeGraphDef->getNodeFactory().get());

        // we affect the topology, but shouldn't require any memory operation
        std::vector<INode*> parents, children;

        // optimization, let's assume each node has one parent and one child from outside of the partition
        parents.reserve(partition->size() * 2);
        children.reserve(partition->size() * 2);

        // we want cost to be linear and for that we are going to avoid searches in the partition
        // we achieve this by collecting all parents/children (some will be in the partition),
        // then invalidating the partition and cleaning up the immediate upstream and downstream
        for (auto nodeInPartition : *partition)
        {
            for (auto parent : nodeInPartition->getParents())
            {
                parents.push_back(parent);
            }

            for (auto child : nodeInPartition->getChildren())
            {
                children.push_back(child);
            }

            // make the node invalid without invalidating the entire topology
            exec::unstable::cast<IGraphBuilderNode>(nodeInPartition)->_invalidateConnections();
        }

        // generate replacement node
        std::string nodeName = carb::fmt::format("Partition_{}", partition->front()->getName().getString().c_str());

        auto newNode = createNode_abi(nodeName.c_str(), definition);
        auto newBuilderNode = exec::unstable::cast<IGraphBuilderNode>(newNode);

        // in one pass: cleanup the topology and reconnect to the new node
        auto rootBuilderNode = exec::unstable::cast<IGraphBuilderNode>(partition->front()->getRoot());
        rootBuilderNode->_removeInvalidChildren();

        for (auto parent : parents)
        {
            if (parent->isValidTopology())
            {
                auto parentBuilderNode = exec::unstable::cast<IGraphBuilderNode>(parent);
                parentBuilderNode->_removeInvalidChildren();
                this->_connect(parentBuilderNode, newBuilderNode);
            }
        }

        for (auto child : children)
        {
            if (child->isValidTopology())
            {
                auto childBuilderNode = exec::unstable::cast<IGraphBuilderNode>(child);
                childBuilderNode->_removeInvalidParents();
                this->_connect(newBuilderNode, childBuilderNode);
            }
        }

        // Need to make sure we are connected to the root (indirectly, or directly if this is an entry node)
        if (newNode->getParents().size() == 0)
        {
            this->_connect(rootBuilderNode, newBuilderNode);
        }
    }

    IGraphBuilderContext* m_context{ nullptr }; //!< All graph builders are operating within a context. We store pointer
                                                //!< to it.
    INodeGraphDef* m_nodeGraphDef{ nullptr }; //!< Graph topology this builder can modify. This is not yet enforced in
                                              //!< code.

    //! Most of the time we won't be needing any space. The size of 2 was chosen arbitrary.
    using NodeArray = SmallVector<INode*, 2>;
    NodeArray m_createdNodes; //!< Collect nodes created dynamically to allow pass pipeline discover them.
};

//! Core GraphBuilder implementation for @ref omni::graph::exec::unstable::IGraphBuilder
using GraphBuilder = GraphBuilderT<IGraphBuilder>;

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