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

#include <omni/core/ITypeFactory.h>
#include <omni/graph/exec/unstable/Assert.h>
#include <omni/graph/exec/unstable/Status.h>

#include <limits>

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

// forward declarations needed by interface declaration
class ExecutionPath;
class IExecutionContext;
class IExecutionCurrentThread;
class IExecutor;
class INode;

//! A task tag can be used by other entities (such as @ref Executor) to group tasks together.
using ExecutionTaskTag = uint64_t;

//! Represents work item generated by an @ref IExecutor and passed to a scheduler for dispatch.
//!
//! ExecutionTask is a utility class that describes a task to be potentially executed on behalf of a @ref INode in a
//! given @ref IExecutionContext.
//!
//! @rst
//!
//! .. image:: /../docs/ef-execution-path-point-k.svg
//!     :align: center
//!
//! @endrst
//!
//! @ref ExecutionTask stores four key pieces of information:
//!
//! - *A pointer to the Node to be executed*. The pointed to @ref INode contains a pointer to either an @ref INodeDef or
//!   @ref NodeGraphDef which contains the computation definition.  See @ref ExecutionTask::getNode().
//!
//! - *The unique path to the node*. In addition to the @ref INode to be executed, an @ref ExecutionPath to the node's
//!   upstream (i.e. containing) node is stored. Combined, these two pieces of information form a unique id for the
//!   node.
//!
//!   Above, if an @ref ExecutionTask is describing the *k* node pointed to by the yellow arrow, @ref
//!   ExecutionTask::getNode() would point to *k* and @ref ExecutionTask::getUpstreamPath() would return */f/p*. Note,
//!   the @ref ExecutionTask::getUpstreamPath() *does not* store */f/p/k*, just */f/p*. This is a micro-optimization
//!   that allows the same |ExecutionPath| to be reused while visiting nodes within the same
//!   @ref INodeGraphDef.
//!
//! - *A pointer to the current execution's* @ref IExecutionContext. Execution always happen in a given context. It's
//!   this context, @ref IExecutionContext, that stores the state of the execution. Multiple entities can be executing a
//!   given @rstref{execution graph <ef_execution_graph>}, each execution using its own @ref IExecutionContext. In order
//!   to understand which of these potentially many executions a task is a part, @ref ExecutionTask stores a reference
//!   to the execution's @ref IExecutionContext. This @ref ExecutionTask::getContext() combined with @ref
//!   ExecutionTask::getUpstreamPath() and @ref ExecutionTask::getNode() can be used to access the per-execution state
//!   for the node (see @ref OMNI_GRAPH_EXEC_GET_NODE_DATA_AS() and @ref OMNI_GRAPH_EXEC_SET_NODE_DATA()).
//!
//! - *A "tag" to identify a task when multiple tasks are associated with a node.*  If an @ref INode generates many
//!   tasks during execution, @ref ExecutionTask::getTag() can be used to uniquely identify each of the node's tasks.
//!   The semantic meaning of @ref ExecutionTask::getTag() is @ref IExecutor dependent and can be used for purposes
//!   other than unique identification.
//!
//! This struct is ABI-safe.
//!
//! See @rstref{Execution Concepts <ef_execution_concepts>} for an in-depth guide on how this object is used during
//! execution.
//!
// @note   Node definition needs to stay unchanged while there are executions to complete. In future we should fully
//         decouple description of compute from compute symbol library. This would allow mutating the authoring side
//         without the need to synchronize with execution.
class ExecutionTask
{
public:
    enum : ExecutionTaskTag
    {
        kEmptyTag = std::numeric_limits<ExecutionTaskTag>::max() //!< Special value to represent an empty tag.
    };

    static_assert(std::numeric_limits<ExecutionTaskTag>::max() == 0xFFFFFFFFFFFFFFFF,
                  "unexpected ExecutionTaskTag max value");

