// 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 <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;

    // 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_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