// Orchestration ggml-side du Talker engine.
//
// Boucle decode : engine Talker fait son vrai forward + greedy sur CB0 ;
// CB1..15 teacher-forcés depuis codes_golden_NxCB.bin (CP non intégré encore).
// next_embed = Σ talker.tok_embd[CB0] + Σ cp_codec_embs[i-1, CB(i)] + tts_pad_embed.
// (En x_vector_only_mode, trailing_text_hidden == tts_pad_embed exactement, vérifié.)
//
// Sortie : codes_engine.bin [N, 16] int32 = ce que l'engine produit. À décoder
// ensuite avec qwen3tts-decoder-test pour obtenir un WAV.
//
// Validation parallèle : à chaque step, on compare next_embed engine-side vs
// talker_step_inputs[s] Python (qui a la même formule). RMSE ~0 attendu (lookups
// dans les mêmes tables, mêmes codes forcés, même sommation).
//
// Usage : tts_orchestrate <gguf> <dump_dir> <out_codes.bin> [cpu|htp] [max_steps]
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <cstdint>
#include <vector>
#include <string>
#include <fstream>
#include <chrono>
#include "llama.h"
#include "ggml-backend.h"
#include "cp_inference.h"

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

static 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;
}

static 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, "open %s\n", p.c_str()); exit(1); }
    size_t n = (size_t)f.tellg() / sizeof(float);
    if (n != n_expected) { fprintf(stderr, "%s: %zu f32, attendu %zu\n", p.c_str(), n, n_expected); exit(1); }
    f.seekg(0);
    std::vector<float> v(n);
    f.read((char*)v.data(), n * sizeof(float));
    return v;
}
static std::vector<int32_t> read_i32(const std::string& p, size_t n_expected) {
    std::ifstream f(p, std::ios::binary | std::ios::ate);
    if (!f) { fprintf(stderr, "open %s\n", p.c_str()); exit(1); }
    size_t n = (size_t)f.tellg() / sizeof(int32_t);
    if (n != n_expected) { fprintf(stderr, "%s: %zu i32, attendu %zu\n", p.c_str(), n, n_expected); exit(1); }
    f.seekg(0);
    std::vector<int32_t> v(n);
    f.read((char*)v.data(), n * sizeof(int32_t));
    return v;
}

