#include "llama-token-cache.h"

#include "llama-impl.h"

#include <sqlite3.h>

#include <algorithm>
#include <array>
#include <atomic>
#include <condition_variable>
#include <chrono>
#include <cstdlib>
#include <cstring>
#include <deque>
#include <filesystem>
#include <limits>
#include <mutex>
#include <thread>

#if defined(_WIN32)
#  include <windows.h>
#else
#  include <sys/stat.h>
#  include <unistd.h>
#endif

namespace fs = std::filesystem;

namespace {

constexpr uint64_t DEFAULT_CAPACITY = 5ull * 1024 * 1024 * 1024;
constexpr uint64_t MAX_QUEUE_BYTES  = 64ull * 1024 * 1024;
constexpr int      SCHEMA_VERSION   = 1;

struct sha256 {
    uint32_t h[8] = {
        0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a,
        0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19,
    };
    uint8_t block[64] = {};
    uint64_t bits = 0;
    size_t used = 0;

    static uint32_t rotr(uint32_t x, uint32_t n) { return (x >> n) | (x << (32 - n)); }

    void compress(const uint8_t * p) {
        static const uint32_t k[64] = {
            0x428a2f98,0x71374491,0xb5c0fbcf,0xe9b5dba5,0x3956c25b,0x59f111f1,0x923f82a4,0xab1c5ed5,
            0xd807aa98,0x12835b01,0x243185be,0x550c7dc3,0x72be5d74,0x80deb1fe,0x9bdc06a7,0xc19bf174,
            0xe49b69c1,0xefbe4786,0x0fc19dc6,0x240ca1cc,0x2de92c6f,0x4a7484aa,0x5cb0a9dc,0x76f988da,
            0x983e5152,0xa831c66d,0xb00327c8,0xbf597fc7,0xc6e00bf3,0xd5a79147,0x06ca6351,0x14292967,
            0x27b70a85,0x2e1b2138,0x4d2c6dfc,0x53380d13,0x650a7354,0x766a0abb,0x81c2c92e,0x92722c85,
            0xa2bfe8a1,0xa81a664b,0xc24b8b70,0xc76c51a3,0xd192e819,0xd6990624,0xf40e3585,0x106aa070,
            0x19a4c116,0x1e376c08,0x2748774c,0x34b0bcb5,0x391c0cb3,0x4ed8aa4a,0x5b9cca4f,0x682e6ff3,
            0x748f82ee,0x78a5636f,0x84c87814,0x8cc70208,0x90befffa,0xa4506ceb,0xbef9a3f7,0xc67178f2,
        };
        uint32_t w[64];
        for (int i = 0; i < 16; ++i) {
            w[i] = uint32_t(p[4*i]) << 24 | uint32_t(p[4*i+1]) << 16 | uint32_t(p[4*i+2]) << 8 | p[4*i+3];
        }
        for (int i = 16; i < 64; ++i) {
            const uint32_t s0 = rotr(w[i-15], 7) ^ rotr(w[i-15], 18) ^ (w[i-15] >> 3);
            const uint32_t s1 = rotr(w[i-2], 17) ^ rotr(w[i-2], 19) ^ (w[i-2] >> 10);
            w[i] = w[i-16] + s0 + w[i-7] + s1;
        }
        uint32_t a=h[0],b=h[1],c=h[2],d=h[3],e=h[4],f=h[5],g=h[6],hh=h[7];
        for (int i = 0; i < 64; ++i) {
            const uint32_t s1 = rotr(e,6)^rotr(e,11)^rotr(e,25);
            const uint32_t ch = (e&f)^((~e)&g);
            const uint32_t t1 = hh+s1+ch+k[i]+w[i];
            const uint32_t s0 = rotr(a,2)^rotr(a,13)^rotr(a,22);
            const uint32_t maj = (a&b)^(a&c)^(b&c);
            const uint32_t t2 = s0+maj;
            hh=g; g=f; f=e; e=d+t1; d=c; c=b; b=a; a=t1+t2;
        }
        h[0]+=a;h[1]+=b;h[2]+=c;h[3]+=d;h[4]+=e;h[5]+=f;h[6]+=g;h[7]+=hh;
    }

