#include "tts_engine.h"
#include "llama.h"
#include "ggml-backend.h"
#include "cp_inference.h"
#include "sampler.h"
#include "kazeia_text_tokenizer.h"
#include "decoder.h"

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <chrono>
#include <fstream>
#include <string>
#include <vector>

using namespace kazeia::tts;

namespace {

double now_s() {
    using clk = std::chrono::steady_clock;
    return std::chrono::duration<double>(clk::now().time_since_epoch()).count();
}

std::vector<float> read_f32(const std::string& p, size_t n_expected) {
    std::ifstream f(p, std::ios::binary | std::ios::ate);
    if (!f) { fprintf(stderr, "tts_engine: open %s\n", p.c_str()); return {}; }
    size_t n = (size_t)f.tellg() / sizeof(float);
    if (n_expected && n != n_expected) {
        fprintf(stderr, "tts_engine: %s: %zu f32, attendu %zu\n", p.c_str(), n, n_expected); return {};
    }
    f.seekg(0); std::vector<float> v(n); f.read((char*)v.data(), n * sizeof(float)); return v;
}

ggml_backend_dev_t find_htp() {
    for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
        auto d = ggml_backend_dev_get(i);
        if (!strcmp(ggml_backend_dev_name(d), "HTP0")) return d;
    }
    return nullptr;
}

inline float silu(float x) { return x / (1.0f + std::exp(-x)); }

void linear(const float* W, const float* b, int out_dim, int in_dim, const float* x, float* out) {
    for (int m = 0; m < out_dim; ++m) {
        float s = b ? b[m] : 0.0f;
        const float* wr = W + (size_t)m * in_dim;
        for (int k = 0; k < in_dim; ++k) s += wr[k] * x[k];
        out[m] = s;
    }
}

void text_projection(const float* in, int N,
                     const float* fc1_w, const float* fc1_b,
                     const float* fc2_w, const float* fc2_b,
                     float* mid_buf, float* out,
                     int text_hidden, int hidden) {
    for (int n = 0; n < N; ++n) {
        linear(fc1_w, fc1_b, text_hidden, text_hidden, in + n * text_hidden, mid_buf);
        for (int d = 0; d < text_hidden; ++d) mid_buf[d] = silu(mid_buf[d]);
        linear(fc2_w, fc2_b, hidden, text_hidden, mid_buf, out + n * hidden);
    }
}

bool write_wav_pcm16_mono(const char* path, const float* samples, size_t N, int sr = 24000) {
    FILE* f = fopen(path, "wb");
    if (!f) return false;
    const uint32_t data_bytes = (uint32_t)(N * 2);
    const uint32_t riff_size  = 36 + data_bytes;
    auto w16 = [&](uint16_t v){ fwrite(&v, 2, 1, f); };
    auto w32 = [&](uint32_t v){ fwrite(&v, 4, 1, f); };
    fwrite("RIFF", 1, 4, f); w32(riff_size); fwrite("WAVE", 1, 4, f);
    fwrite("fmt ", 1, 4, f); w32(16); w16(1); w16(1);
    w32((uint32_t)sr); w32((uint32_t)sr * 2); w16(2); w16(16);
    fwrite("data", 1, 4, f); w32(data_bytes);
    for (size_t i = 0; i < N; ++i) {
        float v = samples[i]; if (v > 1.f) v = 1.f; else if (v < -1.f) v = -1.f;
        int16_t pcm = (int16_t)std::lround(v * 32767.0f);
        fwrite(&pcm, 2, 1, f);
    }
    fclose(f); return true;
}

} // namespace

struct TtsEngine {
    // --- Constantes manifest ---
    int text_vocab = 151936, text_hidden = 2048, hidden = 1024;
    int tts_bos = 151672, tts_eos = 151673, tts_pad = 151671;
    int codec_bos = 2149, codec_eos = 2150, codec_pad = 2148;
    int codec_think = 2154, codec_think_bos = 2156, codec_think_eos = 2157;
    int lang_fr = 2061;

    // --- Fixtures (load-once) ---
    std::vector<float> text_embed;
    std::vector<float> tp_fc1_w, tp_fc1_b, tp_fc2_w, tp_fc2_b;
    std::vector<float> tok_embd;
    std::vector<float> xvector;

