// 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 ParallelScheduler.h
//!
//! @brief Defines @ref omni::kit::exec::core::unstable::ParallelSpawner.
#pragma once

#include <omni/graph/exec/unstable/AtomicBackoff.h>
#include <omni/kit/exec/core/unstable/ITbbSchedulerState.h>
#include <omni/kit/exec/core/unstable/TbbSchedulerState.h>

#ifndef __TBB_ALLOW_MUTABLE_FUNCTORS
//! Allow task functors to mutate captured by value state
#    define __TBB_ALLOW_MUTABLE_FUNCTORS 1
#endif

#include <tbb/task.h>
#include <tbb/task_group.h>

#include <atomic>
#include <thread>

namespace omni
{
namespace kit
{
namespace exec
{
namespace core
{
namespace unstable
{

#ifndef DOXYGEN_BUILD

//! Parallel scheduler for task dispatch.
//!
//! Currently only needs to handle special requests, actual parallel work is enqueue directly to TBB
class ParallelScheduler
{
public:
    //! Access static per-DLL singleton.
    static ParallelScheduler& getSingleton()
    {
        // this static is per-DLL, but the actual private data (s_state) is shared across DLLs
        static ParallelScheduler sScheduler;
        return sScheduler;
    }

    //! Class wrapping lambda into serial and parallel tasks. Additionally, this class keeps a track of all
    //! created tasks to allow activating isolation pass when all serial and parallel work is completed.
    template <typename Fn>
    class TaskBody : public tbb::task
    {
    public:
        //! Expected function signature is void(). Increments total task count.
        TaskBody(Fn&& f) : m_fn(std::forward<Fn>(f)), m_sharedState(ParallelScheduler::getSingleton().s_state)
        {
            ++m_sharedState->totalTasksInFlight;
        }
        //! Decrements total task count and starts processing isolate tasks if no more tasks are running.
        ~TaskBody()
        {
            --m_sharedState->totalTasksInFlight;
        }

        //! TBB task execute method calling into lambda. To track nested executions, we increment and decrement
        //! tasks per thread atomic counter.
        task* execute() override
        {
            struct ExecutingCountedTaskScope
            {
                ExecutingCountedTaskScope(tbb::enumerable_thread_specific<int>& executingPerThread)
                    : _executingPerThread(executingPerThread)
                {
                    _executingPerThread.local()++;
                }
                ~ExecutingCountedTaskScope()
                {
                    _executingPerThread.local()--;
                }

                tbb::enumerable_thread_specific<int>& _executingPerThread;
            } executingScope(m_sharedState->tasksExecutingPerThread);

            return m_fn(); // allow for potential continuation
        };

    private:
        Fn m_fn; //!< Lambda function to call from the execute() function.
        TbbSchedulerState* m_sharedState{ nullptr }; //!< Scheduler shared state used to track total task count.
    };

    //! Class wrapping lambda into an isolated TBB task. No other tasks will be running when this task executes.
    template <typename Fn>
    class IsolateTaskBody : public tbb::task
    {
    public:
        //! Constructor capturing the lambda function
        //!
        //! Expected function signature is void()
        IsolateTaskBody(Fn&& f) : m_fn(std::forward<Fn>(f))
        {
        }

        //! TBB task execute method calling into lambda.
        task* execute() override
        {
            // isolate tasks are not counted in total in flight tasks to allow us detect when
            // all other work is finished. Once we are entering scope of isolate task, we make sure
            // PauseTaskScope doesn't consider this task as taking part in totalTasksInFlight
            struct ExecutingNotCountedTaskScope
            {
                ExecutingNotCountedTaskScope(tbb::enumerable_thread_specific<int>& executingPerThread)
                    : _executingPerThread(executingPerThread)
                {
                    auto& perThreadCount = _executingPerThread.local();
                    _originalPerThread = perThreadCount;
                    perThreadCount = 0;
                }
                ~ExecutingNotCountedTaskScope()
                {
                    _executingPerThread.local() = _originalPerThread;
                }

                tbb::enumerable_thread_specific<int>& _executingPerThread;
                int _originalPerThread;
            } executingScope(ParallelScheduler::getSingleton().s_state->tasksExecutingPerThread);

            return m_fn(); // allow for potential continuation
        };

    private:
        Fn m_fn; //!< Lambda function to call from the execute() function
    };

    //! RAII class used to pause and resume total tasks in flight when thread is about to pick up new tasks to work on.
    class PauseTaskScope
    {
    public:
        PauseTaskScope()
            : m_scheduler(ParallelScheduler::getSingleton()),
              m_isCountedTask(m_scheduler.s_state->tasksExecutingPerThread.local() > 0)
        {
            if (m_isCountedTask)
            {
                --m_scheduler.s_state->totalTasksInFlight;
            }
        }
        ~PauseTaskScope()
        {
            if (m_isCountedTask)
            {
                ++m_scheduler.s_state->totalTasksInFlight;
            }
        }

