#include "timer_service.hpp"
#include <algorithm>

namespace agent_smith::internal
{

TimerService::TimerService()
    : worker_([this](std::stop_token stop) { worker(stop); })
{
}

TimerService::~TimerService()
{
    {
        std::lock_guard lock(mutex_);
        worker_.request_stop();
    }
    condition_.notify_all();
}

TimerService::SleepAwaiter::SleepAwaiter(TimerService& service,
                                         std::chrono::milliseconds duration,
                                         std::stop_token stop)
    : service_(&service), duration_(duration), entry_(std::make_shared<Entry>())
{
    entry_->stop = stop;
    callback_ = std::make_unique<std::stop_callback<std::function<void()>>>(
        stop, [&service]()
        {
            std::lock_guard lock(service.mutex_);
            service.condition_.notify_one();
        });
}

bool TimerService::SleepAwaiter::await_ready() const
{
    return duration_ <= std::chrono::milliseconds::zero() ||
           entry_->stop.stop_requested();
}

bool TimerService::SleepAwaiter::await_suspend(
    std::coroutine_handle<> continuation)
{
    entry_->deadline = std::chrono::steady_clock::now() + duration_;
    service_->schedule(entry_, continuation);
    return true;
}

TimerService::SleepAwaiter
TimerService::sleepFor(std::chrono::milliseconds duration,
                      std::stop_token stop)
{
    return SleepAwaiter(*this, duration, stop);
}

void TimerService::schedule(const std::shared_ptr<Entry>& entry,
                            std::coroutine_handle<> continuation)
{
    {
        std::lock_guard lock(mutex_);
        entry->continuation = continuation;
        entries_.push_back(entry);
    }
    condition_.notify_one();
}

void TimerService::worker(std::stop_token stop)
{
    std::unique_lock lock(mutex_);
    while(!stop.stop_requested())
    {
        if(entries_.empty())
        {
            condition_.wait(
                lock,
                [&]() { return stop.stop_requested() || !entries_.empty(); });
            continue;
        }

        const auto now = std::chrono::steady_clock::now();
        auto ready = std::find_if(entries_.begin(), entries_.end(),
            [now](const auto& entry)
            {
                return entry->stop.stop_requested() || entry->deadline <= now;
            });
        if(ready == entries_.end())
        {
            const auto earliest = std::min_element(
                entries_.begin(), entries_.end(),
                [](const auto& lhs, const auto& rhs)
                { return lhs->deadline < rhs->deadline; });
            condition_.wait_until(lock, (*earliest)->deadline);
            continue;
        }

        auto entry = *ready;
        entries_.erase(ready);
        auto continuation = entry->continuation;
        lock.unlock();
        if(continuation != nullptr)
        {
            continuation.resume();
        }
        lock.lock();
    }
}

} // namespace agent_smith::internal