    // Embeds spéciaux pré-projetés (constants une fois le modèle chargé)
    std::vector<float> spec_proj;      // [3 * hidden] = {tts_bos_emb, tts_eos_emb, tts_pad_emb}
    std::vector<float> codec_input_emb;  // [7 * hidden] = {think, think_bos, lang_fr, think_eos, xvec, codec_pad, codec_bos}
    std::vector<float> role_proj;      // [3 * hidden] = projection de <|im_start|>, assistant, \n  (constant car template fixe)
    std::vector<float> codec_pad_emb;  // [hidden] = tok_embd[codec_pad]

    // --- Talker ---
    llama_model   * talker_m   = nullptr;
    llama_context * talker_ctx = nullptr;
    int n_embd  = 0;
    int n_vocab = 0;
    int npe     = 1;   // 1 (RoPE) ou 4 (M-RoPE/I-MRoPE)
    int n_threads = 6;

    // --- CP / Decoder / Tokenizer ---
    CPState        cp_state;
    Decoder        decoder;
    KzTextTokenizer kz_tok;
    bool           has_kz_tok = false;
    bool           cp_use_cache = true;
};

// ===========================================================================
// LOAD
// ===========================================================================

TtsEngine * tts_engine_load(const TtsEngineLoadCfg & cfg) {
    if (!cfg.talker_gguf || !cfg.dump_dir) {
        fprintf(stderr, "tts_engine_load: talker_gguf et dump_dir requis\n"); return nullptr;
    }

    auto eng = new TtsEngine();
    eng->n_threads = cfg.n_threads;
    eng->cp_use_cache = cfg.cp_use_cache;

    std::string D = cfg.dump_dir;
    if (D.back() != '/') D += '/';

    // --- 0) Constantes manifest ---
    {
        std::ifstream f(D + "manifest_text.txt");
        if (!f) { fprintf(stderr, "tts_engine_load: no manifest_text\n"); delete eng; return nullptr; }
        std::string line;
        auto eq = [&](const char* k){ return line.rfind(k, 0) == 0; };
        auto val = [&](size_t off){ return atoi(line.c_str() + off); };
        while (std::getline(f, line)) {
            if      (eq("text_vocab_size:"))     eng->text_vocab = val(16);
            else if (eq("text_hidden_size:"))    eng->text_hidden = val(17);
            else if (eq("hidden_size:"))         eng->hidden = val(12);
            else if (eq("tts_bos_token_id:"))    eng->tts_bos = val(17);
            else if (eq("tts_eos_token_id:"))    eng->tts_eos = val(17);
            else if (eq("tts_pad_token_id:"))    eng->tts_pad = val(17);
            else if (eq("codec_bos_id:"))        eng->codec_bos = val(13);
            else if (eq("codec_eos_id:"))        eng->codec_eos = val(13);
            else if (eq("codec_pad_id:"))        eng->codec_pad = val(13);
            else if (eq("codec_think_id:"))      eng->codec_think = val(15);
            else if (eq("codec_think_bos_id:")) eng->codec_think_bos = val(19);
            else if (eq("codec_think_eos_id:")) eng->codec_think_eos = val(19);
            else if (eq("codec_language_french:")) eng->lang_fr = val(22);
        }
    }

    // --- 1) Fixtures ---
    const int text_hidden = eng->text_hidden;
    const int hidden      = eng->hidden;
    const int text_vocab  = eng->text_vocab;
    const double t_load0 = now_s();
    eng->text_embed = read_f32(D + "text_embed.bin",      (size_t)text_vocab * text_hidden);
    eng->tp_fc1_w   = read_f32(D + "tp_fc1_w.bin",        (size_t)text_hidden * text_hidden);
    eng->tp_fc1_b   = read_f32(D + "tp_fc1_b.bin",        (size_t)text_hidden);
    eng->tp_fc2_w   = read_f32(D + "tp_fc2_w.bin",        (size_t)hidden * text_hidden);
    eng->tp_fc2_b   = read_f32(D + "tp_fc2_b.bin",        (size_t)hidden);
    eng->tok_embd   = read_f32(D + "talker_tok_embd.bin", (size_t)3072 * hidden);
    eng->xvector    = read_f32(D + "damien_xvector.bin",  (size_t)hidden);
    if (eng->text_embed.empty() || eng->tp_fc1_w.empty() || eng->tok_embd.empty() || eng->xvector.empty()) {
        fprintf(stderr, "tts_engine_load: fixtures FAIL\n"); delete eng; return nullptr;
    }

    // --- 2) Backend talker (avant le 1er init backend pour HMX off, etc.) ---
    setenv("GGML_HEXAGON_USE_HMX", "0", 1);
    llama_backend_init();

    // --- 3) Tokenizer optionnel (vocab_only -> ~50 MB RAM) ---
    if (cfg.vocab_gguf && *cfg.vocab_gguf) {
        if (!kz_tok_load(eng->kz_tok, cfg.vocab_gguf)) {
            fprintf(stderr, "tts_engine_load: vocab load FAIL %s\n", cfg.vocab_gguf);
            delete eng; return nullptr;
        }
        eng->has_kz_tok = true;
    }

    // --- 4) Talker (libllama) ---
    auto mp = llama_model_default_params();
    ggml_backend_dev_t devs[2] = { cfg.use_htp ? find_htp() : nullptr, nullptr };
    if (devs[0]) { mp.n_gpu_layers = 99; mp.devices = devs; fprintf(stderr, "talker: HTP0\n"); }
    else         { mp.n_gpu_layers = 0;  fprintf(stderr, "talker: CPU\n"); }
    eng->talker_m = llama_model_load_from_file(cfg.talker_gguf, mp);
    if (!eng->talker_m) { fprintf(stderr, "talker load FAIL %s\n", cfg.talker_gguf); delete eng; return nullptr; }

    auto cp = llama_context_default_params();
    // Borne ctx généreuse : prefill max ~50 + max_steps 512 + marge. Pas critique en CPU.
    cp.n_ctx = 1024;
    cp.n_batch = 1024;
    cp.n_threads = cfg.n_threads;
    cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
    cp.embeddings = true;
    eng->talker_ctx = llama_init_from_model(eng->talker_m, cp);
    if (!eng->talker_ctx) { fprintf(stderr, "talker ctx FAIL\n"); delete eng; return nullptr; }
    eng->n_embd  = llama_model_n_embd(eng->talker_m);
    eng->n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(eng->talker_m));
    if (eng->n_embd != hidden || eng->n_vocab != 3072) {
        fprintf(stderr, "talker dims mismatch (n_embd=%d, n_vocab=%d, attendu %d / 3072)\n",
                eng->n_embd, eng->n_vocab, hidden); delete eng; return nullptr;
    }
    auto rt = llama_model_rope_type(eng->talker_m);
    eng->npe = (rt == LLAMA_ROPE_TYPE_MROPE || rt == LLAMA_ROPE_TYPE_IMROPE) ? 4 : 1;

    // --- 5) CP ---
    // KZTTS_CP_HTP=1 + KZTTS_CP_SCHED=1 active le path sched HMX :
    // poids CP sur HTP buffer, forward via ggml_backend_sched (MUL_MAT -> HMX).
    // CP est un transformer 5L 1024-hidden, batches n_new in {1,2}, donc en-dessous de la
    // limite hexagon nrows(src1) ≤ 1024 qui bloquait BigVGAN. Pattern identique au talker
    // option C déjà validé.
    eng->cp_state.n_threads = cfg.n_threads;
    const bool cp_htp = (getenv("KZTTS_CP_HTP") && atoi(getenv("KZTTS_CP_HTP")) != 0);
    bool cp_ok;
    if (cp_htp) {
        std::vector<ggml_backend_t> cp_backends;
        if (auto htp_dev = find_htp()) {
            ggml_backend_t htp = ggml_backend_dev_init(htp_dev, nullptr);
            if (htp) { cp_backends.push_back(htp); fprintf(stderr, "CP: HTP backend initialisé\n"); }
            else     { fprintf(stderr, "CP: HTP init FAIL, fallback CPU only\n"); }
        }
        ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
        if (!cpu) { fprintf(stderr, "CP: CPU backend init FAIL\n"); delete eng; return nullptr; }
        cp_backends.push_back(cpu);  // CPU TOUJOURS en dernier (assertion sched)
        cp_ok = cp_load_with_backends(eng->cp_state,
                                       (D + "cp_f16.gguf").c_str(),
                                       (D + "cp_heads.bin").c_str(),
                                       (D + "cp_codec_embs.bin").c_str(),
                                       cp_backends, /*owns=*/true);
    } else {
        cp_ok = cp_load(eng->cp_state,
                        (D + "cp_f16.gguf").c_str(),
                        (D + "cp_heads.bin").c_str(),
                        (D + "cp_codec_embs.bin").c_str());
    }
    if (!cp_ok) { fprintf(stderr, "CP load FAIL\n"); delete eng; return nullptr; }

    // --- 6) Decoder ---
    // KZTTS_DECODER_HTP=1 -> charger via load_with_backends({HTP, CPU}) pour activer
    // le sched ggml qui route les MUL_MAT (im2col + matmul issu de ggml_conv_1d) sur HMX.
    // IM2COL, CONV_TRANSPOSE_1D, SIN restent CPU (sched fait le split + copies auto).
    // Sinon : load() classique = CPU pur via ggml_graph_compute_with_ctx (path actuel).
    const bool dec_htp = (getenv("KZTTS_DECODER_HTP") && atoi(getenv("KZTTS_DECODER_HTP")) != 0);
    bool dec_ok;
    if (dec_htp) {
        std::vector<ggml_backend_t> dec_backends;
        // 1er backend = HTP si dispo, sinon CPU seul (équivalent au load classique).
        if (auto htp_dev = find_htp()) {
            ggml_backend_t htp = ggml_backend_dev_init(htp_dev, nullptr);
            if (htp) {
                dec_backends.push_back(htp);
                fprintf(stderr, "decoder: HTP backend initialisé\n");
            } else {
                fprintf(stderr, "decoder: ggml_backend_dev_init(HTP0) FAIL, fallback CPU\n");
            }
        }
        // CPU TOUJOURS en DERNIER (ggml_backend_sched_new assert que backends[end]=CPU,
        // c'est le fallback universel pour les ops non-supportées par les autres).
        ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
        if (!cpu) { fprintf(stderr, "decoder: CPU backend init FAIL\n"); delete eng; return nullptr; }
        dec_backends.push_back(cpu);
        eng->decoder.owns_backends = true;
        dec_ok = eng->decoder.load_with_backends((D + "qwen3tts_decoder.gguf").c_str(), dec_backends);
    } else {
        dec_ok = eng->decoder.load((D + "qwen3tts_decoder.gguf").c_str());
    }
    if (!dec_ok) {
        fprintf(stderr, "decoder load FAIL\n"); delete eng; return nullptr;
    }

    // --- 7) Pre-compute des embeds constants ---
    std::vector<float> mid(text_hidden);

    //   special tokens projetés (tts_bos/eos/pad)
    std::vector<float> spec_text(3 * text_hidden);
    int sp[3] = { eng->tts_bos, eng->tts_eos, eng->tts_pad };
    for (int i = 0; i < 3; ++i)
        std::memcpy(spec_text.data() + i * text_hidden,
                    eng->text_embed.data() + (size_t)sp[i] * text_hidden,
                    text_hidden * sizeof(float));
    eng->spec_proj.resize(3 * hidden);
    text_projection(spec_text.data(), 3,
                    eng->tp_fc1_w.data(), eng->tp_fc1_b.data(),
                    eng->tp_fc2_w.data(), eng->tp_fc2_b.data(),
                    mid.data(), eng->spec_proj.data(), text_hidden, hidden);

    //   codec prefix [think, think_bos, lang_fr, think_eos, x_vec, codec_pad, codec_bos]
    int codec_prefill[4] = { eng->codec_think, eng->codec_think_bos, eng->lang_fr, eng->codec_think_eos };
    eng->codec_input_emb.resize(7 * hidden);
    for (int i = 0; i < 4; ++i)
        std::memcpy(eng->codec_input_emb.data() + i * hidden,
                    eng->tok_embd.data() + (size_t)codec_prefill[i] * hidden,
                    hidden * sizeof(float));
    std::memcpy(eng->codec_input_emb.data() + 4 * hidden, eng->xvector.data(), hidden * sizeof(float));
    std::memcpy(eng->codec_input_emb.data() + 5 * hidden,
                eng->tok_embd.data() + (size_t)eng->codec_pad * hidden, hidden * sizeof(float));
    std::memcpy(eng->codec_input_emb.data() + 6 * hidden,
                eng->tok_embd.data() + (size_t)eng->codec_bos * hidden, hidden * sizeof(float));

    //   role tokens (<|im_start|>, assistant, \n) — projection constante car template fixe.
    //   Pré-calculé une fois, indépendant du texte d'entrée.
    const int role_ids[3] = { 151644, 77091, 198 };
    std::vector<float> role_text(3 * text_hidden);
    for (int i = 0; i < 3; ++i)
        std::memcpy(role_text.data() + i * text_hidden,
                    eng->text_embed.data() + (size_t)role_ids[i] * text_hidden,
                    text_hidden * sizeof(float));
    eng->role_proj.resize(3 * hidden);
    text_projection(role_text.data(), 3,
                    eng->tp_fc1_w.data(), eng->tp_fc1_b.data(),
                    eng->tp_fc2_w.data(), eng->tp_fc2_b.data(),
                    mid.data(), eng->role_proj.data(), text_hidden, hidden);

    //   codec_pad_emb (utilisé partout dans le body)
    eng->codec_pad_emb.assign(eng->tok_embd.data() + (size_t)eng->codec_pad * hidden,
                              eng->tok_embd.data() + (size_t)(eng->codec_pad + 1) * hidden);

    fprintf(stderr, "tts_engine_load: %.2fs (talker_n_embd=%d, npe=%d, kz_tok=%d, cp_cache=%d)\n",
            now_s() - t_load0, eng->n_embd, eng->npe, (int)eng->has_kz_tok, (int)cfg.cp_use_cache);
    return eng;
}

