#include "sched/governor.hpp"

#include <algorithm>
#include <string>

#include "check.hpp"

using namespace velox::daemon::sched;

namespace {

bool has(const std::vector<std::string>& v, const std::string& x) {
    return std::find(v.begin(), v.end(), x) != v.end();
}

TaskView task(std::string id, RunState st, std::optional<std::string> queue = std::nullopt,
              std::int64_t pos = 0, std::string host = "", std::int64_t rank = 0) {
    TaskView t;
    t.task_id = std::move(id);
    t.run_state = st;
    t.queue_id = std::move(queue);
    t.queue_position = pos;
    t.host = std::move(host);
    t.admit_rank = rank;
    return t;
}

QueueView queue(std::string id, bool running, std::int64_t cap, bool window = true) {
    QueueView q;
    q.queue_id = std::move(id);
    q.running = running;
    q.max_concurrent = cap;
    q.window_open = window;
    return q;
}

}  // namespace

void run() {
    // --- global cap: 5 queued, max 3 -> start 3, order is all 3 --------------------
    {
        Governor g(GovernorConfig{.max_concurrent_downloads = 3, .max_active_segments = 32});
        std::vector<TaskView> t;
        for (int i = 0; i < 5; ++i)
            t.push_back(task("t" + std::to_string(i), RunState::Queued, std::nullopt, 0, "", i));
        const auto d = g.evaluate(t, {});
        CHECK_EQ(d.to_start.size(), 3u);
        CHECK(has(d.to_start, "t0") && has(d.to_start, "t1") && has(d.to_start, "t2"));
        CHECK_EQ(d.priority_order.size(), 3u);
        CHECK_EQ(d.priority_order.front(), std::string("t0"));  // FIFO by admit_rank
        CHECK(d.to_pause.empty());
    }

    // --- the ADR 0011 §2 clamp: maxActiveSegments below maxConcurrentDownloads ------
    {
        Governor g(GovernorConfig{.max_concurrent_downloads = 10, .max_active_segments = 2});
        std::vector<TaskView> t;
        for (int i = 0; i < 6; ++i) t.push_back(task("t" + std::to_string(i), RunState::Queued));
        const auto d = g.evaluate(t, {});
        CHECK_EQ(d.to_start.size(), 2u);  // clamped to the segment budget
    }

    // --- per-queue cap: queue running, cap 2, 4 queued in it -----------------------
    {
        Governor g(GovernorConfig{.max_concurrent_downloads = 10, .max_active_segments = 32});
        std::vector<TaskView> t;
        for (int i = 0; i < 4; ++i)
            t.push_back(task("q" + std::to_string(i), RunState::Queued, "Q", i));
        const auto d = g.evaluate(t, {queue("Q", true, 2)});
        CHECK_EQ(d.to_start.size(), 2u);
        CHECK(has(d.to_start, "q0") && has(d.to_start, "q1"));  // lowest queue_position first
    }

    // --- a stopped queue: its running tasks are paused (QueueStopped) -------------
    {
        Governor g(GovernorConfig{.max_concurrent_downloads = 10, .max_active_segments = 32});
        std::vector<TaskView> t{
            task("r0", RunState::Running, "Q", 0),
            task("r1", RunState::Running, "Q", 1),
            task("m0", RunState::Running, std::nullopt, 0, "", 5),  // main list, unaffected
        };
        const auto d = g.evaluate(t, {queue("Q", /*running=*/false, 4)});
        CHECK_EQ(d.to_pause.size(), 2u);
        CHECK(has(d.to_pause, "r0") && has(d.to_pause, "r1"));
        CHECK(d.pause_reasons.at("r0") == PauseReason::QueueStopped);
        CHECK(!has(d.to_pause, "m0"));
    }

    // --- schedule window closed: pause reason is Schedule, not QueueStopped -------
    {
        Governor g(GovernorConfig{.max_concurrent_downloads = 10, .max_active_segments = 32});
        std::vector<TaskView> t{task("r0", RunState::Running, "Q", 0)};
        const auto d = g.evaluate(t, {queue("Q", /*running=*/true, 4, /*window=*/false)});
        CHECK_EQ(d.to_pause.size(), 1u);
        CHECK(d.pause_reasons.at("r0") == PauseReason::Schedule);
    }

    // --- resume: a task paused for Schedule comes back when the window reopens ----
    {
        Governor g(GovernorConfig{.max_concurrent_downloads = 10, .max_active_segments = 32});
        TaskView p = task("p0", RunState::Paused, "Q", 0);
        p.pause_reason = PauseReason::Schedule;
        TaskView u = task("u0", RunState::Paused, "Q", 1);
        u.pause_reason = PauseReason::User;  // must NOT be auto-resumed
        const auto d = g.evaluate({p, u}, {queue("Q", true, 4, true)});
        CHECK_EQ(d.to_resume.size(), 1u);
        CHECK(has(d.to_resume, "p0"));
        CHECK(!has(d.to_resume, "u0"));
    }

    // --- an Auto (auth) pause is never touched, even with slots free -------------
    {
        Governor g(GovernorConfig{.max_concurrent_downloads = 10, .max_active_segments = 32});
        TaskView a = task("a0", RunState::Paused, std::nullopt);
        a.pause_reason = PauseReason::Auto;
        const auto d = g.evaluate({a}, {});
        CHECK(d.to_resume.empty());
        CHECK(d.to_pause.empty());
    }

    // --- per-host cap: 3 queued on the same host, cap 1 ------------------------
    {
        GovernorConfig cfg{.max_concurrent_downloads = 10, .max_active_segments = 32};
        cfg.host_caps["cdn.example"] = 1;
        Governor g(cfg);
        std::vector<TaskView> t{
            task("h0", RunState::Queued, std::nullopt, 0, "cdn.example", 0),
            task("h1", RunState::Queued, std::nullopt, 0, "cdn.example", 1),
            task("h2", RunState::Queued, std::nullopt, 0, "other.example", 2),
        };
        const auto d = g.evaluate(t, {});
        CHECK_EQ(d.to_start.size(), 2u);              // one per host
        CHECK(has(d.to_start, "h0") && has(d.to_start, "h2"));
        CHECK(!has(d.to_start, "h1"));
    }

    // --- admission reconcile: 4 running, cap drops to 2 -> pause the 2 lowest -----
    {
        Governor g(GovernorConfig{.max_concurrent_downloads = 2, .max_active_segments = 32});
        std::vector<TaskView> t;
        for (int i = 0; i < 4; ++i)
            t.push_back(task("t" + std::to_string(i), RunState::Running, std::nullopt, 0, "", i));
        const auto d = g.evaluate(t, {});
        CHECK_EQ(d.to_pause.size(), 2u);
        CHECK(has(d.to_pause, "t2") && has(d.to_pause, "t3"));   // lowest priority (highest rank)
        CHECK(d.pause_reasons.at("t2") == PauseReason::AdmissionReconcile);
        CHECK_EQ(d.priority_order.size(), 2u);
    }

    // --- determinism: shuffled input, identical Decision -----------------------
    {
        Governor g(GovernorConfig{.max_concurrent_downloads = 2, .max_active_segments = 32});
        std::vector<TaskView> a{
            task("b", RunState::Queued, std::nullopt, 0, "", 1),
            task("a", RunState::Queued, std::nullopt, 0, "", 0),
            task("c", RunState::Queued, std::nullopt, 0, "", 2),
        };
        std::vector<TaskView> b{a[2], a[0], a[1]};
        const auto da = g.evaluate(a, {});
        const auto db = g.evaluate(b, {});
        CHECK(da.to_start == db.to_start);
        CHECK(da.priority_order == db.priority_order);
        CHECK_EQ(da.priority_order.front(), std::string("a"));
    }

    // --- main list outranks queues in priority_order -------------------------
    {
        Governor g(GovernorConfig{.max_concurrent_downloads = 10, .max_active_segments = 32});
        std::vector<TaskView> t{
            task("qtask", RunState::Running, "Q", 0),
            task("mtask", RunState::Running, std::nullopt, 0, "", 99),
        };
        const auto d = g.evaluate(t, {queue("Q", true, 4)});
        CHECK_EQ(d.priority_order.front(), std::string("mtask"));
    }
}

TEST_MAIN()