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

#include <omni/graph/exec/unstable/ConstName.h>
#include <omni/graph/exec/unstable/IBase.h>
#include <omni/graph/exec/unstable/Span.h>
#include <omni/graph/exec/unstable/Stamp.h>
#include <omni/graph/exec/unstable/Status.h>
#include <omni/graph/exec/unstable/Types.h>

#include <memory>

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

class ExecutionPath;
class ExecutionTask;
class IApplyOnEachFunction;
class IDef;
class IExecutionContext_abi;
class IExecutionContext;
class IExecutionStateInfo;
class INode;

//! Stores and provides access to the execution state of the graph.
//!
//! The @rstref{execution graph <ef_execution_graph>} is only a description of what needs to be executed. The actual
//! graph state is stored separately in an instance of this object.
//!
//! The execution context allows computing the same graph description within multiple contexts. It also enables the
//! ability to perform this computation concurrently.  Some example use cases of this class:
//!
//! - Computing the state of a graph at a time different than the current time (e.g. asynchronous caching, fake
//!   dynamics)
//!
//! - Computing the state of a graph with a inputs different than the current input state (e.g. double solve)
//!
//! All execution begins with a call to one of the execute methods on this interface. @ref
//! omni::graph::exec::unstable::IExecutionContext::execute() is used to execute the entire execution graph while @ref
//! omni::graph::exec::unstable::IExecutionContext::executeNode() can be used to execute only a part of the graph.
//!
//! Part of this interface defines a key/value store.  The *key* in this store is an @ref
//! omni::graph::exec::unstable::ExecutionPath.  The *value* is an implementation of @ref
//! omni::graph::exec::unstable::IExecutionStateInfo, which in addition to storing computation state, can also store
//! user defined data.  The computation state and user data can be access with one of the `getStateInfo` /
//! `setStateInfo` methods though if you wish to access the user data using one of the `getNodeData` / `setNodeData` is
//! slightly faster.
//!
//! Another feature of @ref omni::graph::exec::unstable::IExecutionContext is the ability to quickly search for nodes
//! using a particular definition and apply a function on them.  Definitions can be searched by name or by pointer (see
//! @ref omni::graph::exec::unstable::IExecutionContext::applyOnEach()).  These methods are used extensively during
//! @rstref{Graph Construction <ef_pass_concepts>}.
//!
//! See @rstref{Execution Concepts <ef_execution_concepts>} for an in-depth guide on how this object is used during
//! execution.
//!
//! Since multiple threads can concurrently traverse a graph, implementors of methods within this class should expect
//! that multiple threads will be accessing this object in parallel.
class IExecutionContext_abi : public omni::core::Inherits<omni::graph::exec::unstable::IBase,
                                                          OMNI_TYPE_ID("omni.graph.exec.unstable.IExecutionContext")>
{
protected:
    //! Current execution version. Incremented with each execution of the context.
    //!
    //! @thread_safety See thread safety information in interface description.
    virtual Stamp getExecutionStamp_abi() noexcept = 0;

    //! Returns `true` if context is currently executing
    //!
    //! See thread safety information in interface description.
    virtual bool inExecute_abi() noexcept = 0;

    //! Returns `true` if the current thread is the one who started this context's execution.
    //!
    //! Note, do not assume that the thread that started the context's execution is the "main" thread.
    //!
    //! @thread_safety See thread safety information in interface description.
    virtual bool isExecutingThread_abi() noexcept = 0;

    //! Main execution method.  Executes the entire execution graph.
    //!
    //! @thread_safety See thread safety information in interface description.
    virtual Status execute_abi() noexcept = 0;

    //! On-demand execution method.  Executes the given node.
    //!
    //! @thread_safety See thread safety information in interface description.
    virtual Status executeNode_abi(OMNI_ATTR("in, not_null, throw_if_null, ref") const ExecutionPath* upstreamPath,
                                   OMNI_ATTR("not_null, throw_if_null") INode* node) noexcept = 0;

    //! Context initialization. Responsible to propagate initialization to graphs.
    //!
    //! @thread_safety See thread safety information in interface description.
    virtual OMNI_ATTR("throw_result") omni::core::Result initialize_abi() noexcept = 0;

    //! Access state for a given execution path.
    //!
    //! If the given node is not @c nullptr, a copy of the given path with the node appended will be used as the lookup
    //! key.
    //!
    //! @thread_safety See thread safety information in interface description.
    //!
    //! @warning    This method should be used for read only access by downstream nodes, example accessing graph state
    //!             when executing downstream nodes. Extra care needs to be taken if this state has to be mutated
    //!             concurrently.
    virtual OMNI_ATTR("throw_result") omni::core::Result
        getStateInfo_abi(OMNI_ATTR("in, not_null, throw_if_null, ref") const ExecutionPath* path,
                         INode* node,
                         OMNI_ATTR("out, *return, not_null, throw_if_null, *no_acquire")
                             IExecutionStateInfo** out) noexcept = 0;

    //! Returns a value from a node's key/value datastore.
    //!
    //! The node from which to grab data is identified by the given path and
    //! node.  The @p node may be @c nullptr.
    //!
    //! The key is used as a look-up in the node's key/value datastore.
    //!
    //! The type of each data item is returned in @p outTypeId.
    //!
    //! @p outPtr will be updated with a pointer to the actual data.
    //!
    //! @p outItemSize store the size of each item in the returned array.
    //!
    //! @p outItemCount contains the number of items returned (i.e. the number
    //! of items @p outPtr points to).  For an array, this will be greater than
    //! 1.
    //!
    //! If the key is not found, @p outPtr is set to @c nullptr and @p
    //! outItemCount is set to 0.
    //!
    //! Accessing the node's key/value datastore is not thread safe.
    //!
    //! An exception is thrown on all other errors.
    virtual OMNI_ATTR("throw_result") omni::core::Result
        getNodeData_abi(OMNI_ATTR("in, not_null, throw_if_null, ref") const ExecutionPath* path,
                        INode* node,
                        NodeDataKey key,
                        OMNI_ATTR("out, not_null, throw_if_null") omni::core::TypeId* outTypeId,
                        OMNI_ATTR("out, not_null, throw_if_null, *out, *in") void** outPtr,
                        OMNI_ATTR("out, not_null, throw_if_null") uint64_t* outItemSize,
                        OMNI_ATTR("out, not_null, throw_if_null") uint64_t* outItemCount) noexcept = 0;

    //! Sets a value in a node's key/value datastore.
    //!
    //! The node in which to set the data is identified by the given path and
    //! node.  The @p node may be @c nullptr.
    //!
    //! The key is used as a look-up in the node's key/value datastore.
    //!
    //! The type of each data item is set with @p typeId.
    //!
    //! @p data points to an array of data items.
    //!
    //! @p itemSize is the size of each item in the given array.
    //!
    //! @p itemCount contains the number of items pointed to by @p data. For an
    //! array, this will be greater than 1.
    //!
    //! @p deleter is a function used to delete @p data when either a new value
    //! is set at the key or the context is invalidated.  If @p deleter is @c
    //! nullptr, it is up to the calling code to manage the lifetime of the @p
    //! data.
    //!
    //! Accessing the node's key/value datastore is not thread safe.
    //!
    //! An exception is thrown on all other errors.
    virtual OMNI_ATTR("throw_result") omni::core::Result
        setNodeData_abi(OMNI_ATTR("in, not_null, throw_if_null, ref") const ExecutionPath* path,
                        INode* node,
                        NodeDataKey key,
                        omni::core::TypeId typeId,
                        OMNI_ATTR("in, out, not_null, throw_if_null") void* data,
                        uint64_t itemSize,
                        uint64_t itemCount,
                        OMNI_ATTR("in, out") NodeDataDeleterFn* deleter) noexcept = 0;

    //! Discover all execution paths leading to given definition and invoke given function with each of them.
    //!
    //! Implementation of this interface will cache results to remove the traversal cost for subsequent call. Change in
    //! execution graph topology will cause invalidation of the cache.
    //!
    //! @thread_safety This method is thread-safe and thread-efficient and can be called during execution, as well as
    //! from within the provided @p callback.
    //!
    //! This method must not be called during graph construction.
    //!
    //! @p def definition to look for
    //! @p callback wrapped function into a callback to execute with each path to given definition
    virtual void applyOnEachDef_abi(OMNI_ATTR("not_null") IDef* def,
                                    OMNI_ATTR("not_null") IApplyOnEachFunction* callback) noexcept = 0;

    //! Discover all execution paths leading to definition with the given name and invoke the given function with each
    //! of them
    //!
    //! Implementation of this interface will cache results to remove the traversal cost for subsequent call. Change in
    //! execution graph topology will cause invalidation of the cache.
    //!
    //! @thread_safety This method is thread-safe and thread-efficient and can be called during execution, as well as
    //! from within the provided @p callback.
    //!
    //! This method must not be called during graph construction.
    //!
    //! @p name definition to look for
    //! @p callback wrapped function into a callback to execute with each path to given definition
    virtual void applyOnEachDefWithName_abi(OMNI_ATTR("in, not_null") const ConstName* name,
                                            OMNI_ATTR("not_null") IApplyOnEachFunction* callback) noexcept = 0;
};

