#include "llama.h"
#include "../src/llama-token-cache.h"

#include <chrono>
#include <filesystem>
#include <string>
#include <vector>

#undef NDEBUG
#include <cassert>

int main() {
    namespace fs = std::filesystem;

    const auto defaults = llama_token_cache_default_params();
    assert(defaults.capacity_bytes == 5ull * 1024 * 1024 * 1024);

    const fs::path root = fs::temp_directory_path() /
        ("llama-token-cache-test-" + std::to_string(
            std::chrono::high_resolution_clock::now().time_since_epoch().count()));
    const fs::path db = root / "cache.sqlite3";

    llama_token_cache_params params = {
        nullptr,
        128,
    };
    // configure copies the path; it must not retain the temporary c_str().
    const std::string db_string = db.string();
    params.path = db_string.c_str();
    assert(llama_token_cache_configure(params));
    assert(llama_token_cache_clear());

    assert(llama_token_cache_hash("abc") ==
        "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad");

    const std::vector<uint8_t> first(80, 0x11);
    const std::vector<uint8_t> second(80, 0x22);
    std::vector<uint8_t> loaded;

    llama_token_cache_store(LLAMA_TOKEN_CACHE_KIND_TOKENIZER, "first", first, true);
    assert(llama_token_cache_lookup(LLAMA_TOKEN_CACHE_KIND_TOKENIZER, "first", loaded));
    assert(loaded == first);

    llama_token_cache_store(LLAMA_TOKEN_CACHE_KIND_TEXT, "second", second, true);
    loaded.clear();
    assert(!llama_token_cache_lookup(LLAMA_TOKEN_CACHE_KIND_TOKENIZER, "first", loaded));
    assert(llama_token_cache_lookup(LLAMA_TOKEN_CACHE_KIND_TEXT, "second", loaded));
    assert(loaded == second);

    const std::vector<uint8_t> async_payload(16, 0x33);
    llama_token_cache_store(LLAMA_TOKEN_CACHE_KIND_TEXT, "async", async_payload);
    assert(llama_token_cache_flush());
    loaded.clear();
    assert(llama_token_cache_lookup(LLAMA_TOKEN_CACHE_KIND_TEXT, "async", loaded));
    assert(loaded == async_payload);

    const auto stats = llama_token_cache_get_stats();
    assert(stats.entries == 2);
    assert(stats.bytes_used == second.size() + async_payload.size());
    assert(stats.evictions >= 1);
    assert(stats.writes >= 3);

    assert(llama_token_cache_clear());
    assert(llama_token_cache_get_stats().entries == 0);

    std::error_code ec;
    fs::remove_all(root, ec);
    return 0;
}