#pragma once

#include "llama.h"

#include <cstdint>
#include <string>
#include <vector>

//
// suffix cache reuse (SCR)
//
// after an edit in the middle of a prompt, the cached attention KV of the text that survived is moved to its new
// positions (K re-rotated), the recurrent state of a hybrid model continues from a snapshot taken inside the last
// moved span, and only the rest is prefilled. the moved rows were computed over the pre-edit text, so the result
// approximates a fresh prefill of the edited prompt.
//

struct common_scr_params {
    int32_t min_span = 256; // shortest span worth moving
    int32_t max_runs = 64;  // most spans moved per request
    int32_t seed     = 32;  // n-gram length that seeds a match
};

// a span of the cached tokens [old_start, old_start + n) that reappears at [new_start, new_start + n)
struct common_scr_run {
    int32_t new_start;
    int32_t old_start;
    int32_t n;
};

struct common_scr_plan {
    int32_t n_prefix = 0;  // leading tokens equal in both prompts
    int32_t restore  = 0;  // position the prefill restarts from: a snapshot at or before n_prefix
    int32_t adopt    = -1; // old position of the snapshot the recurrent state continues from, -1 without one

    std::vector<common_scr_run> runs; // in new order, all at or after n_prefix

    bool    empty()    const { return runs.empty(); }
    int32_t resume()   const; // new position the final prefill starts at
    int32_t n_reused() const;

    std::string str() const;
};

// plans the reuse of a state that holds old[0, n_old) for the prompt cur.
// snapshots: sorted old positions q at which the recurrent state after old[0, q) can be restored (the live state
// counts as old.size()); with recurrent = false every position can be resumed from.
// LLAMA_TOKEN_NULL never matches. the plan always leaves at least one token of cur to prefill.
common_scr_plan common_scr_plan_build(
        const std::vector<llama_token> & old,
        const std::vector<llama_token> & cur,
        const std::vector<int32_t>     & snapshots,
        bool                             recurrent,
        const common_scr_params        & params);

// rearranges the attention KV of seq_id for a plan: drops every cell at or after plan.restore that no run moves,
// then moves the runs. returns 1 when done, 0 when the runs cannot be moved in any order (nothing changed),
// -1 when a move failed after the cache was changed (the sequence must then be cleared).
int32_t common_scr_apply_kv(llama_context * ctx, llama_seq_id seq_id, const common_scr_plan & plan);