// JNI bridge pour libkazeia_tts.so. Thin wrapper autour de tts_engine.{h,cpp}.
// Le handle exposé côté Kotlin est un jlong = (jlong)(TtsEngine*).
//
// Signatures (à matcher dans le .kt) :
//   nativeLoad(talkerGguf, vocabGguf, dumpDir, useHtp, nThreads, cpUseCache) : Long
//   nativeSynthesize(handle, text, outWavPath, maxSteps, seed,
//                    cpTemp, cpTopK, cpTopP, cpRepPenalty,
//                    talkerTemp, talkerTopK, talkerTopP, talkerRepPenalty)
//                    : IntArray { err, frames, total_ms, prefill_ms, talker_ms, cp_ms, decoder_ms }
//   nativeFree(handle)
#include "tts_engine.h"
#include <jni.h>
#include <string>
#include <cstdio>
#include <vector>
#include <map>
#include <mutex>

namespace {

// Helper : récupère une C-string UTF-8 depuis un jstring (rendre à java avec release).
struct JStr {
    JNIEnv * env;
    jstring j;
    const char * c;
    JStr(JNIEnv* e, jstring s) : env(e), j(s), c(s ? e->GetStringUTFChars(s, nullptr) : nullptr) {}
    ~JStr() { if (j && c) env->ReleaseStringUTFChars(j, c); }
};

// Voix courante par handle : x_vector dérivé d'un WAV via le speaker encoder
// (nativeSetVoiceWav). Appliqué en xvector_override à chaque nativeSynthesize.
// Garde l'API synthesize inchangée (pas de param voix à propager partout).
std::mutex g_voice_mtx;
std::map<uintptr_t, std::vector<float>> g_voice_xvec;

} // namespace

