// Implementation stt_engine — voir stt_engine.h pour l'API publique.
//
// Port C++ complet de WhisperHybridEngine.kt (527L Kotlin -> ce fichier).
// Backend : ONNX Runtime 1.24.3 + QNN ExecutionProvider (HTP V79 sur SM8750).
// Modèle : Qualcomm AI Hub HfWhisper KV-cache decoder.
//
// Pipeline transcribe :
//   1. PCM16 -> mel via kazeia_mel (cfg Whisper) -> [1, 80, 3000] fp16
//   2. Encoder (NPU) : input_features -> N cross KV caches fp16
//   3. Decoder (NPU) autoregressif, KV-cache, mask R-to-L, 200 steps max :
//        input_ids[1,1] + attention_mask[1,1,1,200] + N self KV + N cross KV + position_ids
//        -> logits[1,51865,1,1] fp16 + N updated self KV
//   4. argmax logits FP16 (override <|translate|> -> <|transcribe|>, EOT -> stop)
//   5. Token decode via BPE byte-level (vocab.json) -> texte UTF-8
//
// Auto-detect dimensions (num_decoder_layers, num_decoder_heads) via session outputs.
// Bench prod ref : Whisper-Small FR 1.6s -> 825 ms total (mel 189 + enc 125 + dec 510), RTF 0.51.
#include "stt_engine.h"
#include "kazeia_mel.h"

#include "onnxruntime/onnxruntime_cxx_api.h"

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <chrono>
#include <fstream>
#include <memory>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>

// ============================================================================
// FP16 helpers (cohabite avec Ort::Float16_t qui est juste un wrapper uint16_t)
// ============================================================================

// IEEE 754 binary16 conversion (round-to-nearest-even).
static uint16_t fp32_to_fp16(float f) {
    uint32_t x; std::memcpy(&x, &f, 4);
    uint32_t sign = (x >> 16) & 0x8000u;
    int32_t  exp  = (int32_t)((x >> 23) & 0xFFu) - 127 + 15;
    uint32_t mant = x & 0x7FFFFFu;
    if (exp >= 31) {                     // inf / nan
        return (uint16_t)(sign | 0x7C00u | (mant ? 1u : 0u));
    } else if (exp <= 0) {               // subnormal / underflow
        if (exp < -10) return (uint16_t)sign;
        mant |= 0x800000u;
        uint32_t shift = (uint32_t)(14 - exp);
        uint32_t round = (mant >> (shift - 1)) & 1u;
        return (uint16_t)(sign | ((mant >> shift) + round));
    } else {
        uint32_t round = (mant >> 12) & 1u;
        return (uint16_t)(sign | (uint32_t)(exp << 10) | (mant >> 13) + round);
    }
}

static float fp16_to_fp32(uint16_t h) {
    uint32_t sign = (uint32_t)(h & 0x8000u) << 16;
    uint32_t exp  = (h >> 10) & 0x1Fu;
    uint32_t mant = h & 0x3FFu;
    uint32_t out;
    if (exp == 0) {
        if (mant == 0) { out = sign; }
        else {
            // subnormal
            int e = -1;
            while (!(mant & 0x400u)) { mant <<= 1; e--; }
            mant &= 0x3FFu;
            out = sign | (uint32_t)((127 + e) << 23) | (mant << 13);
        }
    } else if (exp == 31) {
        out = sign | 0x7F800000u | (mant << 13);
    } else {
        out = sign | (uint32_t)((127 - 15 + (int)exp) << 23) | (mant << 13);
    }
    float f; std::memcpy(&f, &out, 4);
    return f;
}

// ============================================================================
// VAD RMS énergie (port C++ de VadStage.kt)  -- identique à S1.3
// ============================================================================

struct SttVadState {
    int sample_rate, frame_size, rms_threshold, min_speech_frames, silence_end_frames;
    int speech_count = 0, silence_count = 0;
    bool in_speech = false, end_of_speech = false;
};

SttVadState * stt_vad_new(int sample_rate, int frame_size, int rms_threshold,
                          int min_speech_frames, int silence_end_frames) {
    auto * s = new SttVadState();
    s->sample_rate = sample_rate; s->frame_size = frame_size;
    s->rms_threshold = rms_threshold;
    s->min_speech_frames = min_speech_frames; s->silence_end_frames = silence_end_frames;
    return s;
}
void stt_vad_reset(SttVadState * s) {
    if (!s) return;
    s->speech_count = 0; s->silence_count = 0;
    s->in_speech = false; s->end_of_speech = false;
}
void stt_vad_free(SttVadState * s) { delete s; }
bool stt_vad_is_speech(const SttVadState * s) { return s && s->in_speech; }
bool stt_vad_is_end_of_speech(const SttVadState * s) { return s && s->end_of_speech; }

static float rms_pcm16(const int16_t * pcm, int n) {
    if (n <= 0) return 0.0f;
    double acc = 0.0;
    for (int i = 0; i < n; ++i) { double v = (double)pcm[i]; acc += v * v; }
    return (float)std::sqrt(acc / (double)n);
}

