// vdm/segment/budget.cpp

#include "vdm/segment/budget.hpp"

#include <algorithm>

namespace vdm::segment {

SegmentBudget::SegmentBudget() : SegmentBudget(Options{}) {}

SegmentBudget::SegmentBudget(Options opts)
    : max_active_(opts.max_active_segments ? opts.max_active_segments : 1),
      notify_period_(opts.notify_period) {
    notifier_ = std::jthread([this](std::stop_token st) { notifier_loop(st); });
}

SegmentBudget::~SegmentBudget() {
    notifier_.request_stop();
    notify_cv_.notify_all();
}

// --- allocation -----------------------------------------------------------------------

std::uint32_t SegmentBudget::effective_cap_locked(const Task &t) const {
    std::uint32_t base = t.resumable ? t.per_task_cap : 1;
    base = std::clamp<std::uint32_t>(base, 1, 32);
    if (auto it = host_caps_.find(t.host); it != host_caps_.end() && it->second > 0)
        base = std::min(base, it->second);
    return base;
}

SegmentBudget::EngineBudget SegmentBudget::snapshot_locked() const {
    std::uint32_t starved = 0;
    for (const auto &[id, t] : tasks_)
        if (t.want >= 1 && t.held == 0)
            ++starved;
    return EngineBudget{max_active_, active_, starved};
}

// The two-pass fairness allocation. Recomputes every task's target from scratch (so a
// live cap cut naturally produces target < held -> yield), diffs against the last
// published target, and collects the callbacks to fire once mu_ is released.
SegmentBudget::Plan SegmentBudget::reallocate_locked() {
    // Priority order: DAEMON's list first, then any registered task not in it (defensive;
    // "a running task absent from the list sorts last").
    std::vector<TaskId> order;
    order.reserve(tasks_.size());
    for (TaskId id : order_)
        if (tasks_.count(id))
            order.push_back(id);
    for (const auto &[id, _] : tasks_)
        if (std::find(order.begin(), order.end(), id) == order.end())
            order.push_back(id);

    std::unordered_map<TaskId, std::uint32_t> target;
    target.reserve(order.size());
    std::uint32_t pool = max_active_;

    auto capped_want = [&](TaskId id) {
        const Task &t = tasks_.at(id);
        return std::min(t.want, effective_cap_locked(t));
    };

    // Guarantee pass: one slot each, in priority order, to anyone who wants one.
    for (TaskId id : order) {
        if (pool == 0)
            break;
        if (capped_want(id) >= 1) {
            target[id] = 1;
            --pool;
        }
    }
    // Growth pass: round-robin the remainder, up to each task's effective cap.
    while (pool > 0) {
        bool granted = false;
        for (TaskId id : order) {
            if (pool == 0)
                break;
            std::uint32_t &tv = target[id];
            if (tv < capped_want(id)) {
                ++tv;
                --pool;
                granted = true;
            }
        }
        if (!granted)
            break;
    }

    const SteadyTime now = std::chrono::steady_clock::now();
    Plan plan;
    for (auto &[id, t] : tasks_) {
        std::uint32_t nt = target.count(id) ? target[id] : 0;
        if (nt != t.target) {
            t.target = nt;
            if (t.on_target)
                plan.targets.emplace_back(t.on_target, nt);
        }
        // starvation timestamp bookkeeping
        bool starved_now = t.want >= 1 && t.held == 0;
        if (starved_now && !t.starved_since)
            t.starved_since = now;
        if (!starved_now)
            t.starved_since.reset();
    }

    EngineBudget eb = snapshot_locked();
    bool starved_edge = (eb.tasks_starved == 0) != (last_starved_ == 0);
    if (eb != last_notified_)
        dirty_ = true;
    last_starved_ = eb.tasks_starved;

    if (starved_edge && on_changed_) {
        plan.notify_now = std::make_pair(on_changed_, eb);
        last_notified_ = eb;
        dirty_ = false;
    }
    if (dirty_)
        notify_cv_.notify_one();
    return plan;
}

void SegmentBudget::run(Plan &p) {
    for (auto &[fn, n] : p.targets)
        if (fn)
            fn(n);
    if (p.notify_now && p.notify_now->first)
        p.notify_now->first(p.notify_now->second);
}

// --- task-facing --------------------------------------------------------------------

void SegmentBudget::register_task(TaskId id, const TaskParams &params, SlotTargetFn on_target) {
    Plan plan;
    {
        std::lock_guard lk(mu_);
        Task t;
        t.host = params.host;
        t.per_task_cap = params.per_task_cap ? params.per_task_cap : 1;
        t.resumable = params.resumable;
        t.on_target = std::move(on_target);
        tasks_[id] = std::move(t);
        plan = reallocate_locked();
    }
    run(plan);
}

void SegmentBudget::deregister_task(TaskId id) {
    Plan plan;
    {
        std::lock_guard lk(mu_);
        auto it = tasks_.find(id);
        if (it == tasks_.end())
            return;
        active_ -= it->second.held;
        tasks_.erase(it);
        plan = reallocate_locked();
    }
    run(plan);
}

void SegmentBudget::set_want(TaskId id, std::uint32_t want) {
    Plan plan;
    {
        std::lock_guard lk(mu_);
        auto it = tasks_.find(id);
        if (it == tasks_.end())
            return;
        if (it->second.want == want)
            return;
        it->second.want = want;
        plan = reallocate_locked();
    }
    run(plan);
}

bool SegmentBudget::confirm_slot(TaskId id) {
    std::lock_guard lk(mu_);
    auto it = tasks_.find(id);
    if (it == tasks_.end())
        return false;
    Task &t = it->second;
    if (t.held >= t.target)
        return false;  // target was cut in the race
    ++t.held;
    ++active_;
    if (snapshot_locked() != last_notified_) {
        dirty_ = true;
        notify_cv_.notify_one();
    }
    return true;
}

void SegmentBudget::release_slot(TaskId id) {
    Plan plan;
    {
        std::lock_guard lk(mu_);
        auto it = tasks_.find(id);
        if (it == tasks_.end() || it->second.held == 0)
            return;
        --it->second.held;
        --active_;
        plan = reallocate_locked();
    }
    run(plan);
}

// --- DAEMON-facing ---------------------------------------------------------------------

void SegmentBudget::set_max_active_segments(std::uint32_t n) {
    Plan plan;
    {
        std::lock_guard lk(mu_);
        n = n ? n : 1;
        if (n == max_active_)
            return;
        max_active_ = n;
        plan = reallocate_locked();
    }
    run(plan);
}

void SegmentBudget::set_host_segment_cap(std::string host, std::uint32_t cap) {
    Plan plan;
    {
        std::lock_guard lk(mu_);
        if (cap == 0)
            host_caps_.erase(host);
        else
            host_caps_[std::move(host)] = cap;
        plan = reallocate_locked();
    }
    run(plan);
}

void SegmentBudget::set_task_order(std::span<const TaskId> priority_order) {
    Plan plan;
    {
        std::lock_guard lk(mu_);
        order_.assign(priority_order.begin(), priority_order.end());
        plan = reallocate_locked();
    }
    run(plan);
}

SegmentBudget::EngineBudget SegmentBudget::budget() const {
    std::lock_guard lk(mu_);
    return snapshot_locked();
}

std::uint32_t SegmentBudget::segments_active(TaskId id) const {
    std::lock_guard lk(mu_);
    auto it = tasks_.find(id);
    return it == tasks_.end() ? 0 : it->second.held;
}

std::vector<TaskId> SegmentBudget::starved_tasks() const {
    std::lock_guard lk(mu_);
    std::vector<TaskId> out;
    for (const auto &[id, t] : tasks_)
        if (t.want >= 1 && t.held == 0)
            out.push_back(id);
    return out;
}

std::optional<SteadyTime> SegmentBudget::starved_since(TaskId id) const {
    std::lock_guard lk(mu_);
    auto it = tasks_.find(id);
    return it == tasks_.end() ? std::nullopt : it->second.starved_since;
}

void SegmentBudget::on_budget_changed(std::function<void(EngineBudget)> cb) {
    std::lock_guard lk(mu_);
    on_changed_ = std::move(cb);
}

// --- notifier thread: coalesced <=4 Hz -----------------------------------------------

void SegmentBudget::notifier_loop(std::stop_token st) {
    std::unique_lock lk(mu_);
    while (!st.stop_requested()) {
        notify_cv_.wait_for(lk, notify_period_, [&] { return dirty_ || st.stop_requested(); });
        if (st.stop_requested())
            break;
        if (!dirty_)
            continue;
        EngineBudget eb = snapshot_locked();
        auto cb = on_changed_;
        last_notified_ = eb;
        last_starved_ = eb.tasks_starved;
        dirty_ = false;
        lk.unlock();
        if (cb)
            cb(eb);
        lk.lock();
    }
}

}  // namespace vdm::segment