#include "vdm/segment/budget.hpp"

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <functional>
#include <mutex>
#include <optional>
#include <thread>
#include <vector>

#include "vtest.hpp"

using namespace vdm;
using namespace vdm::segment;
using EB = SegmentBudget::EngineBudget;

namespace {

TaskId tid(std::uint64_t v) {
    return TaskId{v};
}

// A test task that reacts to slot targets the way stage 8's download_task will: start
// workers up to the target, release them when the target drops. Purely bookkeeping.
//
// Production's real equivalent (register_task()'s on_target lambda, download_task.cpp)
// never calls back in synchronously -- it posts through host.schedule() (engine.cpp), so
// apply_slot_target() never runs on the same call stack as whatever budget call triggered
// it. wake_one_waiter_locked()'s wakeups, though, mean on_target() here *can* now be
// reentered on the same thread (e.g. this task's own release_slot() call, deep in the grow
// loop below, can cascade into waking a *different* task whose own release, in turn,
// cascades back to this one) -- a plain std::mutex would self-deadlock on that, the way
// FakeTask::mu almost did until this was written. A recursive_mutex plus coalescing the
// reentrant call's target into `pending` (processed by the outer call's loop once its
// current pass finishes) keeps this a faithful, deterministic stand-in without actually
// needing real threads: no callback ever runs nested inside a still-running instance of
// itself, and a burst of reentrant retargeting converges to the latest value instead of
// each one doing its own redundant grow/shrink pass.
struct FakeTask {
    SegmentBudget *budget = nullptr;
    TaskId id{};
    std::recursive_mutex mu;
    std::uint32_t workers = 0;
    std::uint32_t target = 0;
    bool running = false;
    std::optional<std::uint32_t> pending;

    FakeTask() = default;
    FakeTask(SegmentBudget *b, TaskId i) : budget(b), id(i) {}

    void on_target(std::uint32_t t) {
        std::lock_guard lk(mu);
        if (running) {
            pending = t;  // reentrant call on this thread -- the running instance picks it up
            return;
        }
        running = true;
        std::uint32_t want = t;
        for (;;) {
            target = want;
            while (workers < target) {
                if (!budget->confirm_slot(id))
                    break;
                ++workers;
            }
            // over target -> yield the excess immediately (a real task waits for a
            // boundary)
            while (workers > target) {
                budget->release_slot(id);
                --workers;
            }
            if (!pending)
                break;
            want = *pending;
            pending.reset();
        }
        running = false;
    }
    std::uint32_t held() {
        std::lock_guard lk(mu);
        return workers;
    }
};

// Mirrors Engine's real timer thread (engine.cpp: timer_loop) and how register_task()'s
// on_target lambda actually reaches a task (download_task.cpp: host.schedule(), never a
// direct call). One dedicated thread drains queued closures one at a time -- the
// serialization that makes production immune to the cross-task deadlock risk FakeTask's
// synchronous, whatever-thread-triggered-it delivery has under heavy *concurrent* driving
// (budget_concurrent_confirm_release_stays_consistent, below, is the one test that
// actually exercises this: multiple threads calling set_want() for *different* tasks at
// once, each able to cascade into a wake for another task's callback -- two such cascades
// landing on two different FakeTask mutexes in opposite orders on two different threads is
// a real AB-BA deadlock a recursive_mutex alone doesn't prevent, since that only guards a
// single thread against re-entering itself). Every other test in this file drives the
// budget from one thread at a time, where that risk can't arise, so plain FakeTask is
// still the simpler, sufficient double there.
class TestTimer {
  public:
    TestTimer() {
        worker_ = std::jthread([this](std::stop_token st) { run(st); });
    }
    ~TestTimer() {
        worker_.request_stop();
        cv_.notify_all();
    }
    void post(std::function<void()> fn) {
        {
            std::lock_guard lk(mu_);
            queue_.push_back(std::move(fn));
        }
        cv_.notify_one();
    }
    // Blocks until the queue is empty and nothing is mid-run -- what a test needs before
    // asserting on state this timer's closures mutate.
    void drain() {
        std::unique_lock lk(mu_);
        cv_done_.wait(lk, [&] { return queue_.empty() && !running_; });
    }