bool stt_vad_push(SttVadState * s, const int16_t * pcm, int n_samples) {
    if (!s || !pcm || n_samples <= 0) return false;
    s->end_of_speech = false;
    int offset = 0;
    while (offset + s->frame_size <= n_samples) {
        float r = rms_pcm16(pcm + offset, s->frame_size);
        if (r >= (float)s->rms_threshold) {
            s->speech_count += 1; s->silence_count = 0;
            if (!s->in_speech && s->speech_count >= s->min_speech_frames) s->in_speech = true;
        } else {
            s->silence_count += 1;
            if (s->in_speech && s->silence_count >= s->silence_end_frames) {
                s->in_speech = false; s->end_of_speech = true; s->speech_count = 0;
            } else if (!s->in_speech) {
                s->speech_count = 0;
            }
        }
        offset += s->frame_size;
    }
    return true;
}

// ============================================================================
// JSON minimaliste (mel_filters.json = flat float array, vocab.json = {"k":v,...})
// ============================================================================

// Parse mel_filters.json : "[1.234e-05, 0.0, ...]" (flat array de N floats).
static std::vector<float> parse_mel_filters_json(const char * path) {
    std::ifstream f(path); if (!f) return {};
    std::stringstream ss; ss << f.rdbuf();
    std::string s = ss.str();
    std::vector<float> out;
    out.reserve(80 * 201);
    size_t i = 0;
    while (i < s.size()) {
        // Trouve prochain digit ou signe
        while (i < s.size() && !(std::isdigit((unsigned char)s[i]) || s[i] == '-' || s[i] == '+' || s[i] == '.')) ++i;
        if (i >= s.size()) break;
        size_t j = i;
        // Avance tant que c'est un caractère numérique
        while (j < s.size() && (std::isdigit((unsigned char)s[j]) || s[j] == '.' || s[j] == 'e' || s[j] == 'E' || s[j] == '+' || s[j] == '-')) {
            // garde-fou : - et + uniquement au début ou après e
            if ((s[j] == '-' || s[j] == '+') && j != i && s[j-1] != 'e' && s[j-1] != 'E') break;
            ++j;
        }
        if (j > i) {
            try { out.push_back(std::stof(s.substr(i, j - i))); } catch (...) {}
        }
        i = j;
    }
    return out;
}

// Parse vocab.json (HF whisper format) : {"!":0,"\"":1,...,"<|endoftext|>":50257,...}
// On retourne id -> token (chaîne UTF-8 brute, qui contient parfois des unicodes BPE).
static bool parse_vocab_json(const char * path,
                             std::unordered_map<int, std::string> & id_to_token) {
    std::ifstream f(path); if (!f) return false;
    std::stringstream ss; ss << f.rdbuf();
    std::string s = ss.str();
    // On parcourt en cherchant des paires "key":number
    size_t i = 0;
    while (i < s.size()) {
        // trouve "
        while (i < s.size() && s[i] != '"') ++i;
        if (i >= s.size()) break;
        // Parse string key (handle \\, \", \uXXXX)
        ++i;
        std::string key;
        while (i < s.size() && s[i] != '"') {
            if (s[i] == '\\' && i + 1 < s.size()) {
                char c = s[i+1];
                if (c == '"')  { key.push_back('"'); i += 2; }
                else if (c == '\\') { key.push_back('\\'); i += 2; }
                else if (c == '/')  { key.push_back('/');  i += 2; }
                else if (c == 'n')  { key.push_back('\n'); i += 2; }
                else if (c == 't')  { key.push_back('\t'); i += 2; }
                else if (c == 'r')  { key.push_back('\r'); i += 2; }
                else if (c == 'b')  { key.push_back('\b'); i += 2; }
                else if (c == 'f')  { key.push_back('\f'); i += 2; }
                else if (c == 'u' && i + 5 < s.size()) {
                    // \uXXXX -> codepoint -> UTF-8
                    unsigned cp = 0;
                    for (int k = 0; k < 4; ++k) {
                        char h = s[i+2+k];
                        cp <<= 4;
                        if (h >= '0' && h <= '9') cp |= (unsigned)(h - '0');
                        else if (h >= 'a' && h <= 'f') cp |= (unsigned)(h - 'a' + 10);
                        else if (h >= 'A' && h <= 'F') cp |= (unsigned)(h - 'A' + 10);
                    }
                    // Encode UTF-8
                    if (cp < 0x80) key.push_back((char)cp);
                    else if (cp < 0x800) {
                        key.push_back((char)(0xC0 | (cp >> 6)));
                        key.push_back((char)(0x80 | (cp & 0x3F)));
                    } else {
                        key.push_back((char)(0xE0 | (cp >> 12)));
                        key.push_back((char)(0x80 | ((cp >> 6) & 0x3F)));
                        key.push_back((char)(0x80 | (cp & 0x3F)));
                    }
                    i += 6;
                } else { key.push_back(s[i]); ++i; }
            } else {
                key.push_back(s[i]); ++i;
            }
        }
        if (i >= s.size()) break;
        ++i;  // skip closing "
        // skip whitespace + ':' + whitespace
        while (i < s.size() && (s[i] == ' ' || s[i] == ':' || s[i] == '\t' || s[i] == '\n')) ++i;
        // Parse integer value
        size_t j = i;
        while (j < s.size() && std::isdigit((unsigned char)s[j])) ++j;
        if (j > i) {
            int id = std::stoi(s.substr(i, j - i));
            id_to_token[id] = key;
        }
        i = j;
        // skip jusqu'à la prochaine ","
        while (i < s.size() && s[i] != ',' && s[i] != '}') ++i;
        if (i < s.size() && s[i] == '}') break;
    }
    return !id_to_token.empty();
}

