#include "vdm/segment/budget.hpp"

#include <atomic>
#include <chrono>
#include <mutex>
#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.
struct FakeTask {
    SegmentBudget *budget = nullptr;
    TaskId id{};
    std::mutex mu;
    std::uint32_t workers = 0;
    std::uint32_t target = 0;

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

    void on_target(std::uint32_t t) {
        std::lock_guard lk(mu);
        target = t;
        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;
        }
    }
    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;
    std::vector<std::unique_ptr<FakeTask>> ts;
    for (int i = 0; i < kTasks; ++i) {
        ts.push_back(std::make_unique<FakeTask>());
        ts.back()->budget = &b;
        ts.back()->id = tid(i + 1);
        FakeTask *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);

    // 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);
}