#pragma once

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

#include <mw/error.hpp>
#include <mw/task.hpp>

namespace agent_smith::internal
{

/// An internal eager wrapper around one mw::Task<mw::E<T>> operation.
template <typename T> class StartedTask
{
public:
    struct promise_type
    {
        StartedTask get_return_object()
        {
            return StartedTask(
                std::coroutine_handle<promise_type>::from_promise(*this));
        }

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

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

                std::coroutine_handle<> await_suspend(
                    std::coroutine_handle<promise_type> handle) noexcept
                {
                    std::lock_guard lock(handle.promise().mutex);
                    handle.promise().complete = true;
                    if(handle.promise().continuation != nullptr)
                    {
                        return handle.promise().continuation;
                    }
                    return std::noop_coroutine();
                }

                void await_resume() noexcept {}
            };

            return FinalAwaiter{};
        }

        void return_value(mw::E<T> value)
        {
            value_.emplace(std::move(value));
        }

        void unhandled_exception()
        {
            try
            {
                throw;
            }
            catch(const std::exception& error)
            {
                value_.emplace(std::unexpected(mw::runtimeError(error.what())));
            }
            catch(...)
            {
                value_.emplace(std::unexpected(mw::runtimeError(
                    "started task threw an unknown exception")));
            }
        }

        std::mutex mutex;
        std::optional<mw::E<T>> value_;
        std::coroutine_handle<> continuation = nullptr;
        bool complete = false;
    };

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

    StartedTask& operator=(StartedTask&& other) noexcept
    {
        if(this == &other)
        {
            return *this;
        }

        destroy();
        handle_ = std::exchange(other.handle_, nullptr);
        return *this;
    }

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

    ~StartedTask()
    {
        destroy();
    }

    /// Start a lazy task immediately and retain its result for later awaiting.
    static StartedTask start(mw::Task<mw::E<T>> task)
    {
        co_return co_await std::move(task);
    }

    /// Check completion without consuming the result.
    bool complete() const
    {
        std::lock_guard lock(handle_.promise().mutex);
        return handle_.promise().complete;
    }

    auto operator co_await() &&
    {
        struct Awaiter
        {
            std::coroutine_handle<promise_type> handle;

            bool await_ready()
            {
                std::lock_guard lock(handle.promise().mutex);
                return handle.promise().complete;
            }

            bool await_suspend(std::coroutine_handle<> continuation)
            {
                std::lock_guard lock(handle.promise().mutex);
                if(handle.promise().complete)
                {
                    return false;
                }

                handle.promise().continuation = continuation;
                return true;
            }

            mw::E<T> await_resume()
            {
                mw::E<T> result = [&]()
                {
                    std::lock_guard lock(handle.promise().mutex);
                    return std::move(*handle.promise().value_);
                }();
                handle.destroy();
                return result;
            }
        };

        return Awaiter{std::exchange(handle_, nullptr)};
    }

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

    void destroy()
    {
        if(handle_ == nullptr)
        {
            return;
        }

        if(!complete())
        {
            std::terminate();
        }

        handle_.destroy();
        handle_ = nullptr;
    }

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

} // namespace agent_smith::internal