// BPE byte-level GPT-2/Whisper : map character code (in vocab token strings) -> raw byte
static std::unordered_map<int, int> build_bpe_unicode_to_byte() {
    std::unordered_map<int, int> map;
    for (int b = 33;  b <= 126; ++b) map[b] = b;
    for (int b = 161; b <= 172; ++b) map[b] = b;
    for (int b = 174; b <= 255; ++b) map[b] = b;
    int n = 0;
    for (int b = 0; b <= 255; ++b) {
        bool found = false;
        for (auto & kv : map) { if (kv.second == b) { found = true; break; } }
        if (!found) { map[256 + n] = b; ++n; }
    }
    return map;
}

// Parse UTF-8 codepoint at position i. Renvoie cp et avance i.
static unsigned utf8_next(const std::string & s, size_t & i) {
    unsigned char c = (unsigned char)s[i++];
    if (c < 0x80) return c;
    if ((c >> 5) == 0x6) {
        unsigned cp = (c & 0x1F);
        cp = (cp << 6) | ((unsigned char)s[i++] & 0x3F);
        return cp;
    }
    if ((c >> 4) == 0xE) {
        unsigned cp = (c & 0x0F);
        cp = (cp << 6) | ((unsigned char)s[i++] & 0x3F);
        cp = (cp << 6) | ((unsigned char)s[i++] & 0x3F);
        return cp;
    }
    if ((c >> 3) == 0x1E) {
        unsigned cp = (c & 0x07);
        cp = (cp << 6) | ((unsigned char)s[i++] & 0x3F);
        cp = (cp << 6) | ((unsigned char)s[i++] & 0x3F);
        cp = (cp << 6) | ((unsigned char)s[i++] & 0x3F);
        return cp;
    }
    return c;
}

// Décode une liste de token ids -> texte UTF-8 via BPE byte-level Whisper.
// Filtre les tokens spéciaux <|...|>.
static std::string decode_tokens(const std::vector<int> & ids,
                                  const std::unordered_map<int, std::string> & vocab,
                                  const std::unordered_map<int, int> & b2u_byte_map) {
    std::string bytes;
    bytes.reserve(ids.size() * 2);
    for (int id : ids) {
        auto it = vocab.find(id); if (it == vocab.end()) continue;
        const std::string & w = it->second;
        if (w.size() >= 4 && w[0] == '<' && w[1] == '|' &&
            w[w.size()-2] == '|' && w[w.size()-1] == '>') continue;
        size_t i = 0;
        while (i < w.size()) {
            unsigned cp = utf8_next(w, i);
            auto bit = b2u_byte_map.find((int)cp);
            if (bit != b2u_byte_map.end()) bytes.push_back((char)bit->second);
        }
    }
    return bytes;
}

// ============================================================================
// SttEngine — ORT QNN HfWhisper KV-cache
// ============================================================================

namespace {
// Constantes Whisper (multilingue) — confirmées dans le code prod Kotlin.
constexpr int SOT             = 50258;
constexpr int EOT             = 50257;
constexpr int TRANSLATE_TOK   = 50358;
constexpr int TRANSCRIBE_TOK  = 50359;
constexpr int VOCAB_SIZE      = 51865;
constexpr int MEAN_DECODE_LEN = 200;
constexpr int HEAD_DIM        = 64;
constexpr float MASK_NEG      = -100.0f;

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

struct SttEngine {
    SttEngineLoadCfg cfg;
    std::vector<float> mel_basis;

    std::unique_ptr<Ort::Env>            ort_env;
    std::unique_ptr<Ort::SessionOptions> enc_opts, dec_opts;
    std::unique_ptr<Ort::Session>        enc_sess, dec_sess;
    std::unique_ptr<Ort::MemoryInfo>     mem_info;
    std::unique_ptr<Ort::AllocatorWithDefaultOptions> allocator;

    // Auto-detected dims (Whisper-Base: 6/8, Small: 12/12, Medium: 24/16)
    int num_decoder_layers = 0;
    int num_decoder_heads  = 0;

    // Cached I/O names (for the decoder loop tight allocation budget)
    std::vector<std::string>           enc_in_names_owned;
    std::vector<const char *>          enc_in_names;
    std::vector<std::string>           enc_out_names_owned;
    std::vector<const char *>          enc_out_names;

