#pragma once

#include <concepts>
#include <condition_variable>
#include <coroutine>
#include <cstddef>
#include <exception>
#include <expected>
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <queue>
#include <thread>
#include <type_traits>
#include <utility>
#include <vector>

#include <mw/error.hpp>

namespace agent_smith::internal
{

template <typename T> struct ExpectedResultType
{
    using Type = mw::E<T>;
};

template <typename T> struct ExpectedResultType<mw::E<T>>
{
    using Type = mw::E<T>;
};

/// The flattened expected result produced by a blocking callable.
template <typename T>
using ExpectedResult =
    typename ExpectedResultType<std::remove_cvref_t<T>>::Type;

/// Private FIFO executor for work that cannot run on the agent coroutine.
class BlockingPool
{
public:
    explicit BlockingPool(std::size_t size);
    ~BlockingPool();

    BlockingPool(const BlockingPool&) = delete;
    BlockingPool& operator=(const BlockingPool&) = delete;

    template <typename F> auto run(F&& function);

private:
    struct WorkItem
    {
        std::function<void()> run;
        std::function<void()> cancel;
    };

    bool enqueue(WorkItem item);
    void worker(std::stop_token stop);

    mutable std::mutex mutex_;
    std::condition_variable condition_;
    std::queue<WorkItem> pending_;
    bool stopping_ = false;
    std::vector<std::jthread> workers_;
};

template <typename Result, typename Function> class BlockingWork
{
public:
    explicit BlockingWork(Function&& function) : function_(std::move(function))
    {
    }

    void run()
    {
        Result result = invoke();
        complete(std::move(result));
    }

    void cancel()
    {
        complete(std::unexpected(mw::runtimeError("blocking pool stopped")));
    }

    bool complete() const
    {
        std::lock_guard lock(mutex_);
        return complete_;
    }

    bool suspend(std::coroutine_handle<> continuation)
    {
        std::lock_guard lock(mutex_);
        if(complete_)
        {
            return false;
        }

        continuation_ = continuation;
        return true;
    }

    Result result()
    {
        std::lock_guard lock(mutex_);
        return std::move(*result_);
    }

private:
    Result invoke()
    {
        try
        {
            using RawResult = std::invoke_result_t<Function&>;
            if constexpr(std::is_void_v<RawResult>)
            {
                function_();
                return {};
            }
            else if constexpr(std::same_as<std::remove_cvref_t<RawResult>,
                                           Result>)
            {
                return function_();
            }
            else
            {
                return Result(function_());
            }
        }
        catch(const std::exception& error)
        {
            return std::unexpected(mw::runtimeError(error.what()));
        }
        catch(...)
        {
            return std::unexpected(mw::runtimeError(
                "blocking callable threw an unknown exception"));
        }
    }

    void complete(Result result)
    {
        std::coroutine_handle<> continuation;
        {
            std::lock_guard lock(mutex_);
            if(complete_)
            {
                return;
            }

            result_.emplace(std::move(result));
            complete_ = true;
            continuation = continuation_;
        }

        if(continuation != nullptr)
        {
            continuation.resume();
        }
    }

    mutable std::mutex mutex_;
    Function function_;
    std::optional<Result> result_;
    std::coroutine_handle<> continuation_ = nullptr;
    bool complete_ = false;
};

template <typename Result, typename Function> class BlockingAwaiter
{
public:
    explicit BlockingAwaiter(
        std::shared_ptr<BlockingWork<Result, Function>> work)
        : work_(std::move(work))
    {
    }

    bool await_ready() const
    {
        return work_->complete();
    }

    bool await_suspend(std::coroutine_handle<> continuation)
    {
        return work_->suspend(continuation);
    }

    Result await_resume()
    {
        return work_->result();
    }

private:
    std::shared_ptr<BlockingWork<Result, Function>> work_;
};

template <typename F> auto BlockingPool::run(F&& function)
{
    using Function = std::decay_t<F>;
    using Result = ExpectedResult<std::invoke_result_t<Function&>>;
    using Work = BlockingWork<Result, Function>;

    auto work = std::make_shared<Work>(std::forward<F>(function));
    WorkItem item{
        .run = [work]() { work->run(); },
        .cancel = [work]() { work->cancel(); },
    };

    if(!enqueue(std::move(item)))
    {
        work->cancel();
    }

    return BlockingAwaiter<Result, Function>(std::move(work));
}

} // namespace agent_smith::internal