#include "scr.h"
#include "testing.h"

#include <algorithm>
#include <cstdint>
#include <vector>

static std::vector<llama_token> make_prompt(int n_blocks, bool reordered) {
    std::vector<llama_token> out;
    out.reserve(n_blocks * 11);
    for (int i = 0; i < n_blocks; ++i) {
        const int block = reordered ? (i * 37 + 1) % n_blocks : i;
        out.push_back((llama_token) (100000 + block));
        for (int j = 0; j < 10; ++j) {
            out.push_back((llama_token) (block * 1000 + j + 1));
        }
    }
    return out;
}

static common_scr_params fixture_params(int max_runs = -1) {
    common_scr_params params;
    params.min_span = 8;
    params.seed = 4;
    if (max_runs >= 0) {
        params.max_runs = max_runs;
    }
    return params;
}

static void check_exact_runs(
        testing & t,
        const std::vector<llama_token> & old,
        const std::vector<llama_token> & cur,
        const common_scr_plan & plan) {
    for (const auto & run : plan.runs) {
        t.assert_true("run leaves one current token",
                      run.new_start + run.n < (int32_t) cur.size());
        t.assert_true("run is in old bounds",
                      run.old_start >= 0 && run.old_start + run.n <= (int32_t) old.size());
        t.assert_true("run is in current bounds",
                      run.new_start >= 0 && run.new_start + run.n <= (int32_t) cur.size());
        for (int32_t i = 0; i < run.n; ++i) {
            t.assert_equal("run matches both ranges",
                           old[run.old_start + i], cur[run.new_start + i]);
            t.assert_true("null token is never reused", old[run.old_start + i] != LLAMA_TOKEN_NULL);
        }
    }
}

int main() {
    testing t;

    t.test("default cap keeps 32 disjoint spans", [](testing & t) {
        const auto old = make_prompt(32, false);
        const auto cur = make_prompt(32, true);
        const auto plan = common_scr_plan_build(old, cur, {}, false, fixture_params());
        t.assert_equal(32, (int) plan.runs.size());
        check_exact_runs(t, old, cur, plan);
    });

    t.test("explicit cap six discards the rest", [](testing & t) {
        const auto old = make_prompt(70, false);
        const auto cur = make_prompt(70, true);
        const auto plan = common_scr_plan_build(old, cur, {}, false, fixture_params(6));
        t.assert_equal(6, (int) plan.runs.size());
        check_exact_runs(t, old, cur, plan);
    });

    t.test("default cap is exactly 64", [](testing & t) {
        const auto old = make_prompt(70, false);
        const auto cur = make_prompt(70, true);
        const auto plan = common_scr_plan_build(old, cur, {}, false, fixture_params());
        t.assert_equal(64, (int) plan.runs.size());
        check_exact_runs(t, old, cur, plan);
    });

    t.test("recurrent reuse ends at a real snapshot", [](testing & t) {
        const auto old = make_prompt(32, false);
        const auto cur = make_prompt(32, true);
        std::vector<int32_t> snapshots;
        for (int i = 0; i < 32; ++i) {
            snapshots.push_back(i * 11 + 11);
        }
        const auto plan = common_scr_plan_build(old, cur, snapshots, true, fixture_params());
        t.assert_true(!plan.runs.empty());
        t.assert_true(std::find(snapshots.begin(), snapshots.end(), plan.adopt) != snapshots.end());
        t.assert_equal(plan.adopt - plan.runs.back().old_start, plan.runs.back().n);
        check_exact_runs(t, old, cur, plan);
    });

    t.test("no snapshot means no recurrent reuse", [](testing & t) {
        const auto old = make_prompt(32, false);
        const auto cur = make_prompt(32, true);
        const auto plan = common_scr_plan_build(old, cur, {}, true, fixture_params());
        t.assert_true(plan.runs.empty());
        t.assert_equal(0, plan.n_reused());
    });

    t.test("null token breaks reuse", [](testing & t) {
        auto old = make_prompt(4, false);
        auto cur = make_prompt(4, true);
        old[1] = LLAMA_TOKEN_NULL;
        const auto plan = common_scr_plan_build(old, cur, {}, false, fixture_params());
        check_exact_runs(t, old, cur, plan);
        t.assert_true(!plan.runs.empty());
        t.assert_true(plan.n_reused() < 4 * 11);
    });

    return t.summary();
}