    private:
        ParallelScheduler& m_scheduler; //!< To avoid calling into getSingleton multiple times
        const bool m_isCountedTask; //!< Know if we are pausing task that was scheduled by us or someone else.
    };

    //! Typically parallel tasks are spawned right away, but when executed from isolation tasks, we defer spawning
    //! until all isolation work is completed.
    void pushParallelTask(tbb::task& t, bool isExecutingThread)
    {
        if (!isProcessingIsolate())
        {
            if (isExecutingThread && (s_state->totalTasksInFlight < 2))
            {
                // the executing thread never joins the arena, so won't process its local queue. add the task to the
                // global queue.  we only do this if there aren't already tasks in flight that can pick up task if we
                // spawn() it.
                tbb::task::enqueue(t);
            }
            else
            {
                tbb::task::spawn(t); // add to local queue
            }
        }
        else
        {
            s_state->queueParallel.push(&t);
        }
    }

    //! Add task /p t for serial dispatch. This tasks are guaranteed to be run one after another, but other thread-safe
    //! tasks are allowed to run together with them
    void pushSerialTask(tbb::task& t)
    {
        s_state->queueSerial.push(&t);
    }

    //! Store isolate task for execution when no other tasks are running.
    void pushIsolateTask(tbb::task& t)
    {
        s_state->queueIsolate.push(&t);
    }

    //! Called to notify completion of a serial task and acquire more work if available.
    tbb::task* completedSerialTask()
    {
        tbb::task* serialTask = nullptr;
        s_state->queueSerial.try_pop(serialTask);
        return serialTask;
    }

    //! Called to notify completion of an isolate task and acquire more work if available.
    tbb::task* completedIsolateTask()
    {
        tbb::task* isolateTask = nullptr;
        s_state->queueIsolate.try_pop(isolateTask);
        return isolateTask;
    }

    //! Returns true if calling thread has been tagged as the one processing serial task. If so, this thread
    //! is in the middle of processing a serial task.
    bool isWithinSerialTask() const
    {
        return (s_state->serialThreadId.load() == std::this_thread::get_id());
    }

    //! Returns true if calling thread has been tagged as the one processing isolate task. If so, this thread
    //! is in the middle of processing a isolate task.
    bool isWithinIsolateTask() const
    {
        return (s_state->isolateThreadId.load() == std::this_thread::get_id());
    }

    //! Returns true if processing of isolate tasks is in progress.
    bool isProcessingIsolate() const
    {
        return (s_state->isolateThreadId.load() != std::thread::id());
    }

    void processContextThread(tbb::empty_task* rootTask)
    {
        // See the note in TbbSchedulerState.h as to why we need this
        // mutex lock here.
        std::lock_guard<tbb::recursive_mutex> lock(s_state->executingThreadMutex);

        graph::exec::unstable::AtomicBackoff backoff;

        if (isWithinIsolateTask())
        {
            while (rootTask->ref_count() > 1)
            {
                if (processQueuedTasks())
                    backoff.reset();
                else
                    backoff.pause();
            }
        }
        else
        {
            while (rootTask->ref_count() > 1)
            {
                if (processSerialTasks() || processIsolateTasks())
                    backoff.reset();
                else
                    backoff.pause();
            }
        }
    }

private:
    //! Explicit call to execute the serial task queueSerial until empty.
    //!
    bool processSerialTasks()
    {
        tbb::task* task = nullptr;
        s_state->queueSerial.try_pop(task);

        // No work
        if (!task)
            return false;

        {
            // Acquire the thread that's supposed to be handling serial task evaluation. Note
            // that this thread can be derived from many different execution contexts kickstarted
            // by multiple different threads.
            std::thread::id currentThreadId = std::this_thread::get_id();
            std::thread::id originalThreadId;
            if (!s_state->serialThreadId.compare_exchange_strong(originalThreadId, currentThreadId) &&
                originalThreadId != currentThreadId)
            {
                return false;
            }

            // Once serial tasks evaluation is complete, we will want to restore the serial thread ID
            // back to its default value; employ an RAII helper to do so.
            struct ScopeRelease
            {
                std::thread::id _originalThreadId;
                std::atomic<std::thread::id>& _serialThreadId;
                ~ScopeRelease()
                {
                    _serialThreadId = _originalThreadId;
                }
            } threadScope = { originalThreadId, s_state->serialThreadId };

            // Run the loop over tasks. We are responsible here to delete the task since it is consumed outside of TBB.
            while (task)
            {
                tbb::task* nextTask = task->execute();
                tbb::task::destroy(*task);
                task = nextTask;
            }
        }

        // Let the caller know that we had something to do
        return true;
    }

