#pragma once

#include <coroutine>
#include <deque>
#include <memory>
#include <mutex>
#include <optional>
#include <utility>

#include <mw/task.hpp>

namespace agent_smith::internal
{

/// A closeable queue that resumes consumers when another thread pushes data.
template <typename T> class AsyncQueue
{
public:
    /// Add an item unless the queue has already been closed.
    void push(T value)
    {
        std::shared_ptr<PopState> waiter;
        std::coroutine_handle<> continuation;
        {
            std::lock_guard lock(mutex_);
            if(closed_)
            {
                return;
            }

            if(waiters_.empty())
            {
                values_.push_back(std::move(value));
                return;
            }

            waiter = std::move(waiters_.front());
            waiters_.pop_front();
            waiter->value.emplace(std::move(value));
            continuation = waiter->continuation;
        }

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

    /// Wait for one item, or nullopt after closure and queue drainage.
    mw::Task<std::optional<T>> pop()
    {
        co_return co_await PopAwaiter(*this);
    }

    /// Prevent further pushes and wake consumers once buffered items are read.
    void close()
    {
        std::deque<std::shared_ptr<PopState>> waiters;
        {
            std::lock_guard lock(mutex_);
            if(closed_)
            {
                return;
            }

            closed_ = true;
            if(!values_.empty())
            {
                return;
            }

            waiters.swap(waiters_);
        }

        for(const auto& waiter : waiters)
        {
            if(waiter->continuation != nullptr)
            {
                waiter->continuation.resume();
            }
        }
    }

private:
    struct PopState
    {
        std::optional<T> value;
        std::coroutine_handle<> continuation = nullptr;
    };

    class PopAwaiter
    {
    public:
        explicit PopAwaiter(AsyncQueue& queue)
            : queue_(&queue), state_(std::make_shared<PopState>())
        {
        }

        bool await_ready()
        {
            std::lock_guard lock(queue_->mutex_);
            return queue_->take(*state_);
        }

        bool await_suspend(std::coroutine_handle<> continuation)
        {
            std::lock_guard lock(queue_->mutex_);
            if(queue_->take(*state_))
            {
                return false;
            }

            state_->continuation = continuation;
            queue_->waiters_.push_back(state_);
            return true;
        }

        std::optional<T> await_resume()
        {
            return std::move(state_->value);
        }

    private:
        AsyncQueue* queue_;
        std::shared_ptr<PopState> state_;
    };

    bool take(PopState& state)
    {
        if(!values_.empty())
        {
            state.value.emplace(std::move(values_.front()));
            values_.pop_front();
            return true;
        }

        return closed_;
    }

    std::mutex mutex_;
    std::deque<T> values_;
    std::deque<std::shared_ptr<PopState>> waiters_;
    bool closed_ = false;
};

} // namespace agent_smith::internal