#include "thread_pool.h"

#include <utility>

namespace
{

// Terminate before packaged_task can capture an unexpected exception.
void invokeTask(std::move_only_function<void()>& task) noexcept
{
    task();
}

}

ThreadPool::ThreadPool(std::size_t worker_count)
    : workers(worker_count)
{}

mw::E<std::unique_ptr<ThreadPool>> ThreadPool::create(std::size_t worker_count)
{
    if(worker_count == 0)
    {
        return std::unexpected(mw::runtimeError(
            "Thread pool capacity must be positive"));
    }
    auto pool = std::unique_ptr<ThreadPool>(new ThreadPool(worker_count));
    for(auto& worker : pool->workers)
    {
        worker.thread = std::thread(&ThreadPool::run, pool.get(),
                                    std::ref(worker));
    }
    return pool;
}

ThreadPool::~ThreadPool()
{
    stop();
}

void ThreadPool::stop()
{
    {
        std::lock_guard lock(mutex);
        stopping = true;
    }
    for(auto& worker : workers)
    {
        worker.work_ready.notify_one();
    }
    for(auto& worker : workers)
    {
        if(worker.thread.joinable())
        {
            worker.thread.join();
        }
    }
}

mw::E<std::optional<std::future<void>>> ThreadPool::trySubmit(
    std::move_only_function<void()> task)
{
    if(!task)
    {
        return std::unexpected(mw::runtimeError(
            "Thread pool task must not be empty"));
    }
    std::unique_lock lock(mutex);
    if(stopping)
    {
        return std::nullopt;
    }
    for(auto& worker : workers)
    {
        if(!worker.busy)
        {
            std::packaged_task<void()> packaged(
                std::bind_front(invokeTask, std::move(task)));
            auto result = packaged.get_future();
            worker.task = std::move(packaged);
            worker.busy = true;
            ++outstanding;
            lock.unlock();
            worker.work_ready.notify_one();
            return std::optional<std::future<void>>(std::move(result));
        }
    }
    return std::nullopt;
}

void ThreadPool::waitIdle()
{
    std::unique_lock lock(mutex);
    while(outstanding != 0)
    {
        idle.wait(lock);
    }
}

void ThreadPool::run(Worker& worker)
{
    while(true)
    {
        std::packaged_task<void()> task;
        {
            std::unique_lock lock(mutex);
            while(!worker.task.valid() && !stopping)
            {
                worker.work_ready.wait(lock);
            }
            if(!worker.task.valid())
            {
                return;
            }
            task = std::move(worker.task);
        }
        task();
        task = {};
        {
            std::lock_guard lock(mutex);
            worker.busy = false;
            --outstanding;
            if(outstanding == 0)
            {
                idle.notify_all();
            }
        }
    }
}