  private:
    void run(std::stop_token st) {
        std::unique_lock lk(mu_);
        while (true) {
            cv_.wait(lk, st, [&] { return !queue_.empty(); });
            if (st.stop_requested())
                return;
            auto fn = std::move(queue_.front());
            queue_.erase(queue_.begin());
            running_ = true;
            lk.unlock();
            fn();
            lk.lock();
            running_ = false;
            if (queue_.empty())
                cv_done_.notify_all();
        }
    }
    std::mutex mu_;
    std::condition_variable_any cv_;  // _any: wait() below takes a stop_token predicate
    std::condition_variable cv_done_;
    std::vector<std::function<void()>> queue_;
    bool running_ = false;
    std::jthread worker_;
};

// Same grow/shrink logic as FakeTask, but on_target() only ever posts through a TestTimer
// instead of running synchronously -- see the comment above TestTimer for why that's the
// faithful model under concurrent driving.
struct AsyncFakeTask {
    SegmentBudget *budget = nullptr;
    TestTimer *timer = nullptr;
    TaskId id{};
    std::mutex mu;  // only the timer thread ever touches workers/target -- plain suffices
    std::uint32_t workers = 0;
    std::uint32_t target = 0;

    void on_target(std::uint32_t t) {
        timer->post([this, t] {
            std::lock_guard lk(mu);
            target = t;
            while (workers < target) {
                if (!budget->confirm_slot(id))
                    break;
                ++workers;
            }
            while (workers > target) {
                budget->release_slot(id);
                --workers;
            }
        });
    }
    std::uint32_t held() {
        std::lock_guard lk(mu);
        return workers;
    }
};

}  // namespace

VT_TEST(budget_single_task_grows_to_cap) {
    SegmentBudget b({.max_active_segments = 32});
    FakeTask t{&b, tid(1)};
    b.register_task(tid(1), {.host = "h", .per_task_cap = 8, .resumable = true},
                    [&](std::uint32_t n) { t.on_target(n); });
    b.set_want(tid(1), 8);
    VT_CHECK_EQ(t.held(), 8u);
    VT_CHECK_EQ(b.segments_active(tid(1)), 8u);
    VT_CHECK_EQ(b.budget().active, 8u);
    VT_CHECK_EQ(b.budget().tasks_starved, 0u);
}

VT_TEST(budget_min_one_before_seconds) {
    // Budget of 3, two tasks each wanting 8. min-1 first: each gets 1, then the higher-
    // priority one grows to 2.
    SegmentBudget b({.max_active_segments = 3});
    FakeTask a{&b, tid(1)}, c{&b, tid(2)};
    b.register_task(tid(1), {.host = "h1", .per_task_cap = 8, .resumable = true},
                    [&](std::uint32_t n) { a.on_target(n); });
    b.register_task(tid(2), {.host = "h2", .per_task_cap = 8, .resumable = true},
                    [&](std::uint32_t n) { c.on_target(n); });
    std::vector<TaskId> order = {tid(1), tid(2)};
    b.set_task_order(order);
    b.set_want(tid(1), 8);
    b.set_want(tid(2), 8);

    VT_CHECK(a.held() >= 1);  // guarantee
    VT_CHECK(c.held() >= 1);  // guarantee — the load-bearing property
    VT_CHECK_EQ(a.held() + c.held(), 3u);
    VT_CHECK_EQ(a.held(), 2u);  // higher priority took the growth slot
}

VT_TEST(budget_new_high_priority_task_gets_min_one_via_yield) {
    SegmentBudget b({.max_active_segments = 4});
    FakeTask a{&b, tid(1)};
    b.register_task(tid(1), {.host = "h", .per_task_cap = 8, .resumable = true},
                    [&](std::uint32_t n) { a.on_target(n); });
    b.set_task_order(std::vector<TaskId>{tid(1)});
    b.set_want(tid(1), 8);
    VT_CHECK_EQ(a.held(), 4u);  // hogging the whole budget

    // a second, higher-priority task arrives
    FakeTask c{&b, tid(2)};
    b.register_task(tid(2), {.host = "h2", .per_task_cap = 8, .resumable = true},
                    [&](std::uint32_t n) { c.on_target(n); });
    b.set_task_order(std::vector<TaskId>{tid(2), tid(1)});
    b.set_want(tid(2), 8);

    // a yields so c gets at least its guaranteed slot; the surplus is shared round-robin.
    VT_CHECK(c.held() >= 1);  // min-1 — the load-bearing guarantee
    VT_CHECK(a.held() >= 1);  // a keeps its own min-1
    VT_CHECK(a.held() < 4);   // a really did yield at least one
    VT_CHECK_EQ(a.held() + c.held(), 4u);
    VT_CHECK_EQ(b.budget().active, 4u);
    VT_CHECK_EQ(b.budget().tasks_starved, 0u);
}

