#pragma once

#include <atomic>
#include <memory>
#include <stop_token>

#include "async_queue.hpp"
#include "background_tasks.hpp"
#include "../cancellation.hpp"
#include "../error.hpp"

namespace agent_smith::internal
{

template <typename T> struct TimeoutState
{
    std::atomic<bool> settled{false};
    AsyncQueue<mw::E<T>> result;
    std::stop_source timer_stop;
    CancellationSource cancel;
};

template <typename T>
mw::Task<mw::E<void>> finishOperation(
    mw::Task<mw::E<T>> task, std::shared_ptr<TimeoutState<T>> state)
{
    auto result = co_await std::move(task);
    if(!state->settled.exchange(true))
    {
        state->timer_stop.request_stop();
        state->result.push(std::move(result));
    }
    co_return mw::E<void>{};
}

template <typename T>
mw::Task<mw::E<void>> finishTimeout(
    Runtime& runtime, std::chrono::milliseconds timeout,
    std::shared_ptr<TimeoutState<T>> state)
{
    auto slept = co_await runtime.sleepFor(
        timeout, state->timer_stop.get_token());
    if(!state->settled.exchange(true))
    {
        state->cancel.cancel();
        if(!slept)
        {
            state->result.push(std::unexpected(slept.error()));
        }
        else
        {
            state->result.push(std::unexpected(mw::Error(TimeoutError{
                .msg = "tool execution timed out",
                .operation = "tool invocation",
                .budget = timeout,
            })));
        }
    }
    co_return mw::E<void>{};
}

/// Race an operation against a timer, retaining the loser in the runtime.
/// The task must own everything it uses after a timeout. The runtime must
/// outlive this call; shutdown drains non-cooperative operations.
template <typename T>
mw::Task<mw::E<T>> withTimeout(
    Runtime& runtime, mw::Task<mw::E<T>> task,
    std::chrono::milliseconds timeout, CancellationSource* cancel = nullptr)
{
    if(timeout <= std::chrono::milliseconds::zero())
    {
        if(cancel != nullptr)
        {
            cancel->cancel();
        }
        co_return std::unexpected(mw::Error(TimeoutError{
            .msg = "tool execution timed out",
            .operation = "tool invocation",
            .budget = timeout,
        }));
    }
    auto state = std::make_shared<TimeoutState<T>>();
    if(cancel != nullptr)
    {
        state->cancel = *cancel;
    }
    BackgroundTasks::start(runtime, finishOperation(std::move(task), state));
    BackgroundTasks::start(runtime, finishTimeout(runtime, timeout, state));
    auto result = co_await state->result.pop();
    co_return std::move(*result);
}

} // namespace agent_smith::internal