#pragma once
#include <ggml.h>
#include <gguf.h>
#include <ggml-backend.h>
#include <string>
#include <vector>
#include <unordered_map>

namespace kazeia::tts {

struct DecoderConfig {
    int latent_dim;            // 1024
    int codebook_dim;          // 512
    int codebook_size;         // 2048
    int decoder_dim;           // 1536
    int hidden_size;           // 512
    int intermediate_size;     // 1024
    int num_attention_heads;   // 16
    int head_dim;              // 64
    int num_hidden_layers;     // 8
    int num_quantizers;        // 16
    std::vector<int> upsample_rates;     // {8, 5, 4, 3}
    std::vector<int> upsampling_ratios;  // {2, 2}
    float layer_scale_initial;           // 0.01
    float rope_theta;                    // 10000.0
    float rms_norm_eps;                  // 1e-5
    int total_upsample;                  // = prod(upsample_rates) * prod(upsampling_ratios) = 1920
};

struct Decoder {
    DecoderConfig cfg;

    // Tensor lookup table — populated from GGUF on load.
    std::unordered_map<std::string, struct ggml_tensor*> tensors;

    // ggml backend + context (CPU baseline ; Vulkan/Hexagon en Phase 3-4).
    // Loader créé un ctx_w avec no_alloc=true puis place les data sur le
    // backend buffer `buf_w` (CPU par défaut). Cette indirection est
    // requise pour que le scheduler ggml_backend_sched puisse transférer
    // les poids vers Hexagon/Vulkan à la demande.
    struct ggml_context * ctx_w   = nullptr;
    struct ggml_context * ctx_g   = nullptr;
    struct gguf_context * gguf    = nullptr;
    struct ggml_backend_buffer * buf_w = nullptr;

    // Backend list + sched. backends[0] est prioritaire (HTP si dispo, sinon CPU).
    // backends[end] doit TOUJOURS être CPU (fallback pour ops non-supportées).
    // load_with_backends() peuple ça, et compute_graph() exécute via le sched.
    std::vector<ggml_backend_t> backends;
    struct ggml_backend_sched * sched = nullptr;
    bool owns_backends = false;  // si true, unload() libère aussi backends[]

    // Snake constants pré-calculées au load_with_backends pour rendre
    // stage_bigvgan sched-friendly. Sans ça, precompute_snake host-side ne
    // marche pas avec ctx no_alloc=true. Map clé "{prefix}.alpha" -> (a_eff, inv_b)
    // avec a_eff = exp(α), inv_b = 1 / (exp(β) + eps).
    // Alloués dans ctx_const + buf_const sur le même backend que buf_w.
    struct ggml_context * ctx_const = nullptr;
    struct ggml_backend_buffer * buf_const = nullptr;
    std::unordered_map<std::string, std::pair<struct ggml_tensor*, struct ggml_tensor*>> snake_consts;

    // load() utilise le backend CPU par défaut. Pour Hexagon/Vulkan, passer
    // un buft explicite via load_with_buft(), OU une liste de backends via
    // load_with_backends() (recommandé pour HTP : crée le sched + buft auto).
    bool load(const std::string & gguf_path);
    bool load_with_buft(const std::string & gguf_path, ggml_backend_buffer_type_t buft);
    bool load_with_backends(const std::string & gguf_path,
                            const std::vector<ggml_backend_t> & devs);
    void unload();

    // Helper : compute un graphe via le sched (split CPU/HTP auto).
    // Si sched=NULL (chargement CPU pur via load_with_buft), tombe sur
    // ggml_graph_compute_with_ctx pour rester compatible.
    void compute_graph(struct ggml_context * ctx, struct ggml_cgraph * gf, int n_threads);

    // Forward : codes [B, num_quantizers=16, T] long → wav [B, 1, T*total_upsample] f32
    // Returns owned buffer, length = T * total_upsample.
    std::vector<float> forward(const std::vector<int32_t>& codes, int T);

    // ---- Tests par étape ----
    // Quantizer seul : codes [16, T] → hidden [512, T]. Permet validation
    // tensor-par-tensor vs scripts/dump_reference.py.
    std::vector<float> test_quantizer(const std::vector<int32_t>& codes_flat, int T);
};

}  // namespace kazeia::tts