// prefix sharing for prompt cache states (tools/server/server-kv-share.cpp): a delta decodes to the original state
// byte for byte, shares only rows whose bytes equal the prefix's, and is refused against any other prefix
#include "../tools/server/server-kv-share.h"

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>

static int n_checks = 0;
static void check(bool ok, const char * what) {
    ++n_checks;
    if (!ok) {
        fprintf(stderr, "FAIL: %s\n", what);
        exit(1);
    }
}

template <typename T> static void put(std::vector<uint8_t> & b, T v) {
    const auto * p = reinterpret_cast<const uint8_t *>(&v);
    b.insert(b.end(), p, p + sizeof(T));
}

// a row's bytes for a cell at `pos` in layer `l`: a function of both, so equal positions give equal rows
static void row(std::vector<uint8_t> & b, size_t n, int32_t pos, uint32_t l, uint8_t salt) {
    for (size_t i = 0; i < n; ++i) {
        b.push_back(uint8_t(pos * 31 + l * 7 + i + salt));
    }
}

struct cell { int32_t pos; uint8_t salt; };

// a sequence state in llama_kv_cache::state_write's layout, then an opaque recurrent part
static std::vector<uint8_t> state(const std::vector<cell> & cells, bool mrope, bool v_trans, size_t tail, uint8_t tail_seed) {
    const uint32_t n_layer = 3;
    const size_t   k_row = 64, v_row = 48, el = 2;
    std::vector<uint8_t> b;
    put<uint32_t>(b, 0x4c4c414d); // magic
    put<int32_t>(b, 1);           // seq id
    put<uint32_t>(b, 1);          // n_stream
    put<uint32_t>(b, (uint32_t) cells.size());
    for (const auto & c : cells) {
        put<int32_t>(b, c.pos);
        put<uint32_t>(b, 1);
        if (mrope) {
            put<int32_t>(b, c.pos % 5);
            put<int32_t>(b, 0);
        }
        put<int32_t>(b, 1);
    }
    put<uint32_t>(b, v_trans ? 1 : 0);
    put<uint32_t>(b, n_layer);
    for (uint32_t l = 0; l < n_layer; ++l) {
        put<int32_t>(b, 1);
        put<uint64_t>(b, k_row);
        for (const auto & c : cells) {
            row(b, k_row, c.pos, l, c.salt);
        }
    }
    for (uint32_t l = 0; l < n_layer; ++l) {
        put<int32_t>(b, 1);
        if (!v_trans) {
            put<uint64_t>(b, v_row);
            for (const auto & c : cells) {
                row(b, v_row, c.pos, l + 100, c.salt);
            }
        } else {
            const uint32_t n_embd = v_row / el;
            put<uint32_t>(b, (uint32_t) el);
            put<uint32_t>(b, n_embd);
            for (uint32_t j = 0; j < n_embd; ++j) {
                for (const auto & c : cells) {
                    row(b, el, c.pos, l + 200 + j, c.salt);
                }
            }
        }
    }
    for (size_t i = 0; i < tail; ++i) {
        b.push_back(uint8_t(i * 13 + tail_seed));
    }
    return b;
}

static std::vector<cell> range(int32_t from, int32_t to, uint8_t salt = 0) {
    std::vector<cell> v;
    for (int32_t p = from; p < to; ++p) {
        v.push_back({ p, salt });
    }
    return v;
}

static bool round_trip(const std::vector<uint8_t> & child, const kv_share_parent & parent, const kv_share_layout & layout, size_t & shared) {
    std::vector<uint8_t> delta, back;
    std::string err;
    return kv_share_encode(child, parent, layout, delta, shared) && kv_share_is_delta(delta) &&
           kv_share_parent_id(delta) == parent.id && delta.size() < child.size() &&
           kv_share_decode(delta, parent, back, err) && back == child;
}

