package com.kazeia.stt

// Engine STT Whisper-Small (Qualcomm AI Hub HfWhisper KV-cache) in-process.
// Backend : ONNX Runtime + QNN ExecutionProvider (HTP V79 sur SM8750).
// Empreinte RAM : ~545 MB (encoder ctx 201 MB + decoder ctx 345 MB).
//
// Bench prod (Whisper-Small, FR, audio 1.6s) : ~825 ms, RTF 0.51.
// Bench port C++ (audio 5s, FR) : ~2000 ms, RTF 0.40 — voir CLAUDE.md/STT.
//
// Le bridge natif est dans dist/jni/kazeia_stt_jni.cpp ; build via
// dist/build_kazeia_stt.sh.

class SttJni {
    external fun nativeLoad(
        modelDir: String,
        useHtp: Boolean,
        nThreads: Int,
        maxDecodeSteps: Int
    ): Long

    /** Renvoie une chaîne : "err|text|mel_ms|enc_ms|dec_ms|total_ms|n_tokens|rtf" */
    external fun nativeTranscribe(
        handle: Long,
        pcm: ShortArray,
        sampleRate: Int,
        language: String,
        forceTranscribe: Boolean
    ): String

    external fun nativeFree(handle: Long)

    external fun nativeVadNew(
        sampleRate: Int, frameSize: Int, rmsThreshold: Int,
        minSpeechFrames: Int, silenceEndFrames: Int
    ): Long

    /** 0 = silence, 1 = speech in progress, 2 = end_of_speech declenched */
    external fun nativeVadPush(handle: Long, pcm: ShortArray): Int
    external fun nativeVadReset(handle: Long)
    external fun nativeVadFree(handle: Long)

    companion object { init { System.loadLibrary("kazeia_stt") } }
}

data class SttResult(
    val text: String,
    val melMs: Int,
    val encoderMs: Int,
    val decoderMs: Int,
    val totalMs: Int,
    val nTokens: Int,
    val rtf: Float
)

class SttEngine(
    modelDir: String,
    useHtp: Boolean = true,
    nThreads: Int = 6,
    maxDecodeSteps: Int = 200
) {
    private val jni = SttJni()
    private val h = jni.nativeLoad(modelDir, useHtp, nThreads, maxDecodeSteps)
    init { require(h != 0L) { "SttEngine: load FAILED (modelDir=$modelDir, useHtp=$useHtp)" } }

    /**
     * Transcrit un buffer PCM16 mono. Le sample rate doit être 16000 (resample externe sinon).
     * @param language "fr", "en", "auto", ... (cf Whisper langues, 99 codes ISO).
     * @param forceTranscribe true -> override <|translate|> en <|transcribe|>.
     */
    @Throws(IllegalStateException::class)
    fun transcribe(
        pcm: ShortArray,
        language: String = "fr",
        sampleRate: Int = 16000,
        forceTranscribe: Boolean = true
    ): SttResult {
        val raw = jni.nativeTranscribe(h, pcm, sampleRate, language, forceTranscribe)
        // Format : "err|text|mel|enc|dec|total|n_tokens|rtf"
        // limit=8 conserve un éventuel '|' dans le texte (au cas où) dans le 2e champ.
        val parts = raw.split('|', limit = 8)
        if (parts.size < 8) error("SttEngine.transcribe: format inattendu '$raw'")
        val err = parts[0].toInt()
        if (err != 0) error("SttEngine.transcribe err=$err raw='$raw'")
        return SttResult(
            text       = parts[1],
            melMs      = parts[2].toInt(),
            encoderMs  = parts[3].toInt(),
            decoderMs  = parts[4].toInt(),
            totalMs    = parts[5].toInt(),
            nTokens    = parts[6].toInt(),
            rtf        = parts[7].toFloat()
        )
    }

    fun release() { jni.nativeFree(h) }
}

/**
 * VAD RMS énergie compatible drop-in avec VadStage.kt prod.
 * Paramètres par défaut : 16 kHz, frame 100 ms (1600 samples), seuil PCM16 brut 150,
 * 3 frames pour déclencher parole (300 ms), 8 frames pour fin (800 ms).
 *
 * Usage typique :
 *   val vad = SttVad()
 *   audioRecord.read(buf) { pcm ->
 *     when (vad.push(pcm)) {
 *       SttVad.SILENCE        -> { ... }
 *       SttVad.SPEECH         -> { ... // bufferise }
 *       SttVad.END_OF_SPEECH  -> { stt.transcribe(buffer) ; vad.reset() ; buffer.clear() }
 *     }
 *   }
 */
class SttVad(
    sampleRate: Int = 16000,
    frameSize: Int = 1600,
    rmsThreshold: Int = 150,
    minSpeechFrames: Int = 3,
    silenceEndFrames: Int = 8
) {
    private val jni = SttJni()
    private val h = jni.nativeVadNew(sampleRate, frameSize, rmsThreshold,
                                      minSpeechFrames, silenceEndFrames)
    fun push(pcm: ShortArray): Int = jni.nativeVadPush(h, pcm)
    fun reset() { jni.nativeVadReset(h) }
    fun release() { jni.nativeVadFree(h) }

    companion object {
        const val SILENCE       = 0
        const val SPEECH        = 1
        const val END_OF_SPEECH = 2
    }
}