//! Smart pointer managing an instance of @ref IExecutionContext.
using ExecutionContextPtr = omni::core::ObjectPtr<IExecutionContext>;

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

#define OMNI_BIND_INCLUDE_INTERFACE_DECL
#include <omni/graph/exec/unstable/IExecutionContext.gen.h>

//! @copydoc omni::graph::exec::unstable::IExecutionContext_abi
//!
//! @ingroup groupOmniGraphExecInterfaces
class omni::graph::exec::unstable::IExecutionContext
    : public omni::core::Generated<omni::graph::exec::unstable::IExecutionContext_abi>
{
public:
    //! @copydoc omni::graph::exec::unstable::IExecutionContext_abi::executeNode_abi
    Status execute(const ExecutionPath& path, omni::core::ObjectParam<INode> node)
    {
        return executeNode(path, node);
    }
    using Generated<IExecutionContext_abi>::execute;

    //! Retrieves the state info for the given task.
    //!
    //! The task's upstream path and node are hashed for the lookup.
    //!
    //! The returned @ref omni::graph::exec::unstable::IExecutionStateInfo will
    //! not have @ref omni::core::IObject::acquire() called on it.
    inline IExecutionStateInfo* getStateInfo(const ExecutionTask& task);
    using Generated<IExecutionContext_abi>::getStateInfo;

    //! Retrieves the state info at the given execution path.
    //!
    //! The returned @ref omni::graph::exec::unstable::IExecutionStateInfo will
    //! not have @ref omni::core::IObject::acquire() called on it.
    inline IExecutionStateInfo* getStateInfo(const ExecutionPath& path);

    //! Access state for a given execution path.
    //!
    //! See thread safety information in interface description.
    template <typename T>
    inline T* getStateInfoAs(const ExecutionTask& info);

    //! Access state for a given execution path.
    //!
    //! See thread safety information in interface description.
    template <typename T>
    inline T* getStateInfoAs(const ExecutionPath& path);

    //! Access state for a given execution path.
    //!
    //! See thread safety information in interface description.
    template <typename T>
    inline T* getStateInfoAs(const ExecutionPath& path, omni::core::ObjectParam<INode> node);

    //! Returns a pointer to a value stored in a node's key/value datastore.
    //!
    //! The node whose key/value datastore should be used is identified by combining the given path and node.  @p node
    //! may be @c nullptr.
    //!
    //! If there is no value stored at the given @p key an empty span is returned.
    //!
    //! Accessing the node's key/value datastore is not thread safe.
    //!
    //! If the type @c T does not match the type of the store data, an exception is thrown.
    //!
    //! An exception is thrown on all other errors.
    //!
    //! Prefer using @ref OMNI_GRAPH_EXEC_GET_NODE_DATA_AS() instead of this method, which will populate the type id for
    //! you.
    template <typename T>
    inline Span<T> getNodeDataAs(omni::core::TypeId desiredType,
                                 const ExecutionPath& path,
                                 omni::core::ObjectParam<INode> node,
                                 NodeDataKey key);

    //! Returns a pointer to a value stored in a node's key/value datastore.
    //!
    //! The node whose key/value datastore should be used is identified by combining the given path and node in the
    //! given task.
    //!
    //! If there is no value stored at the given @p key an empty span is returned.
    //!
    //! Accessing the node's key/value datastore is not thread safe.
    //!
    //! If the type @c T does not match the type of the store data, an exception is thrown.
    //!
    //! An exception is thrown on all other errors.
    //!
    //! Prefer using @ref OMNI_GRAPH_EXEC_GET_NODE_DATA_AS() instead of this method, which will populate the type id for
    //! you.
    template <typename T>
    inline Span<T> getNodeDataAs(omni::core::TypeId desiredType, const ExecutionTask& path, NodeDataKey key);

    //! Stores a value in a node's key/value datastore.
    //!
    //! The node whose key/value datastore should be used is identified by combining the given path and node.  @p node
    //! may be @c nullptr.
    //!
    //! If a value is already stored at the given @p key it will be replaced.
    //!
    //! Accessing the node's key/value datastore is not thread safe.
    //!
    //! An exception is thrown on all errors.
    //!
    //! Prefer using @ref OMNI_GRAPH_EXEC_SET_NODE_DATA() instead of this method, which will populate the type id for
    //! you.
    template <typename SpecifiedT, typename DataT>
    inline void setNodeData(omni::core::TypeId itemType,
                            const ExecutionPath& path,
                            omni::core::ObjectParam<INode> node,
                            NodeDataKey key,
                            std::unique_ptr<DataT> data);

    //! Stores a value in a node's key/value datastore.
    //!
    //! The node whose key/value datastore should be used is identified by combining the given path and node in the
    //! given task.
    //!
    //! If a value is already stored at the given @p key it will be replaced.
    //!
    //! Accessing the node's key/value datastore is not thread safe.
    //!
    //! An exception is thrown on all errors.
    //!
    //! Prefer using @ref OMNI_GRAPH_EXEC_SET_NODE_DATA() instead of this method, which will populate the type id for
    //! you.
    template <typename SpecifiedT, typename DataT>
    inline void setNodeData(omni::core::TypeId itemType,
                            const ExecutionTask& path,
                            NodeDataKey key,
                            std::unique_ptr<DataT> data);

    //! Discover all execution paths leading to given definition and invoke given function with each of them
    //!
    //! This inline implementation wraps lambda into IApplyOnEachFunction
    //!
    //! The supplied function should have the signature of `void(const ExecutionPath&)`.
    template <typename FN>
    inline void applyOnEach(omni::core::ObjectParam<omni::graph::exec::unstable::IDef> def, FN&& callback);

    //! Discover all execution paths leading to definition with the given name and invoke the given function with each
    //! of them
    //!
    //! This inline implementation wraps lambda into IApplyOnEachFunction
    //!
    //! The supplied function should have the signature of `void(const ExecutionPath&)`.
    template <typename FN>
    inline void applyOnEach(const ConstName& name, FN&& callback);
};