    //! Constructor for execution task
    //!
    //! @param  context  Context in which execution task is created. Task can only access state from this context.
    //!                  @ref omni::core::IObject::acquire() is not called on this context.  It is up to the calling
    //!                  code to ensure the context remains valid for the lifetime of the ExecutionTask.
    //!
    //! @param  node     Node holding the execution definition omni::core::IObject::acquire() is not called on this
    //!                  context.  It is up to the calling code to ensure the context remains valid for the lifetime of
    //!                  the ExecutionTask.
    //!
    //! @param  upPath   Execution path to the graph owning the node. Node can be executed multiple times with different
    //!                  paths when graph definition is shared.
    //!
    //! @param  tag      Used to identify dynamically generated work items that node can compute.
    ExecutionTask(IExecutionContext* context, INode* node, const ExecutionPath& upPath, ExecutionTaskTag tag = kEmptyTag) noexcept
        : m_context(context), m_node(node), m_upstreamPath(&upPath), m_tag(tag)
    {
        OMNI_GRAPH_EXEC_ASSERT(context);
        OMNI_GRAPH_EXEC_ASSERT(node);

        static_assert(std::is_standard_layout<ExecutionTask>::value, "ExecutionTask is expected to be abi safe");
        static_assert(offsetof(ExecutionTask, m_context) == 0, "unexpected context offset");
        static_assert(offsetof(ExecutionTask, m_node) == 8, "unexpected node offset");
        static_assert(offsetof(ExecutionTask, m_upstreamPath) == 16, "unexpected upstream path offset");
        static_assert(offsetof(ExecutionTask, m_tag) == 24, "unexpected tag offset");
        static_assert(offsetof(ExecutionTask, m_userIndex) == 32, "unexpected status offset");
        static_assert(offsetof(ExecutionTask, m_status) == 40, "unexpected status offset");
        static_assert(48 == sizeof(ExecutionTask), "ExecutionTask is an unexpected size");
    }

    //! Return context for this task.
    //!
    //! The returned @ref IExecutionContext will *not* have @ref omni::core::IObject::acquire() called before being
    //! returned.
    IExecutionContext* getContext() const noexcept
    {
        return m_context;
    }

    //! Return node for this task.
    //!
    //! The returned @ref INode will *not* have @ref omni::core::IObject::acquire() called before being returned.
    INode* getNode() const noexcept
    {
        return m_node;
    }

    //! Return execution path to graph owning the node.
    const ExecutionPath& getUpstreamPath() const noexcept
    {
        return *m_upstreamPath;
    }

    //! Return tag.
    ExecutionTaskTag getTag() const noexcept
    {
        return m_tag;
    }

    //! Check if this task has a valid tag set. This will mean that a node generates more than one task.
    bool hasValidTag() const noexcept
    {
        return (m_tag != kEmptyTag);
    }

    //! Return execution status for this task
    Status getExecutionStatus() const noexcept
    {
        return m_status;
    }

    //! Execute the task. Will be called by the scheduler when task is dispatched for execution.
    inline Status execute(omni::core::ObjectParam<IExecutor> executor) noexcept;

    //! This index will never be used by the framework, but is a way to pass something into
    //! user code via generated task. Mutating this value is allowed as long as it is done
    //! via only legal way to access task, i.e. getCurrentTask
    //!
    //! Setter for user index
    void setUserIndex(uint64_t userIndex) noexcept
    {
        m_userIndex = userIndex;
    }

    //! This index will never be used by the framework, but is a way to pass something into
    //! user code via generated task. Mutating this value is allowed as long as it is done
    //! via only legal way to access task, i.e. getCurrentTask
    //!
    //! Getter for user index
    uint64_t getUserIndex() const noexcept
    {
        return m_userIndex;
    }

    //! Sets the status of the task.
    //!
    //! This is an internal method and should not be called by users.
    void setExecutionStatus(Status status) noexcept
    {
        m_status = status;
    }

private:
    //! Context in which this task was created. This context needs to live as long as there are still executions to
    //! complete.
    IExecutionContext* m_context;

    //! Node holding the execution definition.
    INode* m_node;

    //! Execution path to the graph owning the node.
    const ExecutionPath* m_upstreamPath;

    //! Used to identify dynamically generated work items that node can compute.
    ExecutionTaskTag m_tag;

    //! User index help with passing data into user code.
    uint64_t m_userIndex{ 0 };

    //! Execution status.
    Status m_status{ Status::eUnknown };

    //! Reserved padding space.
    uint32_t m_reserved;
};

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

#include <omni/graph/exec/unstable/IExecutionCurrentThread.h>

//! Execute the task. Will be called by the scheduler when task is dispatched for execution.
inline omni::graph::exec::unstable::Status omni::graph::exec::unstable::ExecutionTask::execute(
    omni::core::ObjectParam<IExecutor> executor) noexcept
{
    if (Status::eUnknown != m_status)
    {
        return m_status;
    }

    return getCurrentThread()->execute(*this, executor, &m_status);
}
