#pragma once

#include <condition_variable>
#include <coroutine>
#include <exception>
#include <mutex>
#include <optional>
#include <utility>

#include <mw/task.hpp>

namespace agent_smith
{
namespace internal
{

template <typename T> class SyncWaitTask
{
public:
    struct promise_type
    {
        SyncWaitTask get_return_object()
        {
            return SyncWaitTask(
                std::coroutine_handle<promise_type>::from_promise(*this));
        }

        std::suspend_always initial_suspend() noexcept
        {
            return {};
        }

        auto final_suspend() noexcept
        {
            struct FinalAwaiter
            {
                bool await_ready() noexcept
                {
                    return false;
                }

                void await_suspend(
                    std::coroutine_handle<promise_type> handle) noexcept
                {
                    auto& promise = handle.promise();
                    std::lock_guard lock(promise.mutex);
                    promise.done = true;
                    promise.condition.notify_one();
                }

                void await_resume() noexcept {}
            };

            return FinalAwaiter{};
        }

        template <typename U> void return_value(U&& value)
        {
            value_.emplace(std::forward<U>(value));
        }

        void unhandled_exception()
        {
            exception_ = std::current_exception();
        }

        std::mutex mutex;
        std::condition_variable condition;
        bool done = false;
        std::optional<T> value_;
        std::exception_ptr exception_;
    };

    explicit SyncWaitTask(std::coroutine_handle<promise_type> handle)
        : handle_(handle)
    {
    }

    SyncWaitTask(SyncWaitTask&& other) noexcept
        : handle_(std::exchange(other.handle_, nullptr))
    {
    }

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

    ~SyncWaitTask()
    {
        if(handle_ != nullptr)
        {
            handle_.destroy();
        }
    }

    void start()
    {
        handle_.resume();
    }

    T get()
    {
        auto& promise = handle_.promise();
        std::unique_lock lock(promise.mutex);
        promise.condition.wait(lock, [&]() { return promise.done; });
        lock.unlock();

        if(promise.exception_ != nullptr)
        {
            std::rethrow_exception(promise.exception_);
        }
        return std::move(*promise.value_);
    }

private:
    std::coroutine_handle<promise_type> handle_ = nullptr;
};

template <typename T> SyncWaitTask<T> makeSyncWaiter(mw::Task<T> task)
{
    co_return co_await std::move(task);
}

} // namespace internal

/// Start a task and block this thread until its result is available.
///
/// Use only at a synchronous application boundary or in tests. Do not call
/// this from a thread that must run work required by the task to complete.
template <typename T> T syncWait(mw::Task<T> task)
{
    auto waiter = internal::makeSyncWaiter(std::move(task));
    waiter.start();
    return waiter.get();
}

} // namespace agent_smith