// ===========================================================================
// SYNTHESIZE
// ===========================================================================

TtsSynthesizeResult tts_engine_synthesize(TtsEngine * eng, const TtsSynthesizeCfg & cfg) {
    TtsSynthesizeResult R{};
    if (!eng || !cfg.text || !cfg.out_wav_path) { R.err = -1; return R; }
    if (!eng->has_kz_tok) {
        fprintf(stderr, "tts_engine_synthesize: pas de tokenizer chargé (vocab_gguf manquant au load)\n");
        R.err = -2; return R;
    }
    const int text_hidden = eng->text_hidden;
    const int hidden      = eng->hidden;
    const int n_embd      = eng->n_embd;
    const int n_vocab     = eng->n_vocab;

    // --- 1) Tokenize ---
    auto input_ids = kz_tok_encode_tts_prompt(eng->kz_tok, cfg.text);
    const int input_ids_len = (int)input_ids.size();
    if (input_ids_len < 8) {
        fprintf(stderr, "tts_engine_synthesize: tokens trop courts (%d, min=8)\n", input_ids_len);
        R.err = -3; return R;
    }
    const int Nt = input_ids_len - 5 - 3;  // role(3) + body(Nt) + trailing(5)

    // --- 2) Body projection ---
    std::vector<float> mid(text_hidden);
    std::vector<float> body_text((size_t)Nt * text_hidden);
    for (int i = 0; i < Nt; ++i)
        std::memcpy(body_text.data() + i * text_hidden,
                    eng->text_embed.data() + (size_t)input_ids[3 + i] * text_hidden,
                    text_hidden * sizeof(float));
    std::vector<float> body_proj((size_t)Nt * hidden);
    text_projection(body_text.data(), Nt,
                    eng->tp_fc1_w.data(), eng->tp_fc1_b.data(),
                    eng->tp_fc2_w.data(), eng->tp_fc2_b.data(),
                    mid.data(), body_proj.data(), text_hidden, hidden);

    // --- 3) Assemble prefill ---
    const float* tts_bos_emb = eng->spec_proj.data() + 0 * hidden;
    const float* tts_eos_emb = eng->spec_proj.data() + 1 * hidden;
    const float* tts_pad_emb = eng->spec_proj.data() + 2 * hidden;
    const float* codec_pad_emb = eng->codec_pad_emb.data();
    const int T_prefill = 3 + 6 + (Nt + 1) + 1;

    std::vector<float> prefill((size_t)T_prefill * hidden, 0.0f);
    std::memcpy(prefill.data(), eng->role_proj.data(), 3 * hidden * sizeof(float));
    for (int i = 0; i < 5; ++i) {
        const float* ce = eng->codec_input_emb.data() + i * hidden;
        float* p = prefill.data() + (3 + i) * hidden;
        for (int d = 0; d < hidden; ++d) p[d] = tts_pad_emb[d] + ce[d];
    }
    { const float* ce = eng->codec_input_emb.data() + 5 * hidden;
      float* p = prefill.data() + 8 * hidden;
      for (int d = 0; d < hidden; ++d) p[d] = tts_bos_emb[d] + ce[d]; }
    for (int i = 0; i < Nt; ++i) {
        const float* bp = body_proj.data() + i * hidden;
        float* p = prefill.data() + (9 + i) * hidden;
        for (int d = 0; d < hidden; ++d) p[d] = bp[d] + codec_pad_emb[d];
    }
    { float* p = prefill.data() + (9 + Nt) * hidden;
      for (int d = 0; d < hidden; ++d) p[d] = tts_eos_emb[d] + codec_pad_emb[d]; }
    { const float* cbe = eng->tok_embd.data() + (size_t)eng->codec_bos * hidden;
      float* p = prefill.data() + (10 + Nt) * hidden;
      for (int d = 0; d < hidden; ++d) p[d] = tts_pad_emb[d] + cbe[d]; }

    // --- 4) Reset KV talker (chaque synthèse repart de pos=0) ---
    llama_memory_clear(llama_get_memory(eng->talker_ctx), true);

    // --- 5) Configure les samplers ---
    eng->cp_state.sampler.temp        = cfg.cp_temp;
    eng->cp_state.sampler.top_k       = cfg.cp_top_k;
    eng->cp_state.sampler.top_p       = cfg.cp_top_p;
    eng->cp_state.sampler.rep_penalty = cfg.cp_rep_penalty;
    eng->cp_state.sampler.rep_window  = 16;
    sampler_seed(eng->cp_state.sampler, cfg.seed + 1);

    Sampler talker_sampler{};
    talker_sampler.temp        = cfg.talker_temp;
    talker_sampler.top_k       = cfg.talker_top_k;
    talker_sampler.top_p       = cfg.talker_top_p;
    talker_sampler.rep_penalty = cfg.talker_rep_penalty;
    talker_sampler.rep_window  = 64;
    sampler_seed(talker_sampler, cfg.seed);

    // --- 6) Talker prefill ---
    const double t_pfill0 = now_s();
    {
        const int npe = eng->npe;
        std::vector<llama_pos>     pos(T_prefill * npe, 0);
        std::vector<int32_t>       nsd(T_prefill, 1);
        std::vector<llama_seq_id>  sid0(T_prefill, 0);
        std::vector<llama_seq_id*> sids(T_prefill);
        std::vector<int8_t>        lg(T_prefill, 0);
        for (int i = 0; i < T_prefill; ++i) {
            if (npe == 4) { pos[i] = i; pos[T_prefill + i] = i; pos[2*T_prefill + i] = i; pos[3*T_prefill + i] = 0; }
            else          { pos[i] = i; }
            sids[i] = &sid0[i];
        }
        lg[T_prefill - 1] = 1;
        llama_batch b{};
        b.n_tokens = T_prefill; b.embd = prefill.data();
        b.pos = pos.data(); b.n_seq_id = nsd.data(); b.seq_id = sids.data(); b.logits = lg.data();
        if (llama_decode(eng->talker_ctx, b) != 0) {
            fprintf(stderr, "tts_engine_synthesize: prefill FAIL\n"); R.err = -4; return R;
        }
    }
    R.prefill_s = now_s() - t_pfill0;

    // --- 7) Decode loop (CB0 -> CP -> next_embed -> talker step) ---
    std::vector<float> logits_buf(n_vocab);
    {
        const float* lp = llama_get_logits_ith(eng->talker_ctx, -1);
        memcpy(logits_buf.data(), lp, n_vocab * sizeof(float));
    }
    int cb0 = sampler_sample(talker_sampler, logits_buf.data(), n_vocab);
    std::vector<float> hidden_for_cp(n_embd);
    { const float* hh = llama_get_embeddings_ith(eng->talker_ctx, -1);
      if (hh) memcpy(hidden_for_cp.data(), hh, n_embd * sizeof(float)); }

    std::vector<int32_t> codes_engine;
    codes_engine.reserve((size_t)cfg.max_steps * 16);
    int N_done = 0;
    const double t_loop0 = now_s();
    double t_decode_total = 0, t_cp_total = 0;

    for (int s = 0; s < cfg.max_steps; ++s) {
        if (cb0 == eng->codec_eos) break;
        codes_engine.push_back(cb0);

        const double tcp0 = now_s();
        const float* cb0_emb = eng->tok_embd.data() + (size_t)cb0 * n_embd;
        int32_t cb15[15];
        if (eng->cp_use_cache) cp_predict_cached(eng->cp_state, hidden_for_cp.data(), cb0_emb, cb15);
        else                    cp_predict       (eng->cp_state, hidden_for_cp.data(), cb0_emb, cb15);
        t_cp_total += now_s() - tcp0;
        for (int i = 0; i < 15; ++i) codes_engine.push_back(cb15[i]);

        std::vector<float> next_embed(n_embd, 0.0f);
        const float* e_cb0 = eng->tok_embd.data() + (size_t)cb0 * n_embd;
        for (int d = 0; d < n_embd; ++d) next_embed[d] = e_cb0[d];
        for (int i = 1; i < 16; ++i) {
            int code = cb15[i - 1];
            const float* e = eng->cp_state.codec_embs.data() + ((size_t)(i-1) * 2048 + code) * n_embd;
            for (int d = 0; d < n_embd; ++d) next_embed[d] += e[d];
        }
        for (int d = 0; d < n_embd; ++d) next_embed[d] += tts_pad_emb[d];

        const int npe = eng->npe;
        llama_pos pos1[4] = {0,0,0,0};
        const llama_pos p = T_prefill + s;
        if (npe == 4) { pos1[0] = p; pos1[1] = p; pos1[2] = p; pos1[3] = 0; }
        else          { pos1[0] = p; }
        int32_t nn = 1; llama_seq_id sd = 0; llama_seq_id* sp = &sd; int8_t l = 1;
        llama_batch b{};
        b.n_tokens = 1; b.embd = next_embed.data();
        b.pos = pos1; b.n_seq_id = &nn; b.seq_id = &sp; b.logits = &l;
        const double td0 = now_s();
        if (llama_decode(eng->talker_ctx, b) != 0) { R.err = -5; break; }
        t_decode_total += now_s() - td0;

        {
            const float* lp = llama_get_logits_ith(eng->talker_ctx, -1);
            memcpy(logits_buf.data(), lp, n_vocab * sizeof(float));
        }
        cb0 = sampler_sample(talker_sampler, logits_buf.data(), n_vocab);
        const float* hh = llama_get_embeddings_ith(eng->talker_ctx, -1);
        if (hh) memcpy(hidden_for_cp.data(), hh, n_embd * sizeof(float));
        N_done = s + 1;
    }
    R.talker_loop_s = (now_s() - t_loop0) - t_cp_total;  // ne pas double-compter CP
    R.cp_loop_s     = t_cp_total;
    (void)t_decode_total;  // déjà compris dans talker_loop_s
    R.frames        = N_done;
    R.audio_s       = N_done / 12.0;

    // --- 8) Decoder -> WAV ---
    std::vector<int32_t> codes_dec(16 * N_done);
    for (int t = 0; t < N_done; ++t)
        for (int c = 0; c < 16; ++c)
            codes_dec[c * N_done + t] = codes_engine[t * 16 + c];

    const double t_dec0 = now_s();
    auto wav = eng->decoder.forward(codes_dec, N_done);
    R.decoder_s = now_s() - t_dec0;

    if (!write_wav_pcm16_mono(cfg.out_wav_path, wav.data(), wav.size(), 24000)) {
        fprintf(stderr, "tts_engine_synthesize: WAV write FAIL %s\n", cfg.out_wav_path);
        R.err = -6; return R;
    }

    R.total_s = R.prefill_s + R.talker_loop_s + R.cp_loop_s + R.decoder_s;
    return R;
}

// ===========================================================================
// FREE
// ===========================================================================

void tts_engine_free(TtsEngine * eng) {
    if (!eng) return;
    cp_free(eng->cp_state);
    if (eng->talker_ctx) llama_free(eng->talker_ctx);
    if (eng->talker_m)   llama_model_free(eng->talker_m);
    if (eng->has_kz_tok) kz_tok_free(eng->kz_tok);
    delete eng;
}