VT_TEST(budget_host_cap_clamps_effective_target) {
    SegmentBudget b({.max_active_segments = 32});
    FakeTask t{&b, tid(1)};
    b.set_host_segment_cap("slowcdn", 4);
    b.register_task(tid(1), {.host = "slowcdn", .per_task_cap = 16, .resumable = true},
                    [&](std::uint32_t n) { t.on_target(n); });
    b.set_want(tid(1), 16);
    VT_CHECK_EQ(t.held(), 4u);  // clamped by the host cap, not per_task_cap

    b.set_host_segment_cap("slowcdn", 0);  // clear
    VT_CHECK_EQ(t.held(), 16u);
}

VT_TEST(budget_non_resumable_task_capped_at_one) {
    SegmentBudget b({.max_active_segments = 32});
    FakeTask t{&b, tid(1)};
    b.register_task(tid(1), {.host = "h", .per_task_cap = 8, .resumable = false},
                    [&](std::uint32_t n) { t.on_target(n); });
    b.set_want(tid(1), 8);
    VT_CHECK_EQ(t.held(), 1u);
}

VT_TEST(budget_live_lower_sheds_via_yield_lowest_priority_first) {
    SegmentBudget b({.max_active_segments = 24});
    FakeTask a{&b, tid(1)}, c{&b, tid(2)}, d{&b, tid(3)};
    for (auto *ft : {&a, &c, &d})
        b.register_task(ft->id, {.host = "h", .per_task_cap = 8, .resumable = true},
                        [ft](std::uint32_t n) { ft->on_target(n); });
    b.set_task_order(std::vector<TaskId>{tid(1), tid(2), tid(3)});
    for (auto id : {tid(1), tid(2), tid(3)})
        b.set_want(id, 8);
    VT_CHECK_EQ(a.held() + c.held() + d.held(), 24u);  // 8 + 8 + 8

    b.set_max_active_segments(10);  // live cut
    VT_CHECK_EQ(a.held() + c.held() + d.held(), 10u);
    VT_CHECK(a.held() >= c.held() && c.held() >= d.held());     // priority order preserved
    VT_CHECK(a.held() >= 1 && c.held() >= 1 && d.held() >= 1);  // min-1 still honoured
}

VT_TEST(budget_live_lower_below_task_count_starves_the_tail) {
    SegmentBudget b({.max_active_segments = 6});
    std::vector<FakeTask> ts(4);
    for (std::uint32_t i = 0; i < 4; ++i) {
        ts[i].budget = &b;
        ts[i].id = tid(i + 1);
    }
    for (auto &ft : ts)
        b.register_task(ft.id, {.host = "h", .per_task_cap = 4, .resumable = true},
                        [&ft](std::uint32_t n) { ft.on_target(n); });
    b.set_task_order(std::vector<TaskId>{tid(1), tid(2), tid(3), tid(4)});
    for (auto &ft : ts)
        b.set_want(ft.id, 4);
    VT_CHECK_EQ(b.budget().tasks_starved, 0u);

    b.set_max_active_segments(3);  // below the running-task count
    VT_CHECK_EQ(ts[0].held(), 1u);
    VT_CHECK_EQ(ts[3].held(), 0u);  // lowest priority shed to zero
    VT_CHECK_EQ(b.budget().tasks_starved, 1u);
    VT_REQUIRE(b.starved_tasks().size() == 1);
    VT_CHECK_EQ(b.starved_tasks()[0], tid(4));
    VT_CHECK(b.starved_since(tid(4)).has_value());
    VT_CHECK(!b.starved_since(tid(1)).has_value());
}

VT_TEST(budget_deregister_frees_slots_to_starved) {
    SegmentBudget b({.max_active_segments = 4});
    FakeTask a{&b, tid(1)}, c{&b, tid(2)};
    b.register_task(tid(1), {.host = "h", .per_task_cap = 8, .resumable = true},
                    [&](std::uint32_t n) { a.on_target(n); });
    b.set_task_order(std::vector<TaskId>{tid(1)});
    b.set_want(tid(1), 8);
    VT_CHECK_EQ(a.held(), 4u);

    b.register_task(tid(2), {.host = "h", .per_task_cap = 8, .resumable = true},
                    [&](std::uint32_t n) { c.on_target(n); });
    b.set_task_order(std::vector<TaskId>{tid(1), tid(2)});
    b.set_want(tid(2), 8);
    VT_CHECK(c.held() >= 1);  // min-1 from a's yield

    b.deregister_task(tid(1));
    VT_CHECK_EQ(c.held(), 4u);  // c grows into the whole freed budget
    VT_CHECK_EQ(b.budget().active, 4u);
}