int main(int argc, char** argv) {
    if (argc < 4) { printf("usage: %s <talker_gguf> <dump_dir> <out_codes.bin> [cpu|htp] [max_steps]\n", argv[0]); return 1; }
    const char* gguf = argv[1];
    std::string D = argv[2]; if (D.back() != '/') D += '/';
    const char* out_codes = argv[3];
    bool force_cpu = (argc >= 5 && !strcmp(argv[4], "cpu"));
    int max_steps_arg = (argc >= 6) ? atoi(argv[5]) : 0;

    // CP runner : charge cp_f16.gguf + cp_heads.bin + cp_codec_embs.bin depuis dump_dir.
    // Si l'un manque -> fallback teacher-forcing (mode legacy de tts_orchestrate).
    CPState cp_state;
    cp_state.n_threads = 4;
    bool cp_ok = cp_load(cp_state,
                         (D + "cp_f16.gguf").c_str(),
                         (D + "cp_heads.bin").c_str(),
                         (D + "cp_codec_embs.bin").c_str());
    printf("CP : %s\n", cp_ok ? "ENABLED" : "DISABLED (teacher-forced CB1..15)");

    // manifest
    int T_prefill = 0, N_steps_golden = 0, n_embd = 1024, n_vocab = 3072;
    int N_codebooks = 16, cp_vocab = 2048, codec_eos_token_id = 2150;
    {
        std::ifstream f(D + "manifest.txt"); if (!f) { fprintf(stderr, "no manifest\n"); return 1; }
        std::string line;
        while (std::getline(f, line)) {
            if      (line.rfind("T_prefill:", 0) == 0)       T_prefill         = atoi(line.c_str() + 10);
            else if (line.rfind("N_steps:", 0) == 0)         N_steps_golden    = atoi(line.c_str() + 8);
            else if (line.rfind("n_embd:", 0) == 0)          n_embd            = atoi(line.c_str() + 7);
            else if (line.rfind("vocab:", 0) == 0)           n_vocab           = atoi(line.c_str() + 6);
            else if (line.rfind("N_codes:", 0) == 0)         N_steps_golden    = atoi(line.c_str() + 8);
            else if (line.rfind("N_codebooks:", 0) == 0)     N_codebooks       = atoi(line.c_str() + 12);
            else if (line.rfind("cp_vocab:", 0) == 0)        cp_vocab          = atoi(line.c_str() + 9);
            else if (line.rfind("codec_eos_token_id:", 0)==0)codec_eos_token_id= atoi(line.c_str() + 19);
        }
    }
    const int N = (max_steps_arg > 0) ? std::min(max_steps_arg, N_steps_golden) : N_steps_golden;
    printf("manifest: T_prefill=%d N_steps_golden=%d N=%d n_embd=%d n_vocab=%d N_cb=%d cp_vocab=%d eos=%d\n",
           T_prefill, N_steps_golden, N, n_embd, n_vocab, N_codebooks, cp_vocab, codec_eos_token_id);

    // fixtures
    auto prefill_embeds = read_f32(D + "talker_prefill_embeds.bin", (size_t)T_prefill * n_embd);
    auto tts_pad        = read_f32(D + "tts_pad_embed.bin",         (size_t)n_embd);
    auto tok_embd       = read_f32(D + "talker_tok_embd.bin",       (size_t)n_vocab * n_embd);
    auto cp_codec_embs  = read_f32(D + "cp_codec_embs.bin",         (size_t)15 * cp_vocab * n_embd);
    auto codes_golden   = read_i32(D + "codes_golden_NxCB.bin",     (size_t)N_steps_golden * N_codebooks);
    // step_inputs sera comparé pour validation
    auto step_inputs_py = read_f32(D + "talker_step_inputs.bin",    (size_t)N_steps_golden * n_embd);
    printf("fixtures loaded\n");

    // engine
    setenv("GGML_HEXAGON_USE_HMX", "0", 1);
    llama_backend_init();
    auto mp = llama_model_default_params();
    ggml_backend_dev_t devs[2] = { force_cpu ? nullptr : find_htp(), nullptr };
    if (devs[0]) { mp.n_gpu_layers = 99; mp.devices = devs; printf("device: HTP0\n"); }
    else         { mp.n_gpu_layers = 0;  printf("device: CPU\n"); }
    auto m = llama_model_load_from_file(gguf, mp);
    if (!m) { printf("model load FAILED\n"); return 1; }
    auto cp = llama_context_default_params();
    cp.n_ctx = std::max(512, T_prefill + N + 16); cp.n_batch = 1024; cp.n_threads = 4;
    cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
    cp.embeddings = true;
    auto ctx = llama_init_from_model(m, cp);
    if (!ctx) { printf("ctx FAILED\n"); return 1; }

    auto rt = llama_model_rope_type(m);
    const int npe = (rt == LLAMA_ROPE_TYPE_MROPE || rt == LLAMA_ROPE_TYPE_IMROPE) ? 4 : 1;

    // PREFILL
    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;
    const double t_prefill0 = now_s();
    {
        llama_batch b{};
        b.n_tokens = T_prefill; b.embd = prefill_embeds.data();
        b.pos = pos.data(); b.n_seq_id = nsd.data(); b.seq_id = sids.data(); b.logits = lg.data();
        if (llama_decode(ctx, b) != 0) { printf("prefill FAILED\n"); return 1; }
    }
    const double t_prefill = now_s() - t_prefill0;
    printf("prefill OK (T=%d) en %.3f s\n", T_prefill, t_prefill);

    // sample CB0 at prefill output (first step's "input_ids")
    auto logits = llama_get_logits_ith(ctx, -1);
    int cb0 = 0; float mx = logits[0];
    for (int i = 1; i < n_vocab; ++i) if (logits[i] > mx) { mx = logits[i]; cb0 = i; }
    printf("prefill argmax CB0 = %d (golden[0,0]=%d)%s\n", cb0, codes_golden[0],
           cb0 == codes_golden[0] ? " ✓" : " (engine greedy ≠ python sample)");

    // boucle decode
    std::vector<int32_t> codes_engine(N * N_codebooks, 0);
    int n_match_cb0 = 0, n_match_full = 0;
    int n_eos = -1;
    double t_compose = 0, t_decode = 0, t_cp = 0;
    const double t_loop0 = now_s();
    // capture du hidden state du Talker (pour passer au CP) = ce que renvoie
    // llama_get_embeddings_ith(ctx,-1) après chaque decode. On l'a déjà au prefill (h ci-dessous).
    const float* h_last = llama_get_embeddings_ith(ctx, -1);
    std::vector<float> hidden_for_cp(n_embd);
    if (h_last) memcpy(hidden_for_cp.data(), h_last, n_embd * sizeof(float));

    for (int s = 0; s < N; ++s) {
        // CB0 = engine greedy
        codes_engine[s * N_codebooks + 0] = cb0;
        // CB1..15 : soit CP réel, soit teacher-forced
        const double tcp0 = now_s();
        if (cp_ok) {
            const float* cb0_emb = tok_embd.data() + (size_t)cb0 * n_embd;
            int32_t out_cb15[15];
            cp_predict(cp_state, hidden_for_cp.data(), cb0_emb, out_cb15);
            for (int i = 1; i < N_codebooks; ++i) codes_engine[s * N_codebooks + i] = out_cb15[i - 1];
        } else {
            for (int i = 1; i < N_codebooks; ++i) {
                codes_engine[s * N_codebooks + i] = codes_golden[s * N_codebooks + i];
            }
        }
        t_cp += now_s() - tcp0;
        if (cb0 == codes_golden[s * N_codebooks + 0]) n_match_cb0++;
        {
            int m = 0;
            for (int i = 0; i < N_codebooks; ++i)
                if (codes_engine[s * N_codebooks + i] == codes_golden[s * N_codebooks + i]) m++;
            if (m == N_codebooks) n_match_full++;
        }
        if (cb0 == codec_eos_token_id && n_eos < 0) { n_eos = s; printf("  step %d: EOS atteint\n", s); }

        // next_embed = tok_embd[cb0] + sum cp_codec_embs[i-1, CB(i)] + tts_pad
        const double tc0 = now_s();
        std::vector<float> next_embed(n_embd, 0.0f);
        const float* e_cb0 = 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 < N_codebooks; ++i) {
            int code = codes_engine[s * N_codebooks + i];
            const float* e = cp_codec_embs.data() + ((size_t)(i-1) * cp_vocab + 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[d];
        t_compose += now_s() - tc0;

        // valid : next_embed (= input à injecter au decode step s) vs step_inputs_py[s]
        // (= ce que Python a injecté au decode step s). Si codes_engine[s] == codes_golden[s]
        // alors next_embed doit être strictement = step_inputs_py[s] (sum + pad sont les mêmes
        // ops sur les mêmes valeurs).
        if (s < N_steps_golden) {
            const float* py = step_inputs_py.data() + (size_t)s * n_embd;
            double rmse = 0, dot = 0, na = 0, nb = 0;
            for (int d = 0; d < n_embd; ++d) {
                float a = next_embed[d], b = py[d];
                double diff = a - b; rmse += diff * diff;
                dot += (double)a * b; na += (double)a * a; nb += (double)b * b;
            }
            rmse = std::sqrt(rmse / n_embd);
            double cos_v = dot / (std::sqrt(na) * std::sqrt(nb) + 1e-30);
            if (s < 3 || s == N - 1) {
                printf("  step %d: cb0=%d golden=%d, next vs py: rmse=%.4e cos=%.6f\n",
                       s, cb0, codes_golden[s * N_codebooks + 0], rmse, cos_v);
            }
        }

        // Decode talker avec next_embed (sauf au dernier step où on n'a plus besoin du suivant)
        if (s == N - 1) break;
        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 n = 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 = &n; b.seq_id = &sp; b.logits = &l;
        const double td0 = now_s();
        if (llama_decode(ctx, b) != 0) { printf("step %d FAILED\n", s); break; }
        t_decode += now_s() - td0;

        // sample CB0 pour le prochain step + capture hidden pour CP
        logits = llama_get_logits_ith(ctx, -1);
        cb0 = 0; mx = logits[0];
        for (int i = 1; i < n_vocab; ++i) if (logits[i] > mx) { mx = logits[i]; cb0 = i; }
        const float* hh = llama_get_embeddings_ith(ctx, -1);
        if (hh) memcpy(hidden_for_cp.data(), hh, n_embd * sizeof(float));
    }
    const double t_loop = now_s() - t_loop0;

    const double audio_s = N / 12.0; // 12 Hz codec
    printf("DONE: N=%d (audio %.2f s @ 12Hz)\n", N, audio_s);
    printf("  CB0 match golden    : %d/%d\n", n_match_cb0, N);
    printf("  All 16 codes match  : %d/%d (info, dépend du sampling Python)\n", n_match_full, N);
    printf("=== TIMING ===\n");
    printf("  Talker prefill (T=%d) : %.3f s\n", T_prefill, t_prefill);
    printf("  Loop (N=%d steps)     : %.3f s   talker_decode=%.3f, cp=%.3f, compose=%.3f, autres=%.3f\n",
           N, t_loop, t_decode, t_cp, t_compose, t_loop - t_decode - t_cp - t_compose);
    printf("  per-step talker       : %.1f ms\n", t_decode * 1000.0 / N);
    printf("  per-step CP (15 pass) : %.1f ms\n", t_cp * 1000.0 / N);
    printf("  Total Talker+CP       : %.3f s -> RTF talker+cp side = %.3f\n",
           t_prefill + t_loop, (t_prefill + t_loop) / audio_s);
    cp_free(cp_state);

    // dump codes_engine
    {
        std::ofstream f(out_codes, std::ios::binary);
        f.write((char*)codes_engine.data(), N * N_codebooks * sizeof(int32_t));
        printf("codes_engine -> %s (%d frames * %d codebooks)\n", out_codes, N, N_codebooks);
    }
    llama_free(ctx);
    llama_model_free(m);
    return 0;
}