#include "scr.h"

#include "../src/llama-ext.h"

#include <algorithm>
#include <unordered_map>

int32_t common_scr_plan::resume() const {
    return runs.empty() ? restore : runs.back().new_start + runs.back().n;
}

int32_t common_scr_plan::n_reused() const {
    int32_t n = 0;
    for (const auto & r : runs) {
        n += r.n;
    }
    return n;
}

std::string common_scr_plan::str() const {
    std::string s = "prefix " + std::to_string(n_prefix) + ", restore " + std::to_string(restore) + ", adopt " + std::to_string(adopt);
    for (const auto & r : runs) {
        s += ", R " + std::to_string(r.new_start) + " <- " + std::to_string(r.old_start) + " x " + std::to_string(r.n);
    }
    return s;
}

// greedy maximal spans of cur[from, ...) found in old[from, ...), seeded by n-gram hashes, each old token used once
static std::vector<common_scr_run> scr_find_spans(
        const std::vector<llama_token> & old,
        const std::vector<llama_token> & cur,
        int32_t from,
        const common_scr_params & params) {
    std::vector<common_scr_run> spans;

    const int32_t n_old = (int32_t) old.size();
    const int32_t n_cur = (int32_t) cur.size();
    const int32_t k     = params.seed;

    if (n_old - from < params.min_span || n_cur - from < params.min_span) {
        return spans;
    }

    auto gram = [k](const std::vector<llama_token> & v, int32_t i) {
        uint64_t h = 1469598103934665603ULL;
        for (int32_t j = 0; j < k; ++j) {
            h = (h ^ (uint32_t) v[i + j]) * 1099511628211ULL;
        }
        return h;
    };

    // the first few old positions of every seed n-gram
    const size_t max_bucket = 16;
    std::unordered_map<uint64_t, std::vector<int32_t>> index;
    index.reserve(n_old - from);
    for (int32_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(n_old, false);

    for (int32_t p = from; p + k <= n_cur; ) {
        const auto it = index.find(gram(cur, p));

        int32_t best_o = -1;
        int32_t best_n = 0;
        if (it != index.end()) {
            for (const int32_t o : it->second) {
                int32_t n = 0;
                while (p + n < n_cur && o + n < n_old && !used[o + n] &&
                        cur[p + n] == old[o + n] && cur[p + n] != LLAMA_TOKEN_NULL) {
                    ++n;
                }
                if (n > best_n) {
                    best_n = n;
                    best_o = o;
                }
            }
        }

        if (best_n >= params.min_span) {
            spans.push_back({ p, best_o, best_n });
            std::fill(used.begin() + best_o, used.begin() + best_o + best_n, true);
            p += best_n;
        } else {
            ++p;
        }
    }

    return spans;
}

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) {
    common_scr_plan plan;

    const int32_t n_cur = (int32_t) cur.size();

    int32_t a = 0;
    while (a < (int32_t) old.size() && a < n_cur && old[a] == cur[a] && cur[a] != LLAMA_TOKEN_NULL) {
        ++a;
    }
    plan.n_prefix = a;
    plan.restore  = a;

    if (recurrent) {
        plan.restore = 0;
        for (const int32_t q : snapshots) {
            if (q <= a) {
                plan.restore = std::max(plan.restore, q);
            }
        }
    }

    auto spans = scr_find_spans(old, cur, a, params);

    // the largest K spans, back in prompt order
    std::sort(spans.begin(), spans.end(), [](const auto & x, const auto & y) { return x.n > y.n; });
    if ((int32_t) spans.size() > params.max_runs) {
        spans.resize(params.max_runs);
    }
    std::sort(spans.begin(), spans.end(), [](const auto & x, const auto & y) { return x.new_start < y.new_start; });

    // the last span must end at a snapshot, which also ends the reuse: everything after it is prefilled from that state
    while (!spans.empty()) {
        auto & last = spans.back();

        int32_t end = -1;
        if (recurrent) {
            for (const int32_t q : snapshots) {
                const int32_t n = q - last.old_start;
                if (n >= params.min_span && n <= last.n && last.new_start + n < n_cur) {
                    end = std::max(end, q);
                }
            }
        } else {
            end = last.old_start + std::min(last.n, n_cur - 1 - last.new_start);
        }

        if (end - last.old_start >= params.min_span) {
            last.n     = end - last.old_start;
            plan.adopt = recurrent ? end : -1;
            break;
        }

        spans.pop_back();
    }

    plan.runs = std::move(spans);

    return plan;
}

int32_t common_scr_apply_kv(llama_context * ctx, llama_seq_id seq_id, const common_scr_plan & plan) {
    // a run may move once its target holds no source of a run still to move
    std::vector<common_scr_run> order;
    {
        std::vector<common_scr_run> pending = plan.runs;
        while (!pending.empty()) {
            auto next = std::find_if(pending.begin(), pending.end(), [&](const common_scr_run & r) {
                return std::none_of(pending.begin(), pending.end(), [&](const common_scr_run & o) {
                    return &o != &r && o.old_start < r.new_start + r.n && r.new_start < o.old_start + o.n;
                });
            });
            if (next == pending.end()) {
                return 0;
            }
            order.push_back(*next);
            pending.erase(next);
        }
    }

    std::vector<common_scr_run> sources = plan.runs;
    std::sort(sources.begin(), sources.end(), [](const auto & x, const auto & y) { return x.old_start < y.old_start; });

    llama_pos pos = plan.restore;
    for (const auto & s : sources) {
        if (s.old_start > pos && !llama_scr_kv_rm(ctx, seq_id, pos, s.old_start)) {
            return -1;
        }
        pos = std::max(pos, (llama_pos) (s.old_start + s.n));
    }
    if (!llama_scr_kv_rm(ctx, seq_id, pos, -1)) {
        return -1;
    }

    for (const auto & r : order) {
        if (!llama_scr_kv_relocate(ctx, seq_id, r.old_start, r.old_start + r.n, r.new_start - r.old_start)) {
            return -1;
        }
    }

    return 1;
}