    std::vector<std::string>           dec_in_names_owned;
    std::vector<const char *>          dec_in_names;
    std::vector<std::string>           dec_out_names_owned;
    std::vector<const char *>          dec_out_names;

    // Vocab + BPE
    std::unordered_map<int, std::string> vocab;
    std::unordered_map<int, int>         bpe_unicode_to_byte;

    bool loaded = false;
};

// Charge un fichier binaire ou JSON pour mel filters et retourne le tableau.
// Tente d'abord mel_filters.bin, sinon mel_filters.json.
static bool load_mel_basis_any(const std::string & model_dir,
                                const kazeia_mel::MelConfig & cfg,
                                std::vector<float> & out) {
    if (kazeia_mel::load_mel_basis((model_dir + "/mel_filters.bin").c_str(), cfg, out)) return true;
    auto v = parse_mel_filters_json((model_dir + "/mel_filters.json").c_str());
    const size_t expected = (size_t)cfg.n_mels * (cfg.n_fft/2+1);
    if (v.size() == expected) { out = std::move(v); return true; }
    fprintf(stderr, "stt_engine: mel basis introuvable (ni .bin ni .json valide) à %s\n",
            model_dir.c_str());
    return false;
}

// Détecte num_decoder_layers (en comptant k_cache_cross_*) et num_decoder_heads (shape[0] du 1er).
static void detect_dims(SttEngine & e) {
    Ort::AllocatorWithDefaultOptions a;
    size_t n_out = e.enc_sess->GetOutputCount();
    int n_layers = 0;
    int n_heads  = 0;
    for (size_t i = 0; i < n_out; ++i) {
        auto name_ptr = e.enc_sess->GetOutputNameAllocated(i, a);
        std::string name = name_ptr.get();
        if (name.rfind("k_cache_cross_", 0) == 0) {
            ++n_layers;
            if (n_heads == 0) {
                auto info = e.enc_sess->GetOutputTypeInfo(i);
                auto tinfo = info.GetTensorTypeAndShapeInfo();
                auto shape = tinfo.GetShape();
                if (!shape.empty() && shape[0] > 0) n_heads = (int)shape[0];
            }
        }
    }
    // Défauts si détection ratée (Whisper-Base = 6/8, Small = 12/12)
    if (n_layers == 0) n_layers = 12;
    if (n_heads  == 0) n_heads  = 12;
    e.num_decoder_layers = n_layers;
    e.num_decoder_heads  = n_heads;
}

// Cache les input/output names (les sessions ORT exigent const char**)
static void cache_io_names(Ort::Session & sess,
                            std::vector<std::string> & in_owned,
                            std::vector<const char *> & in_ptrs,
                            std::vector<std::string> & out_owned,
                            std::vector<const char *> & out_ptrs) {
    Ort::AllocatorWithDefaultOptions a;
    size_t ni = sess.GetInputCount(), no = sess.GetOutputCount();
    in_owned.reserve(ni); in_ptrs.reserve(ni);
    for (size_t i = 0; i < ni; ++i) {
        auto p = sess.GetInputNameAllocated(i, a);
        in_owned.emplace_back(p.get());
    }
    for (auto & s : in_owned) in_ptrs.push_back(s.c_str());
    out_owned.reserve(no); out_ptrs.reserve(no);
    for (size_t i = 0; i < no; ++i) {
        auto p = sess.GetOutputNameAllocated(i, a);
        out_owned.emplace_back(p.get());
    }
    for (auto & s : out_owned) out_ptrs.push_back(s.c_str());
}

SttEngine * stt_engine_load(const SttEngineLoadCfg & cfg) {
    if (!cfg.model_dir) {
        fprintf(stderr, "stt_engine_load: model_dir requis\n"); return nullptr;
    }
    auto * eng = new SttEngine();
    eng->cfg = cfg;
    std::string D = cfg.model_dir;

    // 1) Mel basis (binaire optimal, JSON fallback)
    auto mel_cfg = kazeia_mel::config_whisper();
    if (!load_mel_basis_any(D, mel_cfg, eng->mel_basis)) {
        delete eng; return nullptr;
    }

    // 2) ORT env
    try {
        eng->ort_env = std::make_unique<Ort::Env>(ORT_LOGGING_LEVEL_WARNING, "kazeia_stt");
        eng->mem_info = std::make_unique<Ort::MemoryInfo>(
            Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault));
        eng->allocator = std::make_unique<Ort::AllocatorWithDefaultOptions>();
    } catch (const std::exception & e) {
        fprintf(stderr, "stt_engine_load: Ort env FAIL: %s\n", e.what());
        delete eng; return nullptr;
    }

