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

#include <omni/graph/exec/unstable/Assert.h>
#include <omni/graph/exec/unstable/IBackgroundTask.h>

#include <future>

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

//! @copydoc omni::graph::exec::unstable::IBackgroundTask
class BackgroundTask : public Implements<IBackgroundTask>
{
public:
    //! Creates a new @ref BackgroundTask.
    //!
    //! May throw
    static omni::core::ObjectPtr<BackgroundTask> create()
    {
        return omni::core::steal(new BackgroundTask);
    }

protected:
    //! Allows access to result of an async operation.
    using Future = std::future<omni::core::ObjectPtr<IBackgroundResultWriter>>;

    //! Allows setting the result of an async operation.
    using Promise = std::promise<omni::core::ObjectPtr<IBackgroundResultWriter>>;

    //! @copydoc IBackgroundTask_abi::getBackgroundResult_abi
    omni::core::Result getBackgroundResult_abi(IBackgroundResult** out) noexcept override
    {
        class Result : public Implements<IBackgroundResult>
        {
        public:
            Result(Future&& future) : m_future(std::move(future))
            {
            }

        protected:
            omni::core::Result isReady_abi(bool* out) noexcept override
            {
                try
                {
                    *out = (m_future.wait_for(std::chrono::milliseconds(0)) == std::future_status::ready); // may throw
                    return omni::core::kResultSuccess;
                }
                OMNI_GRAPH_EXEC_CATCH_ABI_EXCEPTION()
            }

            omni::core::Result cancel_abi(bool blocking) noexcept override
            {
                try
                {
                    if (blocking)
                    {
                        m_future.wait(); // may throw
                    }
                    return omni::core::kResultSuccess;
                }
                OMNI_GRAPH_EXEC_CATCH_ABI_EXCEPTION()
            }

            omni::core::Result write_abi(ExecutionTask* info, Status* out) noexcept override
            {
                try
                {
                    *out = Status::eUnknown;
                    *out = m_future.get()->write(*info); // may throw
                    return omni::core::kResultSuccess;
                }
                OMNI_GRAPH_EXEC_CATCH_ABI_EXCEPTION()
            }

            omni::core::Result waitFor_abi(uint64_t nanoseconds, BackgroundResultStatus* out) noexcept override
            {
                try
                {
                    auto result = m_future.wait_for(std::chrono::nanoseconds(nanoseconds)); // may throw
                    switch (result)
                    {
                    case std::future_status::deferred: // ?
                    case std::future_status::ready:
                        *out = BackgroundResultStatus::eReady;
                        break;
                    case std::future_status::timeout:
                        *out = BackgroundResultStatus::eTimeout;
                        break;
                    default:
                        throw std::logic_error("unknown future state");
                    }
                    return omni::core::kResultSuccess;
                }
                OMNI_GRAPH_EXEC_CATCH_ABI_EXCEPTION()
            }

        private:
            Future m_future;
        };

        try
        {
            *out = new Result(m_promise.get_future());
            return omni::core::kResultSuccess;
        }
        OMNI_GRAPH_EXEC_CATCH_ABI_EXCEPTION()
    }

    //! @copydoc IBackgroundTask_abi::setResultWriter_abi
    omni::core::Result setResultWriter_abi(IBackgroundResultWriter* writer) noexcept override
    {
        try
        {
            m_promise.set_value(omni::core::borrow(writer));
            return omni::core::kResultSuccess;
        }
        OMNI_GRAPH_EXEC_CATCH_ABI_EXCEPTION()
    }

private:
    Promise m_promise;
};

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