#include <omni/graph/exec/unstable/ExecutionPath.h>
#include <omni/graph/exec/unstable/ExecutionTask.h>
#include <omni/graph/exec/unstable/IApplyOnEachFunction.h>
#include <omni/graph/exec/unstable/IDef.h>
#include <omni/graph/exec/unstable/IExecutionStateInfo.h>
#include <omni/graph/exec/unstable/INode.h>

#ifndef DOXYGEN_BUILD // templates are hard to understand, even for doxygen...

inline omni::graph::exec::unstable::IExecutionStateInfo* omni::graph::exec::unstable::IExecutionContext::getStateInfo(
    const ExecutionTask& info)
{
    IExecutionStateInfo* out;
    OMNI_THROW_IF_FAILED(getStateInfo_abi(&(info.getUpstreamPath()), info.getNode(), &out));
    return out;
}

inline omni::graph::exec::unstable::IExecutionStateInfo* omni::graph::exec::unstable::IExecutionContext::getStateInfo(
    const ExecutionPath& path)
{
    IExecutionStateInfo* out;
    OMNI_THROW_IF_FAILED(getStateInfo_abi(&path, nullptr, &out));
    return out;
}

template <typename T>
inline T* omni::graph::exec::unstable::IExecutionContext::getStateInfoAs(const ExecutionTask& info)
{
    IExecutionStateInfo* out;
    OMNI_THROW_IF_FAILED(getStateInfo_abi(&(info.getUpstreamPath()), info.getNode(), &out));
    return omni::graph::exec::unstable::cast<T>(out);
}

