// GGUF loader : lit le .gguf produit par convert_to_gguf.py et peuple la
// table `tensors` du Decoder. Le forward pass dans decoder.cpp accède
// ensuite aux poids par leur clé string.
//
// Phase 1 : chargement seul, pas de forward.

#include "decoder.h"
#include <ggml-backend.h>
#include <ggml-cpu.h>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <vector>

namespace kazeia::tts {

static int gguf_int(gguf_context * g, const char * key, int def = 0) {
    int64_t idx = gguf_find_key(g, key);
    if (idx < 0) return def;
    return (int)gguf_get_val_u32(g, idx);
}
static float gguf_float(gguf_context * g, const char * key, float def = 0.0f) {
    int64_t idx = gguf_find_key(g, key);
    if (idx < 0) return def;
    return gguf_get_val_f32(g, idx);
}

bool Decoder::load(const std::string & path) {
    return load_with_buft(path, ggml_backend_cpu_buffer_type());
}

bool Decoder::load_with_buft(const std::string & path, ggml_backend_buffer_type_t buft) {
    // Étape 1 : lire la metadata + créer les tenseurs SANS allouer (no_alloc).
    struct gguf_init_params params;
    params.no_alloc = true;
    params.ctx = &ctx_w;
    gguf = gguf_init_from_file(path.c_str(), params);
    if (!gguf) {
        fprintf(stderr, "Decoder::load failed to open %s\n", path.c_str());
        return false;
    }

    // Lire metadata
    cfg.latent_dim         = gguf_int(gguf, "qwen3-tts-decoder-12hz.latent_dim", 1024);
    cfg.codebook_dim       = gguf_int(gguf, "qwen3-tts-decoder-12hz.codebook_dim", 512);
    cfg.codebook_size      = gguf_int(gguf, "qwen3-tts-decoder-12hz.codebook_size", 2048);
    cfg.decoder_dim        = gguf_int(gguf, "qwen3-tts-decoder-12hz.decoder_dim", 1536);
    cfg.hidden_size        = gguf_int(gguf, "qwen3-tts-decoder-12hz.hidden_size", 512);
    cfg.intermediate_size  = gguf_int(gguf, "qwen3-tts-decoder-12hz.intermediate_size", 1024);
    cfg.num_attention_heads = gguf_int(gguf, "qwen3-tts-decoder-12hz.num_attention_heads", 16);
    cfg.head_dim           = gguf_int(gguf, "qwen3-tts-decoder-12hz.head_dim", 64);
    cfg.num_hidden_layers  = gguf_int(gguf, "qwen3-tts-decoder-12hz.num_hidden_layers", 8);
    cfg.num_quantizers     = gguf_int(gguf, "qwen3-tts-decoder-12hz.num_quantizers", 16);
    cfg.layer_scale_initial = gguf_float(gguf, "qwen3-tts-decoder-12hz.layer_scale_initial", 0.01f);
    cfg.rope_theta         = gguf_float(gguf, "qwen3-tts-decoder-12hz.rope_theta", 10000.0f);
    cfg.rms_norm_eps       = gguf_float(gguf, "qwen3-tts-decoder-12hz.rms_norm_eps", 1e-5f);
    cfg.upsample_rates     = {8, 5, 4, 3};   // TODO lire array depuis GGUF
    cfg.upsampling_ratios  = {2, 2};
    cfg.total_upsample = 1;
    for (int r : cfg.upsample_rates)    cfg.total_upsample *= r;
    for (int r : cfg.upsampling_ratios) cfg.total_upsample *= r;

    // Étape 2 : allouer les data sur le backend buffer demandé (CPU par défaut,
    // ou Hexagon/Vulkan si précisé). Indispensable pour le scheduler
    // ggml_backend_sched (assertion buffer_id >= 0).
    buf_w = ggml_backend_alloc_ctx_tensors_from_buft(ctx_w, buft);
    if (!buf_w) {
        fprintf(stderr, "Decoder::load: failed to allocate backend buffer\n");
        return false;
    }
    // Marquer ce buffer comme "weights" pour que sched sache qu'il peut
    // (et préfère) router les ops weight-side sur le backend qui possède ces
    // données — ou bien accepter le transfert vers un autre backend.
    ggml_backend_buffer_set_usage(buf_w, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);

    // Étape 3 : lire les data depuis le fichier GGUF dans le backend buffer.
    // Ouverture en lecture binaire ; on slurp les data section par section.
    FILE * f = fopen(path.c_str(), "rb");
    if (!f) {
        fprintf(stderr, "Decoder::load: cannot reopen %s for reading data\n", path.c_str());
        return false;
    }
    const size_t data_off = gguf_get_data_offset(gguf);

    int64_t n_tensors = gguf_get_n_tensors(gguf);
    std::vector<uint8_t> tmp;
    for (int64_t i = 0; i < n_tensors; i++) {
        const char * name = gguf_get_tensor_name(gguf, i);
        ggml_tensor * t = ggml_get_tensor(ctx_w, name);
        if (!t) continue;
        const size_t nbytes = ggml_nbytes(t);
        const size_t offset = data_off + gguf_get_tensor_offset(gguf, i);
        tmp.resize(nbytes);
        if (fseek(f, offset, SEEK_SET) != 0 ||
            fread(tmp.data(), 1, nbytes, f) != nbytes) {
            fprintf(stderr, "Decoder::load: read failed for %s\n", name);
            fclose(f);
            return false;
        }
        ggml_backend_tensor_set(t, tmp.data(), 0, nbytes);
        tensors[name] = t;
    }
    fclose(f);

    fprintf(stderr, "Decoder loaded : %lld tensors, total_upsample=%d "
            "(buf_w %.1f MB, %s)\n",
            (long long)n_tensors, cfg.total_upsample,
            ggml_backend_buffer_get_size(buf_w) / 1048576.0,
            ggml_backend_buffer_name(buf_w));
    return true;
}

// Charge le modèle avec une liste de backends + crée un sched qui split les ops
// entre eux. backends[0] est prioritaire ; backends[last] doit être CPU (fallback).
// Les poids sont alloués sur le buffer du backend[0] (HTP si présent).
bool Decoder::load_with_backends(const std::string & path,
                                  const std::vector<ggml_backend_t> & devs) {
    if (devs.empty()) {
        fprintf(stderr, "load_with_backends: empty backend list\n");
        return false;
    }
    backends = devs;
    // Allouer les poids sur le buft du 1er backend (HMX-friendly si HTP).
    ggml_backend_buffer_type_t buft_w = ggml_backend_get_default_buffer_type(backends[0]);
    if (!load_with_buft(path, buft_w)) return false;

    // --- Précalcul snake constants (Phase 2b) ---
    // Snake activation = x + sin²(α*x)/β. Le code historique faisait host-side
    // exp() à chaque frame dans precompute_snake (incompatible avec sched
    // no_alloc=true). On précalcule a_eff = exp(α), inv_b = 1/(exp(β)+eps) UNE FOIS
    // ici, stockés dans ctx_const + buf_const sur le même backend que buf_w.
    // Le forward stage_bigvgan les utilise comme des poids normaux.
    {
        const float SNAKE_EPS = 1e-9f;
        // 1. Recenser tous les (alpha_key, beta_key) en parcourant les tensors.
        struct SnakePair { std::string prefix; ggml_tensor * alpha; ggml_tensor * beta; };
        std::vector<SnakePair> pairs;
        for (auto & kv : tensors) {
            const std::string & name = kv.first;
            const std::string suf = ".alpha";
            if (name.size() > suf.size() && name.compare(name.size()-suf.size(), suf.size(), suf) == 0) {
                std::string prefix = name.substr(0, name.size() - suf.size());
                std::string beta_name = prefix + ".beta";
                auto bit = tensors.find(beta_name);
                if (bit != tensors.end()) {
                    pairs.push_back({prefix, kv.second, bit->second});
                }
            }
        }
        // 2. Allouer ctx_const : 2 tensors par paire (a_eff, inv_b), tous f32 [C].
        size_t mem_meta = (2 * pairs.size() + 4) * ggml_tensor_overhead();
        ggml_init_params pc{mem_meta, nullptr, /*no_alloc=*/true};
        ctx_const = ggml_init(pc);
        std::vector<std::pair<ggml_tensor*, ggml_tensor*>> created;
        created.reserve(pairs.size());
        for (auto & p : pairs) {
            const int C = (int)p.alpha->ne[0];
            ggml_tensor * a_eff = ggml_new_tensor_1d(ctx_const, GGML_TYPE_F32, C);
            ggml_tensor * inv_b = ggml_new_tensor_1d(ctx_const, GGML_TYPE_F32, C);
            ggml_set_name(a_eff, (p.prefix + ".a_eff").c_str());
            ggml_set_name(inv_b, (p.prefix + ".inv_b").c_str());
            created.push_back({a_eff, inv_b});
        }
        // 3. Allouer buf_const sur le même backend que buf_w.
        buf_const = ggml_backend_alloc_ctx_tensors_from_buft(ctx_const, buft_w);
        if (!buf_const) {
            fprintf(stderr, "load_with_backends: snake buf_const alloc FAIL\n");
            return false;
        }
        ggml_backend_buffer_set_usage(buf_const, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
        // 4. Calculer host-side puis tensor_set.
        std::vector<float> tmp_a, tmp_b;
        for (size_t i = 0; i < pairs.size(); ++i) {
            const int C = (int)pairs[i].alpha->ne[0];
            // Lire α et β depuis leur backend buffer (opt_hostbuf=1 -> mappable CPU).
            tmp_a.resize(C); tmp_b.resize(C);
            ggml_backend_tensor_get(pairs[i].alpha, tmp_a.data(), 0, C * sizeof(float));
            ggml_backend_tensor_get(pairs[i].beta,  tmp_b.data(), 0, C * sizeof(float));
            std::vector<float> a_eff(C), inv_b(C);
            for (int c = 0; c < C; ++c) {
                a_eff[c] = std::exp(tmp_a[c]);
                inv_b[c] = 1.0f / (std::exp(tmp_b[c]) + SNAKE_EPS);
            }
            ggml_backend_tensor_set(created[i].first,  a_eff.data(), 0, C * sizeof(float));
            ggml_backend_tensor_set(created[i].second, inv_b.data(), 0, C * sizeof(float));
            snake_consts[pairs[i].prefix] = created[i];
        }
        fprintf(stderr, "Decoder: %zu snake constants pré-calculées (buf_const %.1f MB sur %s)\n",
                pairs.size(),
                ggml_backend_buffer_get_size(buf_const) / 1048576.0,
                ggml_backend_buffer_name(buf_const));
    }

    // Mode minimal d'essai : KZTTS_DECODER_SCHED=1 active un vrai sched (refactor stages requis).
    // Par défaut : pas de sched, on profite juste de opt_hostbuf=1 (HTP buffer CPU-mappé)
    // pour que ggml_graph_compute_with_ctx CPU pur lise directement les poids HTP.
    // -> aucun gain HMX dans ce mode, mais valide la chaîne load HTP + accès CPU.
    if (getenv("KZTTS_DECODER_SCHED") && atoi(getenv("KZTTS_DECODER_SCHED")) != 0) {
        std::vector<ggml_backend_buffer_type_t> bufts;
        bufts.reserve(backends.size());
        for (auto b : backends) bufts.push_back(ggml_backend_get_default_buffer_type(b));
        const size_t graph_size = 16384;
        sched = ggml_backend_sched_new(backends.data(), bufts.data(),
                                        (int)backends.size(), graph_size,
                                        /*parallel=*/false, /*op_offload=*/true);
        if (!sched) {
            fprintf(stderr, "load_with_backends: ggml_backend_sched_new FAILED\n");
            return false;
        }
        fprintf(stderr, "Decoder: sched créé (refactor stages requis, KZTTS_DECODER_SCHED=1)\n");
    } else {
        fprintf(stderr, "Decoder: poids sur %s, compute CPU pur (sched OFF)\n",
                ggml_backend_name(backends[0]));
    }
    return true;
}

void Decoder::compute_graph(struct ggml_context * ctx, struct ggml_cgraph * gf, int n_threads) {
    if (sched) {
        // Path backend abstrait : split auto CPU/HTP.
        ggml_backend_sched_reset(sched);
        if (!ggml_backend_sched_alloc_graph(sched, gf)) {
            fprintf(stderr, "compute_graph: sched_alloc_graph FAILED\n"); return;
        }
        if (ggml_backend_sched_graph_compute(sched, gf) != GGML_STATUS_SUCCESS) {
            fprintf(stderr, "compute_graph: sched_graph_compute FAILED\n");
        }
    } else {
        // Path legacy CPU pur (load via load() ou load_with_buft sans backend list).
        ggml_graph_compute_with_ctx(ctx, gf, n_threads);
    }
}

void Decoder::unload() {
    if (sched)     { ggml_backend_sched_free(sched); sched = nullptr; }
    if (buf_const) { ggml_backend_buffer_free(buf_const); buf_const = nullptr; }
    if (ctx_const) { ggml_free(ctx_const); ctx_const = nullptr; }
    snake_consts.clear();
    if (buf_w) { ggml_backend_buffer_free(buf_w); buf_w = nullptr; }
    if (ctx_w) { ggml_free(ctx_w); ctx_w = nullptr; }
    if (ctx_g) { ggml_free(ctx_g); ctx_g = nullptr; }
    if (gguf)  { gguf_free(gguf);  gguf  = nullptr; }
    if (owns_backends) {
        for (auto b : backends) if (b) ggml_backend_free(b);
    }
    backends.clear();
    tensors.clear();
}

}  // namespace kazeia::tts