extern "C" {

JNIEXPORT jlong JNICALL
Java_com_kazeia_tts_TtsJni_nativeLoad(JNIEnv* env, jobject /*thiz*/,
                                       jstring talker_gguf, jstring vocab_gguf,
                                       jstring dump_dir,
                                       jboolean use_htp, jint n_threads,
                                       jboolean cp_use_cache) {
    // DEBUG opt-in : redirige stderr natif vers un fichier app-lisible quand
    // KZTTS_STDERR_LOG est défini. En prod (env absent) on laisse stderr tel quel.
    if (const char * dbg = getenv("KZTTS_STDERR_LOG")) {
        freopen(dbg, "w", stderr);
        setvbuf(stderr, nullptr, _IONBF, 0);
    }
    JStr talker(env, talker_gguf);
    JStr vocab (env, vocab_gguf);
    JStr dump  (env, dump_dir);
    TtsEngineLoadCfg lc;
    lc.talker_gguf  = talker.c;
    lc.vocab_gguf   = vocab.c;
    lc.dump_dir     = dump.c;
    lc.use_htp      = (use_htp == JNI_TRUE);
    lc.n_threads    = n_threads;
    lc.cp_use_cache = (cp_use_cache == JNI_TRUE);
    // Speaker encoder embarqué (clonage vocal "se baser sur les WAV"). Les assets
    // vivent dans le dump dir ; si présents, on les passe pour activer
    // tts_engine_encode_speaker_wav(). Absents -> x_vector fixture (damien).
    const std::string se_path = std::string(dump.c ? dump.c : "") + "/speaker_encoder.gguf";
    const std::string mb_path = std::string(dump.c ? dump.c : "") + "/mel_basis_qwen3tts.bin";
    if (FILE* tf = fopen(se_path.c_str(), "rb")) { fclose(tf); lc.speaker_encoder_gguf = se_path.c_str(); }
    if (FILE* tf = fopen(mb_path.c_str(), "rb")) { fclose(tf); lc.mel_basis_path = mb_path.c_str(); }
    auto * eng = tts_engine_load(lc);
    return (jlong)(uintptr_t)eng;
}

JNIEXPORT jintArray JNICALL
Java_com_kazeia_tts_TtsJni_nativeSynthesize(JNIEnv* env, jobject /*thiz*/,
                                             jlong handle,
                                             jstring text, jstring out_wav_path,
                                             jint max_steps, jint seed,
                                             jfloat cp_temp, jint cp_top_k, jfloat cp_top_p, jfloat cp_rep_penalty,
                                             jfloat talker_temp, jint talker_top_k, jfloat talker_top_p, jfloat talker_rep_penalty) {
    auto * eng = (TtsEngine*)(uintptr_t)handle;
    JStr text_s(env, text);
    JStr wav_s (env, out_wav_path);

    TtsSynthesizeCfg sc;
    sc.text           = text_s.c;
    sc.out_wav_path   = wav_s.c;
    sc.max_steps      = max_steps;
    sc.seed           = (uint32_t)seed;
    sc.cp_temp        = cp_temp;
    sc.cp_top_k       = cp_top_k;
    sc.cp_top_p       = cp_top_p;
    sc.cp_rep_penalty = cp_rep_penalty;
    sc.talker_temp        = talker_temp;
    sc.talker_top_k       = talker_top_k;
    sc.talker_top_p       = talker_top_p;
    sc.talker_rep_penalty = talker_rep_penalty;

    // Voix courante (si un nativeSetVoiceWav a réussi pour ce handle).
    std::vector<float> voice;
    {
        std::lock_guard<std::mutex> lk(g_voice_mtx);
        auto it = g_voice_xvec.find((uintptr_t)handle);
        if (it != g_voice_xvec.end()) voice = it->second;
    }
    if (!voice.empty()) {
        sc.xvector_override     = voice.data();
        sc.xvector_override_len = (int)voice.size();
    }

    TtsSynthesizeResult R{ -100, 0, 0, 0, 0, 0, 0, 0 };
    if (eng) R = tts_engine_synthesize(eng, sc);

    // Pack résultats dans IntArray :
    //   [0] err
    //   [1] frames
    //   [2] total_ms      (audio_s codé séparément si besoin = frames/12)
    //   [3] prefill_ms
    //   [4] talker_loop_ms
    //   [5] cp_loop_ms
    //   [6] decoder_ms
    jint out[7];
    out[0] = R.err;
    out[1] = R.frames;
    out[2] = (jint)(R.total_s        * 1000.0);
    out[3] = (jint)(R.prefill_s      * 1000.0);
    out[4] = (jint)(R.talker_loop_s  * 1000.0);
    out[5] = (jint)(R.cp_loop_s      * 1000.0);
    out[6] = (jint)(R.decoder_s      * 1000.0);
    jintArray arr = env->NewIntArray(7);
    env->SetIntArrayRegion(arr, 0, 7, out);
    return arr;
}

// Dérive un x_vector depuis un WAV (mono/stéréo, n'importe quelle fréquence —
// l'encoder downmixe + resample à 24k) et le mémorise comme voix courante pour
// ce handle. Renvoie true si l'encodage a produit un x_vector de longueur valide.
// Le swap prend effet au prochain nativeSynthesize. Requiert que nativeLoad ait
// trouvé speaker_encoder.gguf + mel_basis dans le dump dir.
JNIEXPORT jboolean JNICALL
Java_com_kazeia_tts_TtsJni_nativeSetVoiceWav(JNIEnv* env, jobject /*thiz*/,
                                             jlong handle, jstring wav_path) {
    auto * eng = (TtsEngine*)(uintptr_t)handle;
    if (!eng) return JNI_FALSE;
    JStr wav(env, wav_path);
    if (!wav.c) return JNI_FALSE;
    std::vector<float> xv = tts_engine_encode_speaker_wav(eng, wav.c);
    if (xv.empty()) {
        fprintf(stderr, "nativeSetVoiceWav: encodage FAIL pour %s\n", wav.c);
        return JNI_FALSE;
    }
    {
        std::lock_guard<std::mutex> lk(g_voice_mtx);
        g_voice_xvec[(uintptr_t)handle] = std::move(xv);
    }
    return JNI_TRUE;
}

JNIEXPORT void JNICALL
Java_com_kazeia_tts_TtsJni_nativeFree(JNIEnv* /*env*/, jobject /*thiz*/, jlong handle) {
    {
        std::lock_guard<std::mutex> lk(g_voice_mtx);
        g_voice_xvec.erase((uintptr_t)handle);
    }
    auto * eng = (TtsEngine*)(uintptr_t)handle;
    tts_engine_free(eng);
}

} // extern "C"