VT_TEST(budget_on_changed_fires_on_starved_edge) {
    SegmentBudget b({.max_active_segments = 1, .notify_period = std::chrono::milliseconds{40}});
    std::mutex m;
    std::vector<EB> seen;
    b.on_budget_changed([&](EB e) {
        std::lock_guard lk(m);
        seen.push_back(e);
    });

    FakeTask a{&b, tid(1)}, c{&b, tid(2)};
    b.register_task(tid(1), {.host = "h", .per_task_cap = 4, .resumable = true},
                    [&](std::uint32_t n) { a.on_target(n); });
    b.register_task(tid(2), {.host = "h", .per_task_cap = 4, .resumable = true},
                    [&](std::uint32_t n) { c.on_target(n); });
    b.set_task_order(std::vector<TaskId>{tid(1), tid(2)});
    b.set_want(tid(1), 4);
    b.set_want(tid(2), 4);  // budget is 1 -> tid(2) is starved: 0 -> nonzero edge

    // the edge fire is synchronous on the triggering call
    bool saw_starved = false;
    {
        std::lock_guard lk(m);
        for (auto &e : seen)
            if (e.tasks_starved > 0)
                saw_starved = true;
    }
    VT_CHECK(saw_starved);

    b.deregister_task(tid(1));  // frees the slot -> tid(2) no longer starved: edge back
    std::this_thread::sleep_for(std::chrono::milliseconds(120));
    bool saw_unstarved_after = false;
    {
        std::lock_guard lk(m);
        VT_CHECK(!seen.empty());
        saw_unstarved_after = seen.back().tasks_starved == 0;
    }
    VT_CHECK(saw_unstarved_after);
}

VT_TEST(budget_concurrent_confirm_release_stays_consistent) {
    SegmentBudget b({.max_active_segments = 16});
    constexpr int kTasks = 6;
    TestTimer timer;
    std::vector<std::unique_ptr<AsyncFakeTask>> ts;
    for (int i = 0; i < kTasks; ++i) {
        ts.push_back(std::make_unique<AsyncFakeTask>());
        ts.back()->budget = &b;
        ts.back()->timer = &timer;
        ts.back()->id = tid(i + 1);
        AsyncFakeTask *ft = ts.back().get();
        b.register_task(ft->id, {.host = "h", .per_task_cap = 6, .resumable = true},
                        [ft](std::uint32_t n) { ft->on_target(n); });
    }
    std::vector<std::jthread> drivers;
    for (int i = 0; i < kTasks; ++i) {
        drivers.emplace_back([&, id = tid(i + 1)] {
            for (int r = 0; r < 4000; ++r)
                b.set_want(id, (r % 7));
        });
    }
    drivers.clear();  // join
    for (auto &ft : ts)
        b.set_want(ft->id, 0);
    timer.drain();  // let every queued on_target actually run before asserting

    // With everyone wanting nothing, the budget must be fully released.
    VT_CHECK_EQ(b.budget().active, 0u);
    std::uint32_t sum = 0;
    for (auto &ft : ts)
        sum += b.segments_active(ft->id);
    VT_CHECK_EQ(sum, 0u);
}

// --- wait-list wakeup: tools/bench heap-profile / load found a real task time out
// waiting on a slot its own target already said it should have (core/docs/m7-baseline.md,
// docs/adr/0012). Root cause: confirm_slot()'s engine-wide cap check (needed so active_
// never exceeds max_active_segments -- a real over-admission bug, not just this liveness
// gap) can deny a task whose target is already correct, when a *different* task is
// legitimately still holding more than its own just-lowered target (yield is deferred to
// a segment boundary, ADR 0011 A1). Nothing in the plain target-changed callback
// mechanism ever revisits a task whose target didn't change -- it was already right. ---

namespace {

// Unlike FakeTask above, on_target() here only enqueues -- it never calls back into the
// budget synchronously. This matches production exactly: register_task()'s on_target
// lambda (download_task.cpp) posts through host.schedule() (engine.cpp), so
// apply_slot_target() never runs on the same call stack as whatever budget call triggered
// it. The test drives delivery explicitly (deliver_one()) instead of a background thread
// so the race this test exists to force -- confirm_slot() denied before the task that's
// over its target has processed its own shrink -- is deterministic, not a timing gamble.
struct QueuedTask {
    SegmentBudget *budget = nullptr;
    TaskId id{};
    std::uint32_t workers = 0;
    std::uint32_t target = 0;
    std::vector<std::uint32_t> pending;

    QueuedTask(SegmentBudget *b, TaskId i) : budget(b), id(i) {}

    void on_target(std::uint32_t n) { pending.push_back(n); }