    //! When no more work is available execute isolated tasks.
    bool processIsolateTasks()
    {
        if (s_state->totalTasksInFlight > 0 || s_state->queueIsolate.empty())
        {
            return false;
        }

        // try acquire the right to process isolated tasks. Need to support nested executions.
        {
            std::thread::id currentThreadId = std::this_thread::get_id();
            std::thread::id originalThreadId;
            if (!s_state->isolateThreadId.compare_exchange_strong(originalThreadId, currentThreadId) &&
                originalThreadId != currentThreadId)
            {
                return false;
            }

            // we acquired the thread, nothing else will be running until the end of this scope
            struct ScopeRelease
            {
                std::thread::id _originalThreadId;
                std::atomic<std::thread::id>& _isolateThreadId;
                ~ScopeRelease()
                {
                    _isolateThreadId = _originalThreadId;
                }
            } threadScope = { originalThreadId, s_state->isolateThreadId };

            // we don't count isolate tasks, so just pick up the next one
            tbb::task* task = completedIsolateTask();

            // Run the loop over tasks. We are responsible here to delete the task since it is consumed outside of TBB.
            while (task)
            {
                tbb::task* nextTask = task->execute();
                tbb::task::destroy(*task);
                task = nextTask;
            }

            // here we will release this thread from processing isolate tasks.
            // We don't worry about synchronization between push tasks and this operation because
            // all push call can only happen from within above loop (indirectly via execute). There is no
            // other concurrent execution happening when we are in here.
        }

        // do NOT start parallel work in nested isolate task. it has to be all consumed on this thread until we can
        // exit isolation
        if (!isWithinIsolateTask())
        {
            // restart parallel task execution
            tbb::task* dispatchTask = nullptr;
            while (s_state->queueParallel.try_pop(dispatchTask))
            {
                if (s_state->totalTasksInFlight < 2)
                {
                    tbb::task::enqueue(*dispatchTask);
                }
                else
                {
                    tbb::task::spawn(*dispatchTask);
                }
            }
        }

        return true;
    }

    //! This processing should be avoided as much as possible but is necessary when work is generated and executed
    //! from within isolate task. We have to empty all the queues and unblock the rest of work.
    bool processQueuedTasks()
    {
        OMNI_ASSERT(isWithinIsolateTask());

        bool ret = false;
        tbb::task* task = nullptr;

        if (s_state->queueIsolate.try_pop(task))
        {
            while (task)
            {
                tbb::task* nextTask = task->execute();
                tbb::task::destroy(*task); // We are responsible here to delete the task since it is consumed manually
                task = nextTask;
            }

            ret |= true;
        }

        if (s_state->queueSerial.try_pop(task))
        {
            while (task)
            {
                tbb::task* nextTask = task->execute();
                tbb::task::destroy(*task); // We are responsible here to delete the task since it is consumed manually
                task = nextTask;
            }

            ret |= true;
        }

        if (s_state->queueParallel.try_pop(task))
        {
            do
            {
                task->execute();
                tbb::task::destroy(*task); // We are responsible here to delete the task since it is consumed manually
                // parallel tasks don't do scheduling continuation, this is why we don't pick up the next task
            } while (s_state->queueParallel.try_pop(task));

            ret |= true;
        }

        return ret;
    }

    //! Constructor
    explicit ParallelScheduler() noexcept
    {
        omni::core::ObjectPtr<ITbbSchedulerState> sInterface = omni::core::createType<ITbbSchedulerState>();
        OMNI_ASSERT(sInterface);
        s_state = sInterface->geState();
        OMNI_ASSERT(s_state);
    }

    TbbSchedulerState* s_state; //!< One scheduler state shared by all DLLs
};

namespace detail
{

//! Utility to construct in-place /p allocTask a task /p TaskT wrapping given lambda /p f
template <template <class...> typename TaskT, typename A, typename Fn>
TaskT<Fn>* makeTask(A&& allocTask, Fn&& f)
{
    return new (allocTask) TaskT<Fn>(std::forward<Fn>(f));
}

} // namespace detail

#endif // DOXYGEN_BUILD

//! Interface executors use to talk to the scheduler
struct ParallelSpawner
{
    //! Constructor
    ParallelSpawner(graph::exec::unstable::IExecutionContext* context) : m_context(context)
    {
        m_root = new (tbb::task::allocate_root()) tbb::empty_task;
        m_root->set_ref_count(1);
    }

    //! Destructor
    ~ParallelSpawner()
    {
        tbb::task::destroy(*m_root);
    }