int main() {
    const int32_t P = 60000, N = 90000; // ~45 MiB of shared rows: above the 16 MiB floor
    for (int mode = 0; mode < 4; ++mode) {
        const bool mrope = mode & 1, v_trans = mode & 2;
        const kv_share_layout layout { mrope ? 8u : 0u };
        const auto parent = state(range(0, P), mrope, v_trans, 4096, 1);
        const kv_share_parent pp { 0x1234, &parent };

        // the conversation continues the prefix: its rows for the prefix's positions are shared, the rest and the
        // recurrent part are its own
        size_t shared = 0;
        const auto child = state(range(0, N), mrope, v_trans, 8192, 9);
        check(round_trip(child, pp, layout, shared), "a continuation decodes to itself");
        check(shared == (size_t) P * (64 + 48) * 3, "every prefix row is shared, nothing else");

        // cells in another order (a unified KV with two slots interleaves them by ubatch) match by position
        auto blocks = range(0, N);
        for (size_t i = 0; i + 1024 <= blocks.size(); i += 1024) {
            std::swap_ranges(blocks.begin() + i, blocks.begin() + i + 512, blocks.begin() + i + 512);
        }
        check(round_trip(state(blocks, mrope, v_trans, 100, 3), pp, layout, shared), "a reordered child decodes to itself");
        check(shared == (size_t) P * (64 + 48) * 3, "and shares the same rows");

        // shared rows scattered one cell at a time are not worth their copy records
        auto pairs = range(0, N);
        for (size_t i = 0; i + 1 < pairs.size(); i += 2) {
            std::swap(pairs[i], pairs[i + 1]);
        }
        std::vector<uint8_t> scattered;
        check(!kv_share_encode(state(pairs, mrope, v_trans, 100, 3), pp, layout, scattered, shared) ||
                  scattered.size() + (16u << 20) <= state(pairs, mrope, v_trans, 100, 3).size(),
              "a delta is only kept when it saves 16 MiB");

        // a prefix computed with other bits is stored literally where it differs
        auto drifted = range(0, N);
        for (int32_t p = 0; p < P; p += 10) {
            drifted[p].salt = 7;
        }
        // (a transposed V splits every cell into one element per row: a delta there costs more than it saves)
        if (!v_trans) {
            check(round_trip(state(drifted, mrope, v_trans, 100, 3), pp, layout, shared), "a child that differs in places decodes to itself");
            check(shared < (size_t) P * (64 + 48) * 3, "and shares only the equal rows");
        } else {
            std::vector<uint8_t> d;
            const auto whole = state(drifted, mrope, v_trans, 100, 3);
            check(!kv_share_encode(whole, pp, layout, d, shared) || d.size() + (16u << 20) <= whole.size(), "a scattered transposed delta is kept only if it saves 16 MiB");
        }

        // too little in common: no delta
        std::vector<uint8_t> delta;
        check(!kv_share_encode(state(range(0, 100), mrope, v_trans, 10, 3), pp, layout, delta, shared) && delta.empty(), "a short child is kept whole");
        check(!kv_share_encode(state(range(P, N), mrope, v_trans, 10, 3), pp, layout, delta, shared), "an unrelated child is kept whole");
    }

    // a delta only decodes against the prefix it was written against
    {
        const kv_share_layout layout { 0 };
        const auto parent = state(range(0, P), false, false, 64, 1);
        const auto child  = state(range(0, N), false, false, 64, 2);
        std::vector<uint8_t> delta, back;
        size_t shared = 0;
        std::string err;
        check(kv_share_encode(child, { 7, &parent }, layout, delta, shared), "encoded");
        check(!kv_share_decode(delta, { 8, &parent }, back, err), "another id is refused");
        auto other = parent;
        other[other.size() - 1] ^= 1;
        check(!kv_share_decode(delta, { 7, &other }, back, err), "other bytes under the same id are refused");
        check(!kv_share_decode(child, { 7, &parent }, back, err) && !kv_share_is_delta(child), "a whole state is not a delta");
        auto cut = delta;
        cut.resize(cut.size() - 5);
        check(!kv_share_decode(cut, { 7, &parent }, back, err), "a truncated delta is refused");
        check(!kv_share_encode(delta, { 7, &parent }, layout, back, shared), "a delta is not encoded again");
    }

    // positions twice (image chunks share one): nothing ambiguous is shared
    {
        const kv_share_layout layout { 0 };
        auto cells = range(0, P);
        cells.push_back({ 5, 0 });
        const auto parent = state(cells, false, false, 64, 1);
        std::vector<uint8_t> delta;
        size_t shared = 0;
        check(!kv_share_encode(state(range(0, N), false, false, 64, 2), { 1, &parent }, layout, delta, shared), "a prefix with a repeated position shares nothing");
    }

    printf("test-kv-share: %d checks passed\n", n_checks);
    return 0;
}