    void update(const void * data, size_t n) {
        const auto * p = static_cast<const uint8_t *>(data);
        bits += uint64_t(n) * 8;
        while (n) {
            const size_t take = std::min(n, sizeof(block) - used);
            memcpy(block + used, p, take);
            used += take; p += take; n -= take;
            if (used == sizeof(block)) { compress(block); used = 0; }
        }
    }

    std::array<uint8_t, 32> finish() {
        const uint64_t original_bits = bits;
        const uint8_t one = 0x80;
        update(&one, 1);
        const uint8_t zero = 0;
        while (used != 56) update(&zero, 1);
        uint8_t len[8];
        for (int i = 0; i < 8; ++i) len[7-i] = uint8_t(original_bits >> (8*i));
        update(len, 8);
        std::array<uint8_t, 32> out;
        for (int i = 0; i < 8; ++i) {
            out[4*i]=uint8_t(h[i]>>24); out[4*i+1]=uint8_t(h[i]>>16);
            out[4*i+2]=uint8_t(h[i]>>8); out[4*i+3]=uint8_t(h[i]);
        }
        return out;
    }
};

struct queued_write {
    llama_token_cache_kind kind;
    std::string key;
    std::vector<uint8_t> payload;
};

struct cache_counters {
    std::atomic<uint64_t> tokenizer_hits{0};
    std::atomic<uint64_t> tokenizer_misses{0};
    std::atomic<uint64_t> text_hits{0};
    std::atomic<uint64_t> text_misses{0};
    std::atomic<uint64_t> writes{0};
    std::atomic<uint64_t> dropped_writes{0};
    std::atomic<uint64_t> evictions{0};
    std::atomic<uint64_t> errors{0};
};

class token_cache {
public:
    token_cache() {
        params.capacity_bytes = DEFAULT_CAPACITY;
        worker = std::thread([this] { run(); });
    }

    ~token_cache() {
        flush();
        {
            std::lock_guard<std::mutex> lock(mutex);
            stopping = true;
        }
        ready.notify_all();
        if (worker.joinable()) worker.join();
        close();
    }

    bool configure(llama_token_cache_params value) {
        std::lock_guard<std::mutex> lock(mutex);
        if (initialized) return false;
        params.capacity_bytes = value.capacity_bytes;
        configured_path = value.path ? value.path : "";
        return true;
    }

    bool lookup(llama_token_cache_kind kind, const std::string & key, std::vector<uint8_t> & out) {
        if (!ensure_open()) return false;
        std::lock_guard<std::mutex> db_lock(db_mutex);
        sqlite3_stmt * stmt = nullptr;
        const char * sql = "SELECT payload FROM entries WHERE kind=?1 AND key=?2";
        if (sqlite3_prepare_v2(db, sql, -1, &stmt, nullptr) != SQLITE_OK) return db_error();
        sqlite3_bind_int(stmt, 1, kind);
        sqlite3_bind_text(stmt, 2, key.data(), int(key.size()), SQLITE_TRANSIENT);
        bool found = false;
        if (sqlite3_step(stmt) == SQLITE_ROW) {
            const int n = sqlite3_column_bytes(stmt, 0);
            const void * p = sqlite3_column_blob(stmt, 0);
            if (n >= 0 && (n == 0 || p != nullptr)) {
                out.assign(static_cast<const uint8_t *>(p), static_cast<const uint8_t *>(p) + n);
                found = true;
            }
        }
        sqlite3_finalize(stmt);
        if (found) {
            sqlite3_stmt * touch = nullptr;
            if (sqlite3_prepare_v2(db, "UPDATE entries SET last_access=?1 WHERE kind=?2 AND key=?3", -1, &touch, nullptr) == SQLITE_OK) {
                sqlite3_bind_int64(touch, 1, now_tick());
                sqlite3_bind_int(touch, 2, kind);
                sqlite3_bind_text(touch, 3, key.data(), int(key.size()), SQLITE_TRANSIENT);
                sqlite3_step(touch);
            }
            sqlite3_finalize(touch);
            if (kind == LLAMA_TOKEN_CACHE_KIND_TEXT) ++stats.text_hits;
        } else if (kind == LLAMA_TOKEN_CACHE_KIND_TEXT) {
            ++stats.text_misses;
        }
        return found;
    }