template <typename T>
inline T* omni::graph::exec::unstable::IExecutionContext::getStateInfoAs(const ExecutionPath& path)
{
    IExecutionStateInfo* out;
    OMNI_THROW_IF_FAILED(getStateInfo_abi(&path, nullptr, &out));
    return omni::graph::exec::unstable::cast<T>(out);
}

template <typename T>
inline T* omni::graph::exec::unstable::IExecutionContext::getStateInfoAs(const ExecutionPath& path,
                                                                         omni::core::ObjectParam<INode> node)
{
    IExecutionStateInfo* out;
    OMNI_THROW_IF_FAILED(getStateInfo_abi(&path, node.get(), &out));
    return omni::graph::exec::unstable::cast<T>(out);
}

template <typename T>
inline omni::graph::exec::unstable::Span<T> omni::graph::exec::unstable::IExecutionContext::getNodeDataAs(
    omni::core::TypeId desiredType, const ExecutionPath& path, omni::core::ObjectParam<INode> node, NodeDataKey key)
{
    omni::core::TypeId outType;
    void* outPtr;
    uint64_t outItemSize, outItemCount;
    OMNI_THROW_IF_FAILED(getNodeData_abi(&path, node.get(), key, &outType, &outPtr, &outItemSize, &outItemCount));

    if (outPtr)
    {
        if (outType != desiredType)
        {
            throw omni::core::ResultError(omni::core::kResultInvalidDataType);
        }

        if (outItemSize != sizeof(T))
        {
            throw omni::core::ResultError(omni::core::kResultInvalidDataSize);
        }
    }

    return Span<T>{ reinterpret_cast<T*>(outPtr), outItemCount };
}

