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

#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/Types.h>

#include <memory>

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

// forward declarations needed by interface declaration
class IBackgroundResult;
class IExecutionStateInfo;
class IExecutionStateInfo_abi;

//! State associated with a given execution task
//!
//! @note   We separated execution state from the execution graph to allow concurrent and/or nested execution
class IExecutionStateInfo_abi : public omni::core::Inherits<omni::graph::exec::unstable::IBase,
                                                            OMNI_TYPE_ID("omni.graph.exec.unstable.IExecutionStateInfo")>
{
protected:
    //! Store a "future" result for this state. The actual computation is running asynchronously outside of execution
    //! frame
    //!
    //! @return \c true if execution state accepts "future" results.
    virtual bool storeBackgroundResult_abi(OMNI_ATTR("not_null, throw_if_null") IBackgroundResult* result) noexcept = 0;

    //! Query used by some executors to determine if computation of a node is necessary
    virtual bool needsCompute_abi(Stamp execVersion) noexcept = 0;

    //! Set to request computation
    virtual void requestCompute_abi() noexcept = 0;

    //! Reset request to compute after computation was performed
    virtual void setComputed_abi() noexcept = 0;

    //! Get current/last exec version set for this node during execution
    virtual SyncStamp getExecutionStamp_abi() noexcept = 0;

    //! Set current exec version for this node. Returns true if version wasn't in sync.
    virtual bool setExecutionStamp_abi(Stamp execVersion) noexcept = 0;

    //! Returns a value from a node's key/value datastore.
    //!
    //! 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(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 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(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;
};

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

// generated API declaration
#define OMNI_BIND_INCLUDE_INTERFACE_DECL
#include <omni/graph/exec/unstable/IExecutionStateInfo.gen.h>

//! @copydoc omni::graph::exec::unstable::IExecutionStateInfo_abi
//!
//! @ingroup groupOmniGraphExecInterfaces
class omni::graph::exec::unstable::IExecutionStateInfo
    : public omni::core::Generated<omni::graph::exec::unstable::IExecutionStateInfo_abi>
{
public:
    //! Returns a pointer to a value stored in the node's key/value datastore.
    //!
    //! 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, NodeDataKey key);

    //! Stores a value in the node's key/value datastore.
    //!
    //! 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, NodeDataKey key, std::unique_ptr<DataT> data);
};

// additional headers needed for API implementation
#include <omni/graph/exec/unstable/IBackgroundResult.h>

#ifndef DOXYGEN_BUILD // templates and doxygen are not friends (remove this line to see why)

template <typename T>
inline omni::graph::exec::unstable::Span<T> omni::graph::exec::unstable::IExecutionStateInfo::getNodeDataAs(
    omni::core::TypeId desiredType, NodeDataKey key)
{
    omni::core::TypeId outType;
    void* outPtr;
    uint64_t outItemSize, outItemCount;
    OMNI_THROW_IF_FAILED(getNodeData_abi(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 SpecifiedT, typename DataT>
inline void omni::graph::exec::unstable::IExecutionStateInfo::setNodeData(omni::core::TypeId desiredType,
                                                                          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(key, desiredType, 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
}

#endif // DOXYGEN_BUILD

//! Calls either @ref omni::graph::exec::unstable::IExecutionContext::getNodeDataAs() or @ref
//! omni::graph::exec::unstable::IExecutionStateInfo::getNodeDataAs() (dependent on the type of the first argument).
//!
//! The purpose of this macro is generate an appropriate @ref omni::core::TypeId at compile time from the data item's
//! type.  The user can manually do this, but this macro is much less error prone.
//!
//! @code
//!    auto data = OMNI_GRAPH_EXEC_GET_NODE_DATA_AS( task->getContext(), GraphContextCacheOverride,
//!        task->getUpstreamPath(), nullptr, tokens::kInstanceContext).data();
//! @endcode
//!
//! The macro itself is a variadic macro and can map to multiple overloads of @c getNodeDataAs() methods in the
//! interface given as the first argument
//!
//! With newer compilers (GCC >= 8), this macro can be replaced with templated methods (without breaking the ABI).
#define OMNI_GRAPH_EXEC_GET_NODE_DATA_AS(context_, type_, ...)                                                         \
    context_->getNodeDataAs<type_>(CARB_HASH_STRING(CARB_STRINGIFY(type_)), __VA_ARGS__)
// the ugly macro above is used to hash the type of the data at compile time.
//
// it's possible to get the type of data at compile time by inspecting the function name (e.g. __FUNCSIG__ and
// __PRETTY_FUNCTION__).  however __PRETTY_FUNCTION__ was not a constexpr until GCC 8.  omniverse currently uses GCC 7
// so were left with this hack.
//
// https://gcc.gnu.org/bugzilla/show_bug.cgi?id=66639

//! Calls either @ref omni::graph::exec::unstable::IExecutionContext::setNodeData() or @ref
//! omni::graph::exec::unstable::IExecutionStateInfo::setNodeData() (dependent on the type of the first argument).
//!
//! The purpose of this macro is generate an appropriate @ref omni::core::TypeId at compile time from the data item's
//! type.  The user can manually do this, but this macro is much less error prone.
//!
//! @code
//! OMNI_GRAPH_EXEC_SET_NODE_DATA(stateInfo, GraphContextCacheOverride, tokens::kInstanceContext,
//! std::move(contextOverridePtr));
//! @endcode
//!
//! The macro itself is a variadic macro and can map to multiple overloads of @c setNodeData() methods in the interface
//! given as the first argument
//!
//! With newer compilers (GCC >= 8), this macro can be replaced with templated methods (without breaking the ABI).
#define OMNI_GRAPH_EXEC_SET_NODE_DATA(context_, type_, ...)                                                            \
    context_->setNodeData<type_>(CARB_HASH_STRING(CARB_STRINGIFY(type_)), __VA_ARGS__)
// the ugly macro above is used to hash the type of the data at compile time.
//
// it's possible to get the type of data at compile time by inspecting the function name (e.g. __FUNCSIG__ and
// __PRETTY_FUNCTION__).  however __PRETTY_FUNCTION__ was not a constexpr until GCC 8.  omniverse currently uses GCC 7
// so were left with this hack.
//
// https://gcc.gnu.org/bugzilla/show_bug.cgi?id=66639

// generated API implementation
#define OMNI_BIND_INCLUDE_INTERFACE_IMPL
#include <omni/graph/exec/unstable/IExecutionStateInfo.gen.h>