    // Delivers the oldest queued target, applying it the way a real task's
    // apply_slot_target()/fill_slots_locked() would: try to grow to it (confirm_slot()
    // may deny), or shed down to it. Returns false (nothing to deliver) if the queue was
    // empty -- the condition VT_REQUIRE checks to prove a wakeup was actually queued.
    bool deliver_one() {
        if (pending.empty())
            return false;
        target = pending.front();
        pending.erase(pending.begin());
        while (workers < target) {
            if (!budget->confirm_slot(id))
                break;
            ++workers;
        }
        while (workers > target) {
            budget->release_slot(id);
            --workers;
        }
        return true;
    }
};

}  // namespace

VT_TEST(budget_release_wakes_a_denied_waiter) {
    // Sanity baseline: max_active=1, A (higher priority) holds it, B wants one too and is
    // fairly denied -- its target stays 0 while A outranks it and still wants its slot.
    // Once A stops wanting one, B's target rises and B is woken via the plain
    // target-changed path -- no wait-list needed for this simple case. The harder case
    // below is what actually needs it.
    SegmentBudget b({.max_active_segments = 1});
    QueuedTask A{&b, tid(1)}, B{&b, tid(2)};
    b.register_task(A.id, {.host = "h", .per_task_cap = 1, .resumable = true},
                    [&](std::uint32_t n) { A.on_target(n); });
    b.set_task_order(std::vector<TaskId>{A.id, B.id});
    b.set_want(A.id, 1);
    VT_REQUIRE(A.deliver_one());
    VT_CHECK_EQ(A.workers, 1u);

    // B registers and wants one too, but with A (higher priority) already holding the
    // only slot and still wanting it, B's fairly computed target stays 0 -- unchanged
    // from its just-registered value, so no callback is queued for it yet.
    b.register_task(B.id, {.host = "h", .per_task_cap = 1, .resumable = true},
                    [&](std::uint32_t n) { B.on_target(n); });
    b.set_want(B.id, 1);
    VT_CHECK(!B.deliver_one());
    VT_CHECK_EQ(B.workers, 0u);

    b.set_want(A.id, 0);          // A is done wanting a slot
    VT_REQUIRE(A.deliver_one());  // A's target dropped to 0 -- sheds its held slot
    VT_CHECK_EQ(A.workers, 0u);
    VT_REQUIRE(B.deliver_one());  // B's target rose to 1 -- the plain target-changed path
    VT_CHECK_EQ(B.workers, 1u);   // B took the freed slot
}

VT_TEST(budget_wait_list_wakes_a_task_whose_target_never_changed) {
    // The real gap. max_active=2. Y alone, holds both (target=2). X arrives wanting 1:
    // this recompute correctly drops Y's target to 1 (giving X its guaranteed slot) and
    // raises X's target to 1 -- both real target changes, both queued. Deliver X's
    // *first*: X's target says grow, but Y still physically holds 2 (hasn't processed
    // its own shrink yet) -- confirm_slot() must deny X here (active_ == max_active_),
    // which is the correctness fix (over-admission is the real RSS bug). Then Y
    // processes its shrink and actually releases. X's target never changes again -- it
    // was already correctly 1 -- so nothing in the plain mechanism ever revisits X.
    SegmentBudget b({.max_active_segments = 2});
    QueuedTask Y{&b, tid(1)}, X{&b, tid(2)};
    b.register_task(Y.id, {.host = "h", .per_task_cap = 2, .resumable = true},
                    [&](std::uint32_t n) { Y.on_target(n); });
    b.set_want(Y.id, 2);
    VT_REQUIRE(Y.deliver_one());
    VT_CHECK_EQ(Y.workers, 2u);

    b.register_task(X.id, {.host = "h", .per_task_cap = 1, .resumable = true},
                    [&](std::uint32_t n) { X.on_target(n); });
    b.set_task_order(std::vector<TaskId>{Y.id, X.id});
    b.set_want(X.id, 1);

    VT_REQUIRE(X.deliver_one());
    VT_CHECK_EQ(X.workers, 0u);  // denied: active_ == max_active_, even though X's target is 1

    VT_REQUIRE(Y.deliver_one());
    VT_CHECK_EQ(Y.workers, 1u);  // Y actually releases its excess now

    // The bug: without a wait-list, X.pending is empty here -- nothing was ever queued
    // for it, because X's target never changed again. X would wait forever despite its
    // target correctly saying it should hold a slot.
    VT_REQUIRE(X.deliver_one());
    VT_CHECK_EQ(X.workers, 1u);
    VT_CHECK_EQ(b.budget().active, 2u);
}