template <typename T>
inline omni::graph::exec::unstable::Span<T> omni::graph::exec::unstable::IExecutionContext::getNodeDataAs(
    omni::core::TypeId desiredType, const ExecutionTask& info, NodeDataKey key)
{
    return getNodeDataAs<T>(desiredType, info.getUpstreamPath(), info.getNode(), key);
}

template <typename SpecifiedT, typename DataT>
inline void omni::graph::exec::unstable::IExecutionContext::setNodeData(omni::core::TypeId itemType,
                                                                        const ExecutionPath& path,
                                                                        omni::core::ObjectParam<INode> node,
                                                                        NodeDataKey key,
                                                                        std::unique_ptr<DataT> data)
{
    static_assert(std::is_same<SpecifiedT, DataT>::value, "given TypeId does not match the data type");
    static_assert(!std::is_array<DataT>::value, "setting arrays as node data via unique_ptr not yet implemented");

    OMNI_THROW_IF_FAILED(setNodeData_abi(&path, node.get(), key, itemType, data.get(), sizeof(DataT), 1,
                                         [](void* p)
                                         {
                                             typename std::unique_ptr<DataT>::deleter_type deleter;
                                             deleter(reinterpret_cast<DataT*>(p));
                                         }));
    data.release(); // now safe to release ownership
}

template <typename SpecifiedT, typename DataT>
inline void omni::graph::exec::unstable::IExecutionContext::setNodeData(omni::core::TypeId itemType,
                                                                        const ExecutionTask& task,
                                                                        NodeDataKey key,
                                                                        std::unique_ptr<DataT> data)
{
    return setNodeData<SpecifiedT>(itemType, task.getUpstreamPath(), task.getNode(), key, std::move(data));
}

template <typename FN>
inline void omni::graph::exec::unstable::IExecutionContext::applyOnEach(
    omni::core::ObjectParam<omni::graph::exec::unstable::IDef> def, FN&& callback)
{
    class Forwarder : public Implements<IApplyOnEachFunction>
    {
    public:
        Forwarder(FN&& fn) : m_fn(std::move(fn))
        {
        }

    protected:
        void invoke_abi(const ExecutionPath* path) noexcept override
        {
            m_fn(*path);
        }

        FN m_fn;
    };

    applyOnEachDef_abi(def.get(), omni::core::steal(new Forwarder(std::forward<FN>(callback))).get());
}

template <typename FN>
inline void omni::graph::exec::unstable::IExecutionContext::applyOnEach(const ConstName& name, FN&& callback)
{
    class Forwarder : public Implements<IApplyOnEachFunction>
    {
    public:
        Forwarder(FN&& fn) : m_fn(std::move(fn))
        {
        }

    protected:
        void invoke_abi(const ExecutionPath* path) noexcept override
        {
            m_fn(*path);
        }

        FN m_fn;
    };

    applyOnEachDefWithName_abi(&name, omni::core::steal(new Forwarder(std::forward<FN>(callback))).get());
}
#endif // DOXYGEN_BUILD

#define OMNI_BIND_INCLUDE_INTERFACE_IMPL
#include <omni/graph/exec/unstable/IExecutionContext.gen.h>
