// Suffix cache reuse (SCR): after an edit in the middle of a prompt, the K/V of the text that survived the edit is
// moved to its new positions instead of being read again.  Ported from llama-cpp-bonsai's common/scr.cpp.
//
// The moved cells were computed over the text before the edit, so the result approximates a fresh read.  The GDN
// layers have no cells to move: their state continues from a conversation checkpoint at the end of the last moved
// run (the plan's `adopt`), and everything after that run is read from it.  Strata's indexer pools 4 cells into one
// key and keeps only the pooled keys, so a run starts and ends on a block boundary of the OLD prompt and moves by a
// multiple of 4: up to 3 pad tokens are read before a run to make its shift one (DeepSeek V4.1's recipe pads to an
// even shift the same way).  The moved blocks then keep their cells, and their pooled keys only turn by the shift.
#pragma once

#include <algorithm>
#include <cstdint>
#include <string>
#include <unordered_map>
#include <vector>

namespace strata::core {

struct ScrParams {
    int64_t min_span = 256;   // shortest run worth moving
    int64_t max_runs = 64;    // most runs moved per request
    int64_t seed = 32;        // n-gram length that seeds a match
    int64_t block = 4;        // the indexer's cells per pooled key: read parts start and end on these
};

// the cells old[old_start, old_start + n) reappear at [new_start, new_start + n) of the prompt, after `pad` pad tokens
// that the read inserts right before the run; `at` is where the run lands in the padded (physical) prompt
struct ScrRun {
    int64_t new_start = 0, old_start = 0, n = 0;
    int64_t pad = 0, at = 0;
    int64_t shift() const { return at - old_start; }
};

struct ScrPlan {
    int64_t n_prefix = 0;   // leading tokens equal in both prompts
    int64_t restore = 0;    // the checkpoint the read restarts from (at or before n_prefix)
    int64_t adopt = -1;     // old length of the checkpoint the GDN state continues from after the last run
    std::vector<ScrRun> runs;   // in new order, at or after restore

