//
// MIT License
// Copyright (c) 2020 Jonathan R. Madsen
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED
// "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT
// LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
// PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
// HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
// ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
// WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
//
//
// ---------------------------------------------------------------
// Tasking class header file
//
// Class Description:
//
// This file creates the a class for handling a group of tasks that
// can be independently joined
//
// ---------------------------------------------------------------
// Author: Jonathan Madsen (Feb 13th 2018)
// ---------------------------------------------------------------

#include "PTL/Types.hh"
#include "PTL/Utility.hh"

namespace PTL
{
template <typename Tp, typename Arg, intmax_t MaxDepth>
template <typename Func>
TaskGroup<Tp, Arg, MaxDepth>::TaskGroup(Func&& _join, ThreadPool* _tp)
: m_join{ std::forward<Func>(_join) }
, m_pool{ _tp }
{
    internal_update();
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
template <typename Up>
TaskGroup<Tp, Arg, MaxDepth>::TaskGroup(ThreadPool* _tp,
                                        enable_if_t<std::is_void<Up>::value, int>)
: m_join{ []() {} }
, m_pool{ _tp }
{
    internal_update();
}

// Destructor
template <typename Tp, typename Arg, intmax_t MaxDepth>
TaskGroup<Tp, Arg, MaxDepth>::~TaskGroup()
{
    {
        // task will decrement counter and then acquire the lock to notify
        // condition variable so acquiring lock here will prevent the
        // task group from being destroyed before this is completed
        AutoLock _lk{ m_task_lock, std::defer_lock };
        if(!_lk.owns_lock())
            _lk.lock();
    }

    if(m_tbb_task_group)
    {
        auto* _arena = m_pool->get_task_arena();
        _arena->execute([this]() { this->m_tbb_task_group->wait(); });
    }
    delete m_tbb_task_group;
    this->clear();
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
template <typename Up>
std::shared_ptr<Up>
TaskGroup<Tp, Arg, MaxDepth>::operator+=(std::shared_ptr<Up>&& _task)
{
    // thread-safe increment of tasks in task group
    operator++();
    // copy the shared pointer to abstract instance
    m_task_list.push_back(_task);
    // return the derived instance
    return std::move(_task);
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
void
TaskGroup<Tp, Arg, MaxDepth>::wait()
{
    auto _dtor = ScopeDestructor{ [&]() {
        if(m_tbb_task_group)
        {
            auto* _arena = m_pool->get_task_arena();
            _arena->execute([this]() { this->m_tbb_task_group->wait(); });
        }
    } };

    ThreadData* data = ThreadData::GetInstance();
    if(!data)
        return;

    // if no pool was initially present at creation
    if(!m_pool)
    {
        // check for master MT run-manager
        m_pool = internal::get_default_threadpool();

        // if no thread pool created
        if(!m_pool)
        {
            if(f_verbose > 0)
            {
                fprintf(stderr, "%s @ %i :: Warning! nullptr to thread-pool (%p)\n",
                        __FUNCTION__, __LINE__, static_cast<void*>(m_pool));
                std::cerr << __FUNCTION__ << "@" << __LINE__ << " :: Warning! "
                          << "nullptr to thread pool!" << std::endl;
            }
            return;
        }
    }

    ThreadPool*     tpool = (m_pool) ? m_pool : data->thread_pool;
    VUserTaskQueue* taskq = (tpool) ? tpool->get_queue() : data->current_queue;

    bool _is_main     = data->is_main;
    bool _within_task = data->within_task;

    auto is_active_state = [&]() {
        return (tpool->state()->load(std::memory_order_relaxed) !=
                thread_pool::state::STOPPED);
    };

    auto execute_this_threads_tasks = [&]() {
        if(!taskq)
            return;

        // only want to process if within a task
        if((!_is_main || tpool->size() < 2) && _within_task)
        {
            int bin = static_cast<int>(taskq->GetThreadBin());
            // const auto nitr = (tpool) ? tpool->size() :
            // Thread::hardware_concurrency();
            while(this->pending() > 0)
            {
                if(!taskq->empty())
                {
                    auto _task = taskq->GetTask(bin);
                    if(_task)
                        (*_task)();
                }
            }
        }
    };

    // checks for validity
    if(!is_native_task_group())
    {
        // for external threads
        if(!_is_main || tpool->size() < 2)
            return;
    }
    else if(f_verbose > 0)
    {
        if(!tpool || !taskq)
        {
            // something is wrong, didn't create thread-pool?
            fprintf(stderr,
                    "%s @ %i :: Warning! nullptr to thread data (%p) or task-queue "
                    "(%p)\n",
                    __FUNCTION__, __LINE__, static_cast<void*>(tpool),
                    static_cast<void*>(taskq));
        }
        // return if thread pool isn't built
        else if(is_native_task_group() && !tpool->is_alive())
        {
            fprintf(stderr, "%s @ %i :: Warning! thread-pool is not alive!\n",
                    __FUNCTION__, __LINE__);
        }
        else if(!is_active_state())
        {
            fprintf(stderr, "%s @ %i :: Warning! thread-pool is not active!\n",
                    __FUNCTION__, __LINE__);
        }
    }

    intmax_t wake_size = 2;
    AutoLock _lock(m_task_lock, std::defer_lock);

    while(is_active_state())
    {
        execute_this_threads_tasks();

        // while loop protects against spurious wake-ups
        while(_is_main && pending() > 0 && is_active_state())
        {
            // auto _wake = [&]() { return (wake_size > pending() ||
            // !is_active_state());
            // };

            // lock before sleeping on condition
            if(!_lock.owns_lock())
                _lock.lock();

            // Wait until signaled that a task has been competed
            // Unlock mutex while wait, then lock it back when signaled
            // when true, this wakes the thread
            if(pending() >= wake_size)
            {
                m_task_cond.wait(_lock);
            }
            else
            {
                m_task_cond.wait_for(_lock, std::chrono::microseconds(100));
            }
            // unlock
            if(_lock.owns_lock())
                _lock.unlock();
        }

        // if pending is not greater than zero, we are joined
        if(pending() <= 0)
            break;
    }

    if(_lock.owns_lock())
        _lock.unlock();

    intmax_t ntask = this->task_count().load();
    if(ntask > 0)
    {
        std::stringstream ss;
        ss << "\nWarning! Join operation issue! " << ntask << " tasks still "
           << "are running!" << std::endl;
        std::cerr << ss.str();
        this->wait();
    }
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
ScopeDestructor
TaskGroup<Tp, Arg, MaxDepth>::get_scope_destructor()
{
    auto& _counter   = m_tot_task_count;
    auto& _task_cond = task_cond();
    auto& _task_lock = task_lock();
    return ScopeDestructor{ [&_task_cond, &_task_lock, &_counter]() {
        auto _count = --(_counter);
        if(_count < 1)
        {
            AutoLock _lk{ _task_lock };
            _task_cond.notify_all();
        }
    } };
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
void
TaskGroup<Tp, Arg, MaxDepth>::notify()
{
    AutoLock _lk{ m_task_lock };
    m_task_cond.notify_one();
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
void
TaskGroup<Tp, Arg, MaxDepth>::notify_all()
{
    AutoLock _lk{ m_task_lock };
    m_task_cond.notify_all();
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
template <typename Func, typename... Args, typename Up>
enable_if_t<std::is_void<Up>::value, void>
TaskGroup<Tp, Arg, MaxDepth>::exec(Func func, Args... args)
{
    if(MaxDepth > 0 && !m_tbb_task_group && ThreadData::GetInstance() &&
       ThreadData::GetInstance()->task_depth > MaxDepth)
    {
        local_exec<Tp>(std::move(func), std::move(args)...);
    }
    else
    {
        auto& _counter   = m_tot_task_count;
        auto& _task_cond = task_cond();
        auto& _task_lock = task_lock();
        auto  _task      = wrap([&_task_cond, &_task_lock, &_counter, func, args...]() {
            auto _tdata = ThreadData::GetInstance();
            if(_tdata)
                ++(_tdata->task_depth);
            func(args...);
            auto _count = --(_counter);
            if(_tdata)
                --(_tdata->task_depth);
            if(_count < 1)
            {
                AutoLock _lk{ _task_lock };
                _task_cond.notify_all();
            }
        });

        if(m_tbb_task_group)
        {
            auto* _arena          = m_pool->get_task_arena();
            auto* _tbb_task_group = m_tbb_task_group;
            auto* _ptask          = _task.get();
            _arena->execute([_tbb_task_group, _ptask]() {
                _tbb_task_group->run([_ptask]() { (*_ptask)(); });
            });
        }
        else
        {
            m_pool->add_task(std::move(_task));
        }
    }
}
template <typename Tp, typename Arg, intmax_t MaxDepth>
template <typename Func, typename... Args, typename Up>
enable_if_t<!std::is_void<Up>::value, void>
TaskGroup<Tp, Arg, MaxDepth>::exec(Func func, Args... args)
{
    if(MaxDepth > 0 && !m_tbb_task_group && ThreadData::GetInstance() &&
       ThreadData::GetInstance()->task_depth > MaxDepth)
    {
        local_exec<Tp>(std::move(func), std::move(args)...);
    }
    else
    {
        auto& _counter   = m_tot_task_count;
        auto& _task_cond = task_cond();
        auto& _task_lock = task_lock();
        auto  _task      = wrap([&_task_cond, &_task_lock, &_counter, func, args...]() {
            auto _tdata = ThreadData::GetInstance();
            if(_tdata)
                ++(_tdata->task_depth);
            auto&& _ret   = func(args...);
            auto   _count = --(_counter);
            if(_tdata)
                --(_tdata->task_depth);
            if(_count < 1)
            {
                AutoLock _lk{ _task_lock };
                _task_cond.notify_all();
            }
            return std::forward<decltype(_ret)>(_ret);
        });

        if(m_tbb_task_group)
        {
            auto* _arena          = m_pool->get_task_arena();
            auto* _tbb_task_group = m_tbb_task_group;
            auto* _ptask          = _task.get();
            _arena->execute([_tbb_task_group, _ptask]() {
                _tbb_task_group->run([_ptask]() { (*_ptask)(); });
            });
        }
        else
        {
            m_pool->add_task(std::move(_task));
        }
    }
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
template <typename Up, typename Func, typename... Args>
enable_if_t<std::is_void<Up>::value, void>
TaskGroup<Tp, Arg, MaxDepth>::local_exec(Func func, Args... args)
{
    auto _tdata = ThreadData::GetInstance();
    if(_tdata)
        ++(_tdata->task_depth);
    promise_type _p{};
    m_future_list.emplace_back(_p.get_future());
    func(args...);
    _p.set_value();
    if(_tdata)
        --(_tdata->task_depth);
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
template <typename Up, typename Func, typename... Args>
enable_if_t<!std::is_void<Up>::value, void>
TaskGroup<Tp, Arg, MaxDepth>::local_exec(Func func, Args... args)
{
    auto _tdata = ThreadData::GetInstance();
    if(_tdata)
        ++(_tdata->task_depth);
    promise_type _p{};
    m_future_list.emplace_back(_p.get_future());
    _p.set_value(func(args...));
    if(_tdata)
        --(_tdata->task_depth);
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
template <typename Up, enable_if_t<!std::is_void<Up>::value, int>>
inline Up
TaskGroup<Tp, Arg, MaxDepth>::join(Up accum)
{
    this->wait();
    for(auto& itr : m_task_list)
    {
        using RetT = decay_t<decltype(itr->get())>;
        accum      = std::move(m_join(std::ref(accum), std::forward<RetT>(itr->get())));
    }
    for(auto& itr : m_future_list)
    {
        using RetT = decay_t<decltype(itr.get())>;
        accum      = std::move(m_join(std::ref(accum), std::forward<RetT>(itr.get())));
    }
    this->clear();
    return accum;
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
template <typename Up, typename Rp,
          enable_if_t<std::is_void<Up>::value && std::is_void<Rp>::value, int>>
inline void
TaskGroup<Tp, Arg, MaxDepth>::join()
{
    this->wait();
    for(auto& itr : m_task_list)
        itr->get();
    for(auto& itr : m_future_list)
        itr.get();
    m_join();
    this->clear();
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
template <typename Up, typename Rp,
          enable_if_t<std::is_void<Up>::value && !std::is_void<Rp>::value, int>>
inline void
TaskGroup<Tp, Arg, MaxDepth>::join()
{
    this->wait();
    for(auto& itr : m_task_list)
    {
        using RetT = decay_t<decltype(itr->get())>;
        m_join(std::forward<RetT>(itr->get()));
    }
    for(auto& itr : m_future_list)
    {
        using RetT = decay_t<decltype(itr.get())>;
        m_join(std::forward<RetT>(itr.get()));
    }
    this->clear();
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
void
TaskGroup<Tp, Arg, MaxDepth>::clear()
{
    m_future_list.clear();
    m_task_list.clear();
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
void
TaskGroup<Tp, Arg, MaxDepth>::internal_update()
{
    if(!m_pool)
        m_pool = internal::get_default_threadpool();

    if(!m_pool)
    {
        std::stringstream ss{};
        ss << "[TaskGroup]> " << __FUNCTION__ << "@" << __LINE__
           << " :: nullptr to thread pool";
        throw std::runtime_error(ss.str());
    }

    if(m_pool->is_tbb_threadpool())
    {
        m_tbb_task_group = new tbb_task_group_t{};
    }
}

template <typename Tp, typename Arg, intmax_t MaxDepth>
int TaskGroup<Tp, Arg, MaxDepth>::f_verbose = GetEnv<int>("PTL_VERBOSE", 0);

}  // namespace PTL