    //! Thread-safe schedule method called by executor to enqueue generated work.
    //! Supports parallel, serial and isolate scheduling constraints.
    template <typename Fn>
    graph::exec::unstable::Status schedule(Fn&& task, graph::exec::unstable::SchedulingInfo schedInfo)
    {
        using namespace detail;

        if (schedInfo == graph::exec::unstable::SchedulingInfo::eParallel)
        {
            auto* dispatchTask = detail::makeTask<ParallelScheduler::TaskBody>(
                tbb::task::allocate_additional_child_of(*m_root),
                [task = graph::exec::unstable::captureScheduleFunction(task), this]() mutable
                {
                    graph::exec::unstable::Status ret = graph::exec::unstable::invokeScheduleFunction(task);
                    this->accumulateStatus(ret);

                    return nullptr;
                });

            ParallelScheduler::getSingleton().pushParallelTask(*dispatchTask, m_context->isExecutingThread());
        }
        else if (schedInfo == graph::exec::unstable::SchedulingInfo::eIsolate)
        {
            auto* dispatchTask = detail::makeTask<ParallelScheduler::IsolateTaskBody>(
                tbb::task::allocate_additional_child_of(*m_root),
                [task = graph::exec::unstable::captureScheduleFunction(task), this]() mutable
                {
                    graph::exec::unstable::Status ret = graph::exec::unstable::invokeScheduleFunction(task);
                    this->accumulateStatus(ret);

                    return ParallelScheduler::getSingleton().completedIsolateTask();
                });

            ParallelScheduler::getSingleton().pushIsolateTask(*dispatchTask);
        }
        else
        {
            auto* dispatchTask = detail::makeTask<ParallelScheduler::TaskBody>(
                tbb::task::allocate_additional_child_of(*m_root),
                [task = graph::exec::unstable::captureScheduleFunction(task), this]() mutable -> tbb::task*
                {
                    graph::exec::unstable::Status ret = graph::exec::unstable::invokeScheduleFunction(task);
                    this->accumulateStatus(ret);

                    return ParallelScheduler::getSingleton().completedSerialTask();
                });

            ParallelScheduler::getSingleton().pushSerialTask(*dispatchTask);
        }

        return graph::exec::unstable::Status::eSuccess;
    }

    //! Thread-safe accumulation of status returned by all spawned tasks
    void accumulateStatus(graph::exec::unstable::Status ret)
    {
        graph::exec::unstable::Status current, newValue = graph::exec::unstable::Status::eUnknown;
        do
        {
            current = m_status.load();
            newValue = ret | current;
        } while (!m_status.compare_exchange_weak(current, newValue));
    }

    //! Blocking call to acquire accumulated status.
    //! When more work is still available, this thread will join the arena and pickup some work OR
    //! will be reserved to process serial tasks if we end up here from a serial task (nested scenario).
    graph::exec::unstable::Status getStatus()
    {
        // We are about to enter nested execution. This has an effect on total tasks in flight, i.e.
        // we will suspend the current task by reducing the counters. All that is done by RAII class below.
        ParallelScheduler::PauseTaskScope pauseTask;

        // Note that in situations where multiple different contexts are running from multiple different
        // threads, just having the first check (m_context->isExecutingThread()) won't be enough because
        // we may be attempting to get the status of a serial/isolate task that was originally created in
        // context A running on thread 1 while context B running on thread 2 is being processed in the
        // below check; this can occur if context A/thread 1 was temporarily suspended in the past after
        // context B/thread 2 beat it to acquiring the s_state->executingThreadMutex, meaning that we
        // are currently processing nested tasks (which can be of any scheduling type) that are derived
        // from some top-level set of serial/isolate tasks in context B. In such situations, we need to
        // additionally check if we are currently within a serial or isolate task scope, since otherwise
        // the task originally created in context A/thread 1 will incorrectly skip processing on context B's
        // thread, despite thread 2 being the only running context thread at the moment, and take a different
        // code-path that leads to hangs. Serial/isolate tasks don't need to be run exclusively on the
        // context thread from which they were eventually scheduled - they can be run on any such context
        // thread as long as that context thread is the only one running/processing serial/isolate tasks.
        if (m_context->isExecutingThread() || ParallelScheduler::getSingleton().isWithinSerialTask() ||
            ParallelScheduler::getSingleton().isWithinIsolateTask())
        {
            ParallelScheduler::getSingleton().processContextThread(m_root);
        }
        else
        {
            if (m_root->ref_count() > 1)
            {
                m_root->wait_for_all();
                m_root->set_ref_count(1);
            }
        }

        return m_status;
    }

private:
    graph::exec::unstable::IExecutionContext* m_context; //!< Execution context used to identify executing thread
    tbb::empty_task* m_root{ nullptr }; //!< Root task for all spawned tasks.
    std::atomic<graph::exec::unstable::Status> m_status{ graph::exec::unstable::Status::eUnknown }; //!< Accumulated
                                                                                                    //!< status
};

} // namespace unable
} // namespace core
} // namespace exec
} // namespace kit
} // namespace omni
