#include "rpc/event_loop.hpp"

#include <poll.h>

#include <algorithm>
#include <cerrno>
#include <cstdint>
#include <stdexcept>
#include <vector>

namespace velox::daemon::rpc {

EventLoop::EventLoop() {
    fds_.emplace(wakeup_.pollfd(), Entry{kRead, [this](int, unsigned) { wakeup_.drain(); }});
}

EventLoop::~EventLoop() = default;

void EventLoop::add_fd(int fd, unsigned interest, Callback cb) {
    fds_[fd] = Entry{interest, std::move(cb)};
}

void EventLoop::mod_fd(int fd, unsigned interest) {
    if (auto it = fds_.find(fd); it != fds_.end()) it->second.interest = interest;
}

void EventLoop::del_fd(int fd) {
    if (fd == wakeup_.pollfd()) return;  // internal, never removed
    fds_.erase(fd);
}

void EventLoop::wake() noexcept { wakeup_.signal(); }

void EventLoop::stop() noexcept {
    stop_requested_ = true;
    wake();
}

void EventLoop::post(std::function<void()> fn) {
    {
        std::lock_guard<std::mutex> lk(post_mu_);
        posts_.push_back(std::move(fn));
    }
    wake();
}

void EventLoop::drain_posts() {
    std::vector<std::function<void()>> batch;
    {
        std::lock_guard<std::mutex> lk(post_mu_);
        batch.swap(posts_);
    }
    for (auto& fn : batch) fn();
}

EventLoop::TimerId EventLoop::add_timer(std::chrono::milliseconds interval,
                                        std::function<void()> cb) {
    const TimerId id = next_timer_id_++;
    timers_.emplace(id, Timer{std::chrono::steady_clock::now() + interval, interval,
                              std::move(cb)});
    wake();  // the loop may already be blocked on a longer timeout
    return id;
}

void EventLoop::remove_timer(TimerId id) { timers_.erase(id); }

int EventLoop::next_timeout_ms() const {
    if (timers_.empty()) return -1;
    auto soonest = timers_.begin()->second.next;
    for (const auto& [id, t] : timers_) soonest = std::min(soonest, t.next);
    const auto now = std::chrono::steady_clock::now();
    const auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(soonest - now);
    return ms.count() > 0 ? static_cast<int>(ms.count()) : 0;
}

void EventLoop::run_due_timers() {
    const auto now = std::chrono::steady_clock::now();
    // Snapshot ids first: a callback may add/remove timers, which would invalidate
    // iteration over timers_ directly.
    std::vector<TimerId> due;
    for (auto& [id, t] : timers_) {
        if (t.next <= now) due.push_back(id);
    }
    for (const TimerId id : due) {
        const auto it = timers_.find(id);
        if (it == timers_.end()) continue;  // removed by an earlier callback this pass
        it->second.next = now + it->second.interval;
        it->second.cb();
    }
}

void EventLoop::run() {
    if (running_) throw std::logic_error("EventLoop::run() is not re-entrant");
    running_ = true;
    stop_requested_ = false;

    std::vector<pollfd> pfds;
    std::vector<int> fired;

    while (!stop_requested_) {
        pfds.clear();
        pfds.reserve(fds_.size());
        for (const auto& [fd, e] : fds_) {
            short ev = 0;
            if (e.interest & kRead) ev |= POLLIN;
            if (e.interest & kWrite) ev |= POLLOUT;
            if (ev == 0 && fd != wakeup_.pollfd()) continue;
            pollfd p{};
            p.fd = fd;
            p.events = ev;
            pfds.push_back(p);
        }

        const int rc = ::poll(pfds.data(), pfds.size(), next_timeout_ms());
        if (rc < 0) {
            if (errno == EINTR) continue;
            throw std::runtime_error("poll() failed");
        }

        drain_posts();
        run_due_timers();

        if (rc == 0) continue;

        // Snapshot the fds that fired before invoking any callback: a callback may erase
        // entries from fds_, which would invalidate iteration over pfds' referents.
        fired.clear();
        for (const auto& p : pfds) {
            if (p.revents != 0) fired.push_back(p.fd);
        }

        for (const int fd : fired) {
            const auto it = fds_.find(fd);
            if (it == fds_.end()) continue;  // removed by an earlier callback this pass

            // Recompute revents for this fd from the snapshot.
            unsigned events = 0;
            for (const auto& p : pfds) {
                if (p.fd != fd) continue;
                if (p.revents & (POLLIN | POLLHUP | POLLERR)) events |= kRead;
                if (p.revents & POLLOUT) events |= kWrite;
                break;
            }
            if (events != 0) it->second.cb(fd, events);
        }
    }

    running_ = false;
}

}  // namespace velox::daemon::rpc