#include "kazeia_text_tokenizer.h"
#include "llama.h"
#include <cstdio>

bool kz_tok_load(KzTextTokenizer & t, const char * gguf_path) {
    auto mp = llama_model_default_params();
    mp.vocab_only = true;
    mp.n_gpu_layers = 0;  // par sécurité (les backends NPU ignorent ça en vocab_only mais explicite).
    t.model = llama_model_load_from_file(gguf_path, mp);
    if (!t.model) {
        fprintf(stderr, "kz_tok_load: failed to load %s\n", gguf_path);
        return false;
    }
    t.vocab = llama_model_get_vocab(t.model);
    if (!t.vocab) {
        fprintf(stderr, "kz_tok_load: no vocab in %s\n", gguf_path);
        llama_model_free(t.model);
        t.model = nullptr;
        return false;
    }
    fprintf(stderr, "kz_tok_load: %s (n_tokens=%d)\n", gguf_path, llama_vocab_n_tokens(t.vocab));
    return true;
}

void kz_tok_free(KzTextTokenizer & t) {
    if (t.model) llama_model_free(t.model);
    t.model = nullptr;
    t.vocab = nullptr;
}

std::vector<int32_t> kz_tok_encode(const KzTextTokenizer & t,
                                    const std::string & text,
                                    bool parse_special) {
    if (!t.vocab) { fprintf(stderr, "kz_tok_encode: vocab not loaded\n"); return {}; }
    // Dimension : 1 token / octet est une borne sup confortable (en pratique on tokenize plutôt
    // à ~4 octets/token). On alloue avec marge et on resize au retour.
    int32_t cap = (int32_t)text.size() + 16;
    std::vector<int32_t> out(cap);
    int32_t n = llama_tokenize(t.vocab, text.c_str(), (int32_t)text.size(),
                                out.data(), cap, /*add_special=*/false, parse_special);
    if (n < 0) {
        // sous-dimensionné : on retente avec la taille demandée.
        out.resize(-n);
        n = llama_tokenize(t.vocab, text.c_str(), (int32_t)text.size(),
                            out.data(), -n, /*add_special=*/false, parse_special);
    }
    if (n < 0) { fprintf(stderr, "kz_tok_encode: tokenize failed (%d)\n", n); return {}; }
    out.resize(n);
    return out;
}

std::vector<int32_t> kz_tok_encode_tts_prompt(const KzTextTokenizer & t,
                                               const std::string & content) {
    // Format chat Qwen3-TTS exact, vérifié sur dump golden (16 tokens pour
    // "Bonjour je m'appelle Kazeia") :
    //   <|im_start|>assistant\n{content}<|im_end|>\n<|im_start|>assistant\n
    //
    // Note : le 2e <|im_start|>assistant\n est le "cue" pour la génération côté talker.
    // Le talker répond par les codes audio à la suite.
    std::string wrapped;
    wrapped.reserve(content.size() + 64);
    wrapped += "<|im_start|>assistant\n";
    wrapped += content;
    wrapped += "<|im_end|>\n<|im_start|>assistant\n";
    return kz_tok_encode(t, wrapped, /*parse_special=*/true);
}