    bool empty() const { return runs.empty(); }
    // the prompt position the final read starts at, and the same in the padded prompt
    int64_t resume() const { return runs.empty() ? restore : runs.back().new_start + runs.back().n; }
    int64_t resume_at() const { return runs.empty() ? restore : runs.back().at + runs.back().n; }
    int64_t pads() const {
        int64_t n = 0;
        for (const ScrRun& r : runs) n += r.pad;
        return n;
    }
    int64_t reused() const {
        int64_t n = 0;
        for (const ScrRun& r : runs) n += r.n;
        return n;
    }
    std::string str() const {
        std::string s = "prefix " + std::to_string(n_prefix) + ", restore " + std::to_string(restore) + ", adopt " +
                        std::to_string(adopt);
        for (const ScrRun& r : runs)
            s += ", R " + std::to_string(r.new_start) + " <- " + std::to_string(r.old_start) + " x " + std::to_string(r.n) +
                 (r.pad ? " pad " + std::to_string(r.pad) : "");
        return s;
    }
};

namespace scr_detail {
// greedy maximal spans of cur[from, ...) found in old[from, ...), seeded by n-gram hashes, each old token used once
template <class A, class B>
std::vector<ScrRun> find_spans(const std::vector<A>& old, const std::vector<B>& cur, int64_t from, const ScrParams& p) {
    std::vector<ScrRun> spans;
    const int64_t n_old = (int64_t) old.size(), n_cur = (int64_t) cur.size(), k = p.seed;
    if (n_old - from < p.min_span || n_cur - from < p.min_span) return spans;
    auto gram = [k](const auto& v, int64_t i) {
        uint64_t h = 1469598103934665603ULL;
        for (int64_t j = 0; j < k; ++j) h = (h ^ (uint32_t) v[(size_t) (i + j)]) * 1099511628211ULL;
        return h;
    };
    constexpr size_t max_bucket = 16;   // the first few old positions of every seed n-gram
    std::unordered_map<uint64_t, std::vector<int64_t>> index;
    index.reserve((size_t) (n_old - from));
    for (int64_t o = from; o + k <= n_old; ++o) {
        auto& b = index[gram(old, o)];
        if (b.size() < max_bucket) b.push_back(o);
    }
    std::vector<bool> used((size_t) n_old, false);
    for (int64_t q = from; q + k <= n_cur;) {
        const auto it = index.find(gram(cur, q));
        int64_t best_o = -1, best_n = 0;
        if (it != index.end())
            for (const int64_t o : it->second) {
                int64_t n = 0;
                while (q + n < n_cur && o + n < n_old && !used[(size_t) (o + n)] &&
                       (int64_t) cur[(size_t) (q + n)] == (int64_t) old[(size_t) (o + n)])
                    ++n;
                if (n > best_n) { best_n = n; best_o = o; }
            }
        if (best_n >= p.min_span) {
            spans.push_back({q, best_o, best_n});
            std::fill(used.begin() + best_o, used.begin() + best_o + best_n, true);
            q += best_n;
        } else {
            ++q;
        }
    }
    return spans;
}
inline int64_t round_up(int64_t x, int64_t b) { return (x + b - 1) / b * b; }
inline int64_t round_down(int64_t x, int64_t b) { return x / b * b; }
}  // namespace scr_detail

// Plans the reuse of a session that holds old[0, n_old) for the prompt `cur`.  `checkpoints`: the old lengths whose
// GDN state can be restored (the live state counts as old.size()).  A run starts on a block boundary of the new
// prompt; every run but the last also ends on one, and the last one ends at a checkpoint.  The plan always leaves at
// least one token of `cur` to read.
template <class A, class B>
ScrPlan scr_plan(const std::vector<A>& old, const std::vector<B>& cur, std::vector<int64_t> checkpoints,
                 const ScrParams& p) {
    using namespace scr_detail;
    ScrPlan plan;
    const int64_t n_cur = (int64_t) cur.size(), bl = std::max<int64_t>(p.block, 1);
    std::sort(checkpoints.begin(), checkpoints.end());
    int64_t a = 0;
    while (a < (int64_t) old.size() && a < n_cur && (int64_t) old[(size_t) a] == (int64_t) cur[(size_t) a]) ++a;
    plan.n_prefix = a;
    for (const int64_t q : checkpoints)
        if (q <= a) plan.restore = std::max(plan.restore, q);
    auto spans = find_spans(old, cur, a, p);
    // a run starts on a block boundary of the old prompt and (but the last) ends on one
    for (ScrRun& s : spans) {
        const int64_t cut = round_up(s.old_start, bl) - s.old_start;
        s.new_start += cut;
        s.old_start += cut;
        s.n -= cut;
        s.n = round_down(s.old_start + s.n, bl) - s.old_start;
    }
    spans.erase(std::remove_if(spans.begin(), spans.end(), [&](const ScrRun& s) { return s.n < p.min_span; }),
                spans.end());
    // the largest runs, back in prompt order
    std::sort(spans.begin(), spans.end(), [](const ScrRun& x, const ScrRun& y) { return x.n > y.n; });
    if ((int64_t) spans.size() > p.max_runs) spans.resize((size_t) p.max_runs);
    std::sort(spans.begin(), spans.end(), [](const ScrRun& x, const ScrRun& y) { return x.new_start < y.new_start; });
    // the last run ends at a checkpoint, which also ends the reuse: everything after it is read from that state.
    // It was cut back to a block boundary, so a checkpoint up to block - 1 cells past it may still match.
    while (!spans.empty()) {
        ScrRun& last = spans.back();
        int64_t end = -1;
        for (const int64_t q : checkpoints) {
            const int64_t n = q - last.old_start;
            if (n < p.min_span || n >= last.n + bl || last.new_start + n >= n_cur || q <= end) continue;
            if (n <= last.n || std::equal(old.begin() + last.old_start + last.n, old.begin() + q,
                                          cur.begin() + last.new_start + last.n,
                                          [](const A& x, const B& y) { return (int64_t) x == (int64_t) y; }))
                end = q;
        }
        if (end >= 0) {
            last.n = end - last.old_start;
            plan.adopt = end;
            break;
        }
        spans.pop_back();
    }
    // the pads: walking the runs in prompt order, each gets what makes its padded position a multiple of the block
    // away from its old one
    int64_t pads = 0;
    for (ScrRun& s : spans) {
        s.pad = ((s.old_start - (s.new_start + pads)) % bl + bl) % bl;
        pads += s.pad;
        s.at = s.new_start + pads;
    }
    plan.runs = std::move(spans);
    return plan;
}

// The pads an earlier splice read (`pads`: their positions in `live`, sorted) that the prompt `cur` still shares with
// the session: the ones before the last token of `cur` that `live` holds at the same place once the pads are
// skipped.  `out` is `cur` with those pads put back, which makes it start with `live` as far as they agree.
template <class A, class B>
std::vector<int64_t> scr_carry(const std::vector<B>& cur, const std::vector<A>& live, const std::vector<int64_t>& pads,
                               int64_t pad_token, std::vector<B>& out) {
    size_t p = 0, i = 0, k = 0, kept = 0;
    while (p < live.size()) {
        if (k < pads.size() && pads[k] == (int64_t) p) {
            if ((int64_t) live[p] != pad_token) break;
            ++k;
            ++p;
            continue;
        }
        if (i + 1 >= cur.size() || (int64_t) live[p] != (int64_t) cur[i]) break;
        ++i;
        ++p;
        kept = k;
    }
    std::vector<int64_t> carried(pads.begin(), pads.begin() + (long) kept);
    out.clear();
    out.reserve(cur.size() + carried.size());
    size_t c = 0;
    for (const int64_t at : carried) {
        while ((int64_t) out.size() < at) out.push_back(cur[c++]);
        out.push_back((B) pad_token);
    }
    out.insert(out.end(), cur.begin() + (long) c, cur.end());
    return carried;
}

// The positions of the pads `plan` inserts, in the padded prompt.
inline std::vector<int64_t> scr_pad_positions(const ScrPlan& plan) {
    std::vector<int64_t> at;
    for (const ScrRun& r : plan.runs)
        for (int64_t j = r.pad; j > 0; --j) at.push_back(r.at - j);
    return at;
}

// The prompt the engine reads under `plan`: `cur` with each run's pads inserted right before it.
template <class B>
std::vector<B> scr_padded(const std::vector<B>& cur, const ScrPlan& plan, B pad_token) {
    std::vector<B> out;
    out.reserve(cur.size() + (size_t) plan.pads());
    int64_t from = 0;
    for (const ScrRun& r : plan.runs) {
        out.insert(out.end(), cur.begin() + from, cur.begin() + r.new_start);
        out.insert(out.end(), (size_t) r.pad, pad_token);
        from = r.new_start;
    }
    out.insert(out.end(), cur.begin() + from, cur.end());
    return out;
}

}  // namespace strata::core