    void store(llama_token_cache_kind kind, const std::string & key, const std::vector<uint8_t> & payload, bool sync) {
        if (params.capacity_bytes == 0 || payload.size() > params.capacity_bytes ||
                payload.size() > static_cast<size_t>(std::numeric_limits<int>::max())) return;
        if (sync) {
            write_one({kind, key, payload});
            return;
        }
        std::lock_guard<std::mutex> lock(mutex);
        if (queue_bytes + payload.size() > MAX_QUEUE_BYTES) {
            ++stats.dropped_writes;
            return;
        }
        queue.push_back({kind, key, payload});
        queue_bytes += payload.size();
        ready.notify_one();
    }

    void flush() {
        std::unique_lock<std::mutex> lock(mutex);
        drained.wait(lock, [&] { return queue.empty() && !writing; });
        if (db) {
            lock.unlock();
            std::lock_guard<std::mutex> db_lock(db_mutex);
            sqlite3_wal_checkpoint_v2(db, nullptr, SQLITE_CHECKPOINT_PASSIVE, nullptr, nullptr);
        }
    }

    bool clear_all() {
        flush();
        if (!ensure_open()) return params.capacity_bytes == 0;
        std::lock_guard<std::mutex> db_lock(db_mutex);
        if (sqlite3_exec(db, "DELETE FROM entries; PRAGMA incremental_vacuum;", nullptr, nullptr, nullptr) != SQLITE_OK) return db_error();
        return true;
    }

    llama_token_cache_stats get_stats() {
        llama_token_cache_stats result = {
            stats.tokenizer_hits.load(),
            stats.tokenizer_misses.load(),
            stats.text_hits.load(),
            stats.text_misses.load(),
            stats.writes.load(),
            stats.dropped_writes.load(),
            stats.evictions.load(),
            stats.errors.load(),
            0,
            0,
        };
        if (ensure_open()) {
            std::lock_guard<std::mutex> db_lock(db_mutex);
            sqlite3_stmt * stmt = nullptr;
            if (sqlite3_prepare_v2(db, "SELECT COUNT(*), COALESCE(SUM(size),0) FROM entries", -1, &stmt, nullptr) == SQLITE_OK &&
                    sqlite3_step(stmt) == SQLITE_ROW) {
                result.entries = sqlite3_column_int64(stmt, 0);
                result.bytes_used = sqlite3_column_int64(stmt, 1);
            }
            sqlite3_finalize(stmt);
        }
        return result;
    }

    void note_tokenizer(bool hit) {
        if (hit) ++stats.tokenizer_hits; else ++stats.tokenizer_misses;
    }

private:
    static int64_t now_tick() {
        return std::chrono::duration_cast<std::chrono::microseconds>(
            std::chrono::system_clock::now().time_since_epoch()).count();
    }

    std::string default_path() {
        if (const char * p = std::getenv("LLAMA_CACHE"); p && *p) return (fs::path(p) / "token-cache.sqlite3").string();
#if defined(_WIN32)
        if (const char * p = std::getenv("LOCALAPPDATA"); p && *p) return (fs::path(p) / "llama.cpp" / "token-cache.sqlite3").string();
#elif defined(__APPLE__)
        if (const char * p = std::getenv("HOME"); p && *p) return (fs::path(p) / "Library" / "Caches" / "llama.cpp" / "token-cache.sqlite3").string();
#else
        if (const char * p = std::getenv("XDG_CACHE_HOME"); p && *p) return (fs::path(p) / "llama.cpp" / "token-cache.sqlite3").string();
        if (const char * p = std::getenv("HOME"); p && *p) return (fs::path(p) / ".cache" / "llama.cpp" / "token-cache.sqlite3").string();
#endif
        return {};
    }

    bool ensure_open() {
        std::lock_guard<std::mutex> init_lock(init_mutex);
        if (initialized) return db != nullptr;
        initialized = true;
        if (params.capacity_bytes == 0) return false;
        path = configured_path.empty() ? default_path() : configured_path;
        if (path.empty()) return false;
        std::error_code ec;
        fs::create_directories(fs::path(path).parent_path(), ec);
        if (ec) return false;
#if !defined(_WIN32)
        chmod(fs::path(path).parent_path().c_str(), 0700);
#endif
        if (sqlite3_open_v2(path.c_str(), &db, SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE | SQLITE_OPEN_FULLMUTEX, nullptr) != SQLITE_OK) {
            warn_once("failed to open token cache database");
            close();
            return false;
        }
#if !defined(_WIN32)
        chmod(path.c_str(), 0600);
#endif
        sqlite3_busy_timeout(db, 5000);
        const char * schema =
            "PRAGMA journal_mode=WAL;"
            "PRAGMA synchronous=NORMAL;"
            "PRAGMA auto_vacuum=INCREMENTAL;"
            "CREATE TABLE IF NOT EXISTS meta(key TEXT PRIMARY KEY,value INTEGER NOT NULL);"
            "CREATE TABLE IF NOT EXISTS entries("
            " kind INTEGER NOT NULL,key TEXT NOT NULL,payload BLOB NOT NULL,size INTEGER NOT NULL,last_access INTEGER NOT NULL,"
            " PRIMARY KEY(kind,key));"
            "CREATE INDEX IF NOT EXISTS entries_lru ON entries(last_access);";
        if (sqlite3_exec(db, schema, nullptr, nullptr, nullptr) != SQLITE_OK) {
            warn_once("failed to initialize token cache database");
            close();
            return false;
        }
        sqlite3_stmt * stmt = nullptr;
        sqlite3_prepare_v2(db, "INSERT OR REPLACE INTO meta(key,value) VALUES('schema',?1)", -1, &stmt, nullptr);
        sqlite3_bind_int(stmt, 1, SCHEMA_VERSION);
        sqlite3_step(stmt);
        sqlite3_finalize(stmt);
        sqlite3_prepare_v2(db, "INSERT OR REPLACE INTO meta(key,value) VALUES('capacity',?1)", -1, &stmt, nullptr);
        sqlite3_bind_int64(stmt, 1, static_cast<sqlite3_int64>(params.capacity_bytes));
        sqlite3_step(stmt);
        sqlite3_finalize(stmt);
        sqlite3_exec(db, "BEGIN IMMEDIATE", nullptr, nullptr, nullptr);
        while (true) {
            sqlite3_prepare_v2(db, "SELECT COALESCE(SUM(size),0) FROM entries", -1, &stmt, nullptr);
            uint64_t total = 0;
            if (sqlite3_step(stmt) == SQLITE_ROW) total = sqlite3_column_int64(stmt, 0);
            sqlite3_finalize(stmt);
            if (total <= params.capacity_bytes) break;
            if (sqlite3_exec(db,
                    "DELETE FROM entries WHERE rowid=(SELECT rowid FROM entries ORDER BY last_access LIMIT 1)",
                    nullptr, nullptr, nullptr) != SQLITE_OK) {
                break;
            }
            ++stats.evictions;
        }
        sqlite3_exec(db, "COMMIT", nullptr, nullptr, nullptr);
        return true;
    }

    void run() {
        while (true) {
            queued_write item;
            {
                std::unique_lock<std::mutex> lock(mutex);
                ready.wait(lock, [&] { return stopping || !queue.empty(); });
                if (stopping && queue.empty()) return;
                item = std::move(queue.front());
                queue.pop_front();
                queue_bytes -= item.payload.size();
                writing = true;
            }
            write_one(item);
            {
                std::lock_guard<std::mutex> lock(mutex);
                writing = false;
                if (queue.empty()) drained.notify_all();
            }
        }
    }

    void write_one(const queued_write & item) {
        if (!ensure_open()) return;
        std::lock_guard<std::mutex> db_lock(db_mutex);
        if (sqlite3_exec(db, "BEGIN IMMEDIATE", nullptr, nullptr, nullptr) != SQLITE_OK) { db_error(); return; }
        uint64_t capacity = params.capacity_bytes;
        sqlite3_stmt * capacity_stmt = nullptr;
        if (sqlite3_prepare_v2(db, "SELECT value FROM meta WHERE key='capacity'", -1, &capacity_stmt, nullptr) == SQLITE_OK &&
                sqlite3_step(capacity_stmt) == SQLITE_ROW) {
            capacity = static_cast<uint64_t>(sqlite3_column_int64(capacity_stmt, 0));
        }
        sqlite3_finalize(capacity_stmt);
        if (item.payload.size() > capacity) {
            sqlite3_exec(db, "ROLLBACK", nullptr, nullptr, nullptr);
            return;
        }
        sqlite3_stmt * existing = nullptr;
        uint64_t old_size = 0;
        sqlite3_prepare_v2(db, "SELECT size FROM entries WHERE kind=?1 AND key=?2", -1, &existing, nullptr);
        sqlite3_bind_int(existing, 1, item.kind);
        sqlite3_bind_text(existing, 2, item.key.data(), int(item.key.size()), SQLITE_TRANSIENT);
        if (sqlite3_step(existing) == SQLITE_ROW) old_size = sqlite3_column_int64(existing, 0);
        sqlite3_finalize(existing);

        sqlite3_stmt * total_stmt = nullptr;
        uint64_t total = 0;
        sqlite3_prepare_v2(db, "SELECT COALESCE(SUM(size),0) FROM entries", -1, &total_stmt, nullptr);
        if (sqlite3_step(total_stmt) == SQLITE_ROW) total = sqlite3_column_int64(total_stmt, 0);
        sqlite3_finalize(total_stmt);
        total -= std::min(total, old_size);

        while (total + item.payload.size() > capacity) {
            sqlite3_stmt * victim = nullptr;
            sqlite3_prepare_v2(db, "SELECT kind,key,size FROM entries WHERE NOT(kind=?1 AND key=?2) ORDER BY last_access LIMIT 1", -1, &victim, nullptr);
            sqlite3_bind_int(victim, 1, item.kind);
            sqlite3_bind_text(victim, 2, item.key.data(), int(item.key.size()), SQLITE_TRANSIENT);
            if (sqlite3_step(victim) != SQLITE_ROW) { sqlite3_finalize(victim); break; }
            const int victim_kind = sqlite3_column_int(victim, 0);
            const std::string victim_key(reinterpret_cast<const char *>(sqlite3_column_text(victim, 1)));
            const uint64_t victim_size = sqlite3_column_int64(victim, 2);
            sqlite3_finalize(victim);
            sqlite3_stmt * del = nullptr;
            sqlite3_prepare_v2(db, "DELETE FROM entries WHERE kind=?1 AND key=?2", -1, &del, nullptr);
            sqlite3_bind_int(del, 1, victim_kind);
            sqlite3_bind_text(del, 2, victim_key.data(), int(victim_key.size()), SQLITE_TRANSIENT);
            sqlite3_step(del);
            sqlite3_finalize(del);
            total -= std::min(total, victim_size);
            ++stats.evictions;
        }

        sqlite3_stmt * put = nullptr;
        sqlite3_prepare_v2(db, "INSERT OR REPLACE INTO entries(kind,key,payload,size,last_access) VALUES(?1,?2,?3,?4,?5)", -1, &put, nullptr);
        sqlite3_bind_int(put, 1, item.kind);
        sqlite3_bind_text(put, 2, item.key.data(), int(item.key.size()), SQLITE_TRANSIENT);
        sqlite3_bind_blob(put, 3, item.payload.data(), int(item.payload.size()), SQLITE_TRANSIENT);
        sqlite3_bind_int64(put, 4, item.payload.size());
        sqlite3_bind_int64(put, 5, now_tick());
        const int rc = sqlite3_step(put);
        sqlite3_finalize(put);
        if (rc == SQLITE_DONE) {
            sqlite3_exec(db, "COMMIT", nullptr, nullptr, nullptr);
            ++stats.writes;
        } else {
            sqlite3_exec(db, "ROLLBACK", nullptr, nullptr, nullptr);
            db_error();
        }
    }

    bool db_error() {
        ++stats.errors;
        warn_once("token cache database operation failed; treating it as a cache miss");
        return false;
    }

    void warn_once(const char * message) {
        bool expected = false;
        if (warned.compare_exchange_strong(expected, true)) {
            LLAMA_LOG_WARN("%s\n", message);
        }
    }

    void close() {
        std::lock_guard<std::mutex> db_lock(db_mutex);
        if (db) sqlite3_close(db);
        db = nullptr;
    }

    llama_token_cache_params params = {nullptr, DEFAULT_CAPACITY};
    std::string configured_path;
    std::string path;
    sqlite3 * db = nullptr;
    std::mutex init_mutex;
    std::mutex db_mutex;
    std::mutex mutex;
    std::condition_variable ready;
    std::condition_variable drained;
    std::deque<queued_write> queue;
    uint64_t queue_bytes = 0;
    bool initialized = false;
    bool stopping = false;
    bool writing = false;
    std::atomic<bool> warned{false};
    std::thread worker;
    cache_counters stats;
};

token_cache & cache() {
    static token_cache instance;
    return instance;
}

} // namespace

std::string llama_token_cache_hash(const void * data, size_t size) {
    sha256 hash;
    hash.update(data, size);
    const auto bytes = hash.finish();
    static const char hex[] = "0123456789abcdef";
    std::string result(64, '0');
    for (size_t i = 0; i < bytes.size(); ++i) {
        result[2*i] = hex[bytes[i] >> 4];
        result[2*i+1] = hex[bytes[i] & 15];
    }
    return result;
}

std::string llama_token_cache_hash(const std::string & data) {
    return llama_token_cache_hash(data.data(), data.size());
}

bool llama_token_cache_lookup(llama_token_cache_kind kind, const std::string & key, std::vector<uint8_t> & payload) {
    return cache().lookup(kind, key, payload);
}

void llama_token_cache_store(llama_token_cache_kind kind, const std::string & key, const std::vector<uint8_t> & payload, bool synchronous) {
    cache().store(kind, key, payload, synchronous);
}

void llama_token_cache_note_tokenizer_hit(bool hit) {
    cache().note_tokenizer(hit);
}

llama_token_cache_params llama_token_cache_default_params() {
    return {nullptr, DEFAULT_CAPACITY};
}

bool llama_token_cache_configure(llama_token_cache_params params) {
    return cache().configure(params);
}

bool llama_token_cache_flush() {
    cache().flush();
    return true;
}

bool llama_token_cache_clear() {
    return cache().clear_all();
}

llama_token_cache_stats llama_token_cache_get_stats() {
    return cache().get_stats();
}
