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

#include <omni/graph/exec/unstable/Assert.h>
#include <omni/graph/exec/unstable/Executor.h>
#include <omni/graph/exec/unstable/ExecutorFactory.h>
#include <omni/graph/exec/unstable/INodeGraphDef.h>
#include <omni/graph/exec/unstable/INodeGraphDefDebug.h>
#include <omni/graph/exec/unstable/Topology.h>

#include <atomic>

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

//! @copydoc omni::graph::exec::unstable::INodeGraphDef
template <typename... Bases>
class NodeGraphDefT : public Implements<Bases...>
{
public:
    //! Construct graph node definition with default executor
    //!
    //! @param  owner  Execution graph having this graph as part of the global topology
    //! @param  definitionName  Definition name is considered as a token that transformation passes can register against
    //!
    //! May throw.
    static omni::core::ObjectPtr<NodeGraphDefT> create(omni::core::ObjectParam<IGraph> owner, const char* definitionName)
    {
        OMNI_THROW_IF_ARG_NULL(owner);
        OMNI_THROW_IF_ARG_NULL(definitionName);
        return omni::core::steal(new NodeGraphDefT(owner, definitionName));
    }

    //! Construct graph node definition with a given executor factory
    //!
    //! @param  owner  Execution graph having this graph as part of the global topology
    //! @param  executorFactory Factory returning executor for this graph
    //! @param  definitionName  Definition name is considered as a token that transformation passes can register against
    //!
    //! May throw.
    static omni::core::ObjectPtr<NodeGraphDefT> create(omni::core::ObjectParam<IGraph> owner,
                                                       const ExecutorFactory& executorFactory,
                                                       const char* definitionName)
    {
        OMNI_THROW_IF_ARG_NULL(owner);
        OMNI_THROW_IF_ARG_NULL(definitionName);
        return omni::core::steal(new NodeGraphDefT(owner, executorFactory, definitionName));
    }

protected:
    //! Core implementation of @ref omni::graph::exec::unstable::IDef::execute_abi for @ref NodeGraphDef
    //!
    //! Execution is delegated to @ref omni::graph::exec::unstable::IExecutor. The lifetime of an executor is only for a
    //! single execution and any state that needs to persist longer than a single execution must be written with @ref
    //! omni::graph::exec::unstable::IExecutionContext::setNodeData_abi()
    Status execute_abi(ExecutionTask* info) noexcept override
    {
        // ef-docs nodegraphdef-execute-begin
        omni::core::ObjectPtr<IExecutor> executor;
        if (m_executorFactory)
        {
            executor = m_executorFactory(m_topology, *info);
        }
        else
        {
            executor = ExecutorFallback::create(m_topology, *info);
        }
        return executor->execute(); // execute the node specified by info->getNode()
        // ef-docs nodegraphdef-execute-end
    }

    //! Core implementation of @ref omni::graph::exec::unstable::IDef::getSchedulingInfo_abi for @ref NodeGraphDef
    SchedulingInfo getSchedulingInfo_abi(const ExecutionTask* info) noexcept override
    {
        return SchedulingInfo::eSerial;
    }

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

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

    //! Core implementation of @ref omni::graph::exec::unstable::INodeGraphDef::initializeState_abi
    omni::core::Result initializeState_abi(ExecutionTask* rootTask) noexcept override
    {
        return omni::core::kResultSuccess;
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INodeGraphDef::preExecute_abi
    Status preExecute_abi(ExecutionTask* info) noexcept override
    {
        return Status::eSuccess;
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INodeGraphDef::postExecute_abi
    Status postExecute_abi(ExecutionTask* info) noexcept override
    {
        return Status::eSuccess;
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INodeGraphDef::getNodeFactory_abi
    INodeFactory* getNodeFactory_abi() noexcept override
    {
        return nullptr;
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INodeGraphDefDebug::getExecutionCount_abi
    uint64_t getExecutionCount_abi() noexcept override
    {
        return m_executionCount;
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INodeGraphDefDebug::incrementExecutionCount_abi
    void incrementExecutionCount_abi() noexcept override
    {
        ++m_executionCount;
    }

    //! Core implementation of @ref omni::graph::exec::unstable::INodeGraphDefDebug::decrementExecutionCount_abi
    void decrementExecutionCount_abi() noexcept override
    {
        --m_executionCount;
    }

    //! Constructor with a default executor
    NodeGraphDefT(omni::core::ObjectParam<IGraph> owner, const char* definitionName) // may throw
        : m_topology{ Topology::create(definitionName) }, m_name{ definitionName }
    {
        _addInvalidationForwarder(owner);
    }

    //! Constructor with a custom executor
    NodeGraphDefT(omni::core::ObjectParam<IGraph> owner,
                  ExecutorFactory executorFactory,
                  const char* definitionName) // may throw
        : m_topology{ Topology::create(definitionName) },
          m_executorFactory(std::move(executorFactory)),
          m_name{ definitionName }
    {
        _addInvalidationForwarder(owner);
    }

private:
    //! Private method that will allow forwarding of topology invalidation to the execution graph.
    //! Invalid global topology will allow pass pipeline to execute and discover invalidated definitions.
    void _addInvalidationForwarder(omni::core::ObjectParam<IGraph> owner) // may throw
    {
        m_topology->addInvalidationForwarder(reinterpret_cast<InvalidationForwarderId>(owner.get()),
                                             [global = owner->getGlobalTopologyStamp()](ITopology*) -> void
                                             { global->next(); });
    }

    omni::core::ObjectPtr<ITopology> m_topology; //!< Graphs topology
    ExecutorFactory m_executorFactory; //!< Executor factory (if empty, default executor will be used)

    std::atomic<std::size_t> m_executionCount{ 0 }; //!< Debugging counter to detect illegal executions.

    ConstName m_name; //!< Definition name
};

//! Core NodeGraphDef implementation for @ref omni::graph::exec::unstable::INodeGraphDef
using NodeGraphDef = NodeGraphDefT<INodeGraphDef, INodeGraphDefDebug>;

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