    // 3) Sessions encoder + decoder. QNN EP si use_htp.
    // Options HTP V79 importantes (cf QNN ExecutionProvider doc) :
    //   backend_path             : libQnnHtp.so (CPU stub vers HTP)
    //   htp_performance_mode     : "burst" = max NPU freq (vs "default" = balanced)
    //   htp_arch                 : "v79" = SM8750 (Snapdragon 8 Elite)
    //   enable_htp_fp16_precision: "1" = NPU fp16 native (Whisper a fp16 inputs/outputs)
    //   profiling_level          : "off" en prod (sinon overhead)
    //   rpc_control_latency      : "100" us pour les transferts CPU↔NPU (défaut peut être plus lent)
    auto make_opts = [&](Ort::SessionOptions & opts) {
        if (cfg.use_htp) {
            opts.AppendExecutionProvider("QNN", {
                {"backend_path",             "libQnnHtp.so"},
                {"htp_performance_mode",     "burst"},
                {"htp_arch",                 "79"},
                {"enable_htp_fp16_precision", "1"},
                {"profiling_level",          "off"},
                {"rpc_control_latency",      "100"},
            });
        }
        opts.SetIntraOpNumThreads(cfg.n_threads);
        // Optim générale : graph opt level max + disable mem pattern (HTP gère sa propre mémoire)
        opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
        opts.DisableMemPattern();
    };

    std::string enc_path = D + "/HfWhisperEncoder.onnx";
    std::string dec_path = D + "/HfWhisperDecoder.onnx";

    try {
        eng->enc_opts = std::make_unique<Ort::SessionOptions>();
        make_opts(*eng->enc_opts);
        double t0 = now_s();
        eng->enc_sess = std::make_unique<Ort::Session>(*eng->ort_env, enc_path.c_str(), *eng->enc_opts);
        fprintf(stderr, "stt_engine_load: encoder loaded in %.0f ms\n", (now_s() - t0) * 1000.0);

        eng->dec_opts = std::make_unique<Ort::SessionOptions>();
        make_opts(*eng->dec_opts);
        double t1 = now_s();
        eng->dec_sess = std::make_unique<Ort::Session>(*eng->ort_env, dec_path.c_str(), *eng->dec_opts);
        fprintf(stderr, "stt_engine_load: decoder loaded in %.0f ms\n", (now_s() - t1) * 1000.0);
    } catch (const std::exception & e) {
        fprintf(stderr, "stt_engine_load: ORT session FAIL: %s\n", e.what());
        delete eng; return nullptr;
    }

    // 4) Auto-détection des dims + cache des noms I/O
    detect_dims(*eng);
    cache_io_names(*eng->enc_sess, eng->enc_in_names_owned, eng->enc_in_names,
                                   eng->enc_out_names_owned, eng->enc_out_names);
    cache_io_names(*eng->dec_sess, eng->dec_in_names_owned, eng->dec_in_names,
                                   eng->dec_out_names_owned, eng->dec_out_names);
    fprintf(stderr, "stt_engine_load: %d layers / %d heads detected, enc I/O=%zu/%zu dec I/O=%zu/%zu\n",
            eng->num_decoder_layers, eng->num_decoder_heads,
            eng->enc_in_names.size(), eng->enc_out_names.size(),
            eng->dec_in_names.size(), eng->dec_out_names.size());

    // 5) Vocab.json + BPE byte map
    if (!parse_vocab_json((D + "/vocab.json").c_str(), eng->vocab)) {
        fprintf(stderr, "stt_engine_load: vocab.json absent ou parse FAIL à %s\n", D.c_str());
        delete eng; return nullptr;
    }
    eng->bpe_unicode_to_byte = build_bpe_unicode_to_byte();
    fprintf(stderr, "stt_engine_load: vocab %zu tokens, BPE map %zu entries\n",
            eng->vocab.size(), eng->bpe_unicode_to_byte.size());

    eng->loaded = true;
    return eng;
}

// Création tensor FP16 owned (heap buffer kept alive jusqu'à la fin de la frame).
// Whisper-decoder loop produit ~200 fwd × (1+1+2L+1) tensors, ~13MB/step au pic. OK CPU.
static Ort::Value make_fp16_tensor(const SttEngine & e,
                                    const std::vector<uint16_t> & data,
                                    const std::vector<int64_t> & shape,
                                    std::vector<std::vector<uint16_t>> & owner) {
    owner.emplace_back(data);
    auto & ref = owner.back();
    return Ort::Value::CreateTensor(
        *e.mem_info, ref.data(), ref.size() * sizeof(uint16_t),
        shape.data(), shape.size(),
        ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16);
}

static Ort::Value make_int32_tensor(const SttEngine & e, int v,
                                     const std::vector<int64_t> & shape,
                                     std::vector<std::vector<int32_t>> & owner) {
    owner.push_back({v});
    auto & ref = owner.back();
    return Ort::Value::CreateTensor<int32_t>(
        *e.mem_info, ref.data(), ref.size(),
        shape.data(), shape.size());
}

// Trouve l'index argmax sur logits FP16 [VOCAB_SIZE].
static int argmax_fp16_logits(const uint16_t * data, int n) {
    int best = 0; float best_v = -1e30f;
    for (int i = 0; i < n; ++i) {
        float v = fp16_to_fp32(data[i]);
        if (v > best_v) { best_v = v; best = i; }
    }
    return best;
}

// Convertit FloatArray mel -> FP16 buffer (taille N_MELS * N_FRAMES = 80*3000 = 240000).
static std::vector<uint16_t> mel_to_fp16(const std::vector<float> & mel) {
    std::vector<uint16_t> out(mel.size());
    for (size_t i = 0; i < mel.size(); ++i) out[i] = fp32_to_fp16(mel[i]);
    return out;
}

SttTranscribeResult stt_engine_transcribe(SttEngine * eng, const SttTranscribeCfg & cfg) {
    SttTranscribeResult R{};
    if (!eng || !eng->loaded || !cfg.pcm16 || cfg.n_samples <= 0) {
        R.err = -1; return R;
    }
    if (cfg.sample_rate != 16000) {
        fprintf(stderr, "stt_engine_transcribe: sample_rate=%d ; 16000 attendu\n", cfg.sample_rate);
        R.err = -10; return R;
    }
    double t_start = now_s();

    // 1) PCM -> wav f32
    std::vector<float> wav((size_t)cfg.n_samples);
    for (int i = 0; i < cfg.n_samples; ++i) wav[i] = (float)cfg.pcm16[i] / 32768.0f;
    const double audio_s = (double)cfg.n_samples / cfg.sample_rate;

    // 2) Mel via kazeia_mel
    auto mel_cfg = kazeia_mel::config_whisper();
    double tm0 = now_s();
    int T = 0;
    auto mel = kazeia_mel::compute(wav, eng->mel_basis, mel_cfg, T);
    if (mel.empty()) { R.err = -3; return R; }
    R.mel_ms = (int)((now_s() - tm0) * 1000.0);

    // 3) Encoder — fp16 input -> N cross KV caches fp16
    double te0 = now_s();
    auto mel_fp16 = mel_to_fp16(mel);
    std::vector<std::vector<uint16_t>> tensor_buffers;
    tensor_buffers.reserve(eng->num_decoder_layers * 4 + 4);
    std::vector<int64_t> mel_shape = {1, 80, 3000};

    std::vector<Ort::Value> enc_inputs;
    enc_inputs.emplace_back(make_fp16_tensor(*eng, mel_fp16, mel_shape, tensor_buffers));

    // ORT expects: const char* const* input_names + Value* + size_t input_count.
    std::vector<Ort::Value> enc_outputs;
    try {
        enc_outputs = eng->enc_sess->Run(
            Ort::RunOptions{nullptr},
            eng->enc_in_names.data(), enc_inputs.data(), enc_inputs.size(),
            eng->enc_out_names.data(), eng->enc_out_names.size());
    } catch (const std::exception & e) {
        fprintf(stderr, "stt_engine_transcribe: encoder Run FAIL: %s\n", e.what());
        R.err = -4; return R;
    }
    R.encoder_ms = (int)((now_s() - te0) * 1000.0);

    // 4) Map outputs (encoder) by name -> get cross KV pointers + shapes
    // Owned copies for the decoder loop (will be reused across all steps).
    struct KvCache {
        std::vector<uint16_t> data;
        std::vector<int64_t> shape;
    };
    std::vector<KvCache> cross_k(eng->num_decoder_layers), cross_v(eng->num_decoder_layers);
    for (size_t i = 0; i < eng->enc_out_names.size(); ++i) {
        std::string nm = eng->enc_out_names_owned[i];
        if (nm.rfind("k_cache_cross_", 0) == 0) {
            int idx = std::stoi(nm.substr(14));
            if (idx < eng->num_decoder_layers) {
                auto & val = enc_outputs[i];
                auto shape = val.GetTensorTypeAndShapeInfo().GetShape();
                size_t n = 1; for (auto d : shape) n *= (size_t)d;
                cross_k[idx].shape = shape;
                cross_k[idx].data.resize(n);
                std::memcpy(cross_k[idx].data.data(), val.GetTensorMutableData<uint16_t>(), n * 2);
            }
        } else if (nm.rfind("v_cache_cross_", 0) == 0) {
            int idx = std::stoi(nm.substr(14));
            if (idx < eng->num_decoder_layers) {
                auto & val = enc_outputs[i];
                auto shape = val.GetTensorTypeAndShapeInfo().GetShape();
                size_t n = 1; for (auto d : shape) n *= (size_t)d;
                cross_v[idx].shape = shape;
                cross_v[idx].data.resize(n);
                std::memcpy(cross_v[idx].data.data(), val.GetTensorMutableData<uint16_t>(), n * 2);
            }
        }
    }
    enc_outputs.clear();

    // 5) Decoder loop autoregressif KV-cache (port optimisé zero-copy)
    //
    // ZERO-COPY : on alloue UNE FOIS tous les buffers, on crée UNE FOIS les Ort::Value
    // qui pointent dedans, et on met juste à jour les valeurs in-place entre steps.
    // Critical path : pas de std::vector::emplace_back(data) qui copie 12 MB/step.
    double td0 = now_s();
    const int kv_slots = MEAN_DECODE_LEN - 1;  // 199
    const size_t self_k_n = (size_t)eng->num_decoder_heads * HEAD_DIM * kv_slots;
    const size_t self_v_n = (size_t)eng->num_decoder_heads * kv_slots * HEAD_DIM;
    const std::vector<int64_t> self_k_shape = {eng->num_decoder_heads, 1, HEAD_DIM, kv_slots};
    const std::vector<int64_t> self_v_shape = {eng->num_decoder_heads, 1, kv_slots, HEAD_DIM};

    std::vector<KvCache> self_k(eng->num_decoder_layers), self_v(eng->num_decoder_layers);
    for (int l = 0; l < eng->num_decoder_layers; ++l) {
        self_k[l].shape = self_k_shape; self_k[l].data.assign(self_k_n, fp32_to_fp16(0.0f));
        self_v[l].shape = self_v_shape; self_v[l].data.assign(self_v_n, fp32_to_fp16(0.0f));
    }

    // Buffers persistants pour input_ids (int32 [1,1]), attention_mask (fp16 [1,1,1,200]),
    // position_ids (int32 [1])
    std::vector<int32_t>  input_ids_buf(1);
    std::vector<uint16_t> mask_buf((size_t)MEAN_DECODE_LEN, fp32_to_fp16(MASK_NEG));
    std::vector<int32_t>  pos_ids_buf(1);
    const std::vector<int64_t> input_ids_shape = {1, 1};
    const std::vector<int64_t> mask_shape      = {1, 1, 1, (int64_t)MEAN_DECODE_LEN};
    const std::vector<int64_t> pos_ids_shape   = {1};

    // Pré-création des Ort::Value DECODER inputs UNE FOIS, alignés avec dec_in_names_owned.
    // Chaque Value pointe sur son buffer permanent : ORT lit le buffer à chaque Run, on
    // n'a qu'à updater le contenu in-place avant Run.
    std::vector<Ort::Value> dec_inputs;
    dec_inputs.reserve(eng->dec_in_names_owned.size());
    for (size_t i = 0; i < eng->dec_in_names_owned.size(); ++i) {
        const std::string & nm = eng->dec_in_names_owned[i];
        if (nm == "input_ids") {
            dec_inputs.emplace_back(Ort::Value::CreateTensor<int32_t>(
                *eng->mem_info, input_ids_buf.data(), input_ids_buf.size(),
                input_ids_shape.data(), input_ids_shape.size()));
        } else if (nm == "attention_mask") {
            dec_inputs.emplace_back(Ort::Value::CreateTensor(
                *eng->mem_info, mask_buf.data(), mask_buf.size() * sizeof(uint16_t),
                mask_shape.data(), mask_shape.size(),
                ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16));
        } else if (nm == "position_ids") {
            dec_inputs.emplace_back(Ort::Value::CreateTensor<int32_t>(
                *eng->mem_info, pos_ids_buf.data(), pos_ids_buf.size(),
                pos_ids_shape.data(), pos_ids_shape.size()));
        } else if (nm.rfind("k_cache_self_", 0) == 0 && nm.find("_in") != std::string::npos) {
            int idx = std::stoi(nm.substr(13));
            dec_inputs.emplace_back(Ort::Value::CreateTensor(
                *eng->mem_info, self_k[idx].data.data(), self_k[idx].data.size() * sizeof(uint16_t),
                self_k[idx].shape.data(), self_k[idx].shape.size(),
                ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16));
        } else if (nm.rfind("v_cache_self_", 0) == 0 && nm.find("_in") != std::string::npos) {
            int idx = std::stoi(nm.substr(13));
            dec_inputs.emplace_back(Ort::Value::CreateTensor(
                *eng->mem_info, self_v[idx].data.data(), self_v[idx].data.size() * sizeof(uint16_t),
                self_v[idx].shape.data(), self_v[idx].shape.size(),
                ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16));
        } else if (nm.rfind("k_cache_cross_", 0) == 0) {
            int idx = std::stoi(nm.substr(14));
            dec_inputs.emplace_back(Ort::Value::CreateTensor(
                *eng->mem_info, cross_k[idx].data.data(), cross_k[idx].data.size() * sizeof(uint16_t),
                cross_k[idx].shape.data(), cross_k[idx].shape.size(),
                ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16));
        } else if (nm.rfind("v_cache_cross_", 0) == 0) {
            int idx = std::stoi(nm.substr(14));
            dec_inputs.emplace_back(Ort::Value::CreateTensor(
                *eng->mem_info, cross_v[idx].data.data(), cross_v[idx].data.size() * sizeof(uint16_t),
                cross_v[idx].shape.data(), cross_v[idx].shape.size(),
                ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16));
        } else {
            fprintf(stderr, "stt_engine_transcribe: input decoder inconnu : %s\n", nm.c_str());
            R.err = -5; return R;
        }
    }

    // Cache l'index du tensor "logits" dans dec_out_names_owned
    int logits_out_idx = -1;
    for (size_t i = 0; i < eng->dec_out_names_owned.size(); ++i) {
        if (eng->dec_out_names_owned[i] == "logits") { logits_out_idx = (int)i; break; }
    }
    if (logits_out_idx < 0) {
        fprintf(stderr, "stt_engine_transcribe: pas de sortie 'logits' au decoder\n");
        R.err = -7; return R;
    }
    // Cache la correspondance self_*_out -> idx layer (parsing 1 fois suffit)
    std::vector<int> k_self_out_layer(eng->dec_out_names_owned.size(), -1);
    std::vector<int> v_self_out_layer(eng->dec_out_names_owned.size(), -1);
    for (size_t i = 0; i < eng->dec_out_names_owned.size(); ++i) {
        const std::string & nm = eng->dec_out_names_owned[i];
        if (nm.rfind("k_cache_self_", 0) == 0 && nm.find("_out") != std::string::npos) {
            k_self_out_layer[i] = std::stoi(nm.substr(13));
        } else if (nm.rfind("v_cache_self_", 0) == 0 && nm.find("_out") != std::string::npos) {
            v_self_out_layer[i] = std::stoi(nm.substr(13));
        }
    }

    std::vector<int> generated;
    generated.reserve(MEAN_DECODE_LEN);
    int current_token = SOT;
    int position_id   = 0;
    int real_vocab_size = -1;  // détecté au 1er step depuis le shape des logits

    for (int step = 0; step < MEAN_DECODE_LEN - 1; ++step) {
        // Update mask : un seul fp16 à écrire par step (le slot R-to-L)
        mask_buf[MEAN_DECODE_LEN - step - 1] = fp32_to_fp16(0.0f);
        input_ids_buf[0] = current_token;
        pos_ids_buf[0]   = position_id;

        // Run decoder — les Ort::Value pointent déjà sur les bons buffers (à jour in-place)
        std::vector<Ort::Value> dec_outputs;
        try {
            dec_outputs = eng->dec_sess->Run(
                Ort::RunOptions{nullptr},
                eng->dec_in_names.data(), dec_inputs.data(), dec_inputs.size(),
                eng->dec_out_names.data(), eng->dec_out_names.size());
        } catch (const std::exception & e) {
            fprintf(stderr, "stt_engine_transcribe: decoder Run step=%d FAIL: %s\n", step, e.what());
            R.err = -6; return R;
        }

        // argmax logits avec shape réelle (détectée au step 0)
        const auto logits_shape = dec_outputs[logits_out_idx].GetTensorTypeAndShapeInfo().GetShape();
        if (real_vocab_size < 0) {
            size_t n = 1; for (auto d : logits_shape) n *= (size_t)d;
            real_vocab_size = (int)n;
            if (step == 0) {
                fprintf(stderr, "stt_engine_transcribe: logits shape detected = %d elements (VOCAB_SIZE const=%d)\n",
                        real_vocab_size, VOCAB_SIZE);
            }
        }
        const uint16_t * logits_data = dec_outputs[logits_out_idx].GetTensorMutableData<uint16_t>();
        int token = argmax_fp16_logits(logits_data, real_vocab_size);

        if (cfg.force_transcribe && token == TRANSLATE_TOK) token = TRANSCRIBE_TOK;

        // Update self KV from outputs : memcpy direct dans nos buffers persistants
        // (les Ort::Value dec_inputs pointent DÉJÀ dessus, prêts pour le step suivant)
        for (size_t i = 0; i < eng->dec_out_names_owned.size(); ++i) {
            int k_idx = k_self_out_layer[i];
            int v_idx = v_self_out_layer[i];
            if (k_idx >= 0) {
                auto & val = dec_outputs[i];
                std::memcpy(self_k[k_idx].data.data(),
                            val.GetTensorMutableData<uint16_t>(),
                            self_k[k_idx].data.size() * sizeof(uint16_t));
            } else if (v_idx >= 0) {
                auto & val = dec_outputs[i];
                std::memcpy(self_v[v_idx].data.data(),
                            val.GetTensorMutableData<uint16_t>(),
                            self_v[v_idx].data.size() * sizeof(uint16_t));
            }
        }
        dec_outputs.clear();

        if (token == EOT) break;
        if (token < 50257) generated.push_back(token);

        current_token = token;
        position_id  += 1;
    }
    R.decoder_ms = (int)((now_s() - td0) * 1000.0);

    // 6) Token decode -> texte
    R.text = decode_tokens(generated, eng->vocab, eng->bpe_unicode_to_byte);
    R.n_tokens = (int)generated.size();
    R.detected_language = cfg.language;
    R.total_ms = (int)((now_s() - t_start) * 1000.0);
    R.rtf = (audio_s > 0) ? (float)((double)R.total_ms / (audio_s * 1000.0)) : 0.0f;
    return R;
}

void stt_engine_free(SttEngine * eng) {
    if (!eng) return;
    eng->enc_sess.reset(); eng->dec_sess.reset();
    eng->enc_opts.reset(); eng->dec_opts.reset();
    eng->ort_env.reset();
    delete eng;
}