eric / piclead-client

package fr.ebii.card2vcf.audio

import android.content.Context
import android.util.Log
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonPrimitive
import kotlinx.serialization.json.jsonObject
import org.vosk.Model
import org.vosk.Recognizer
import java.io.File

/**
 * Transcription locale via le moteur Vosk (modèle français small ~41 Mo).
 *
 * Sur le premier appel à [accepterEchantillons], le modèle est extrait des assets
 * vers [Context.filesDir] si ce n'est pas déjà fait.
 *
 * Si le modèle n'est ni dans les assets ni dans `filesDir` (saveur `fdroid` sans
 * téléchargement), [etat] vaut
 * [TranscripteurLocal.Etat.INDISPONIBLE] et toutes les méthodes de transcription
 * renvoient des valeurs vides/null sans planter.
 */
class TranscripteurVosk(private val context: Context) : TranscripteurLocal {

    companion object {
        private const val TAG = "TranscripteurVosk"
        private const val NOM_MODELE = "vosk-model-small-fr-0.22"
        private const val FREQUENCE = 16_000f

        private val json = Json { ignoreUnknownKeys = true }

        /**
         * Extrait la valeur de la clé "partial" du JSON Vosk.
         * Retourne null si le texte est vide ou le JSON invalide.
         */
        fun parsagePartiel(jsonStr: String): String? = runCatching {
            (json.parseToJsonElement(jsonStr).jsonObject["partial"] as? JsonPrimitive)
                ?.content
                ?.takeIf { it.isNotEmpty() }
        }.getOrNull()

        /**
         * Extrait la valeur de la clé "text" du JSON Vosk.
         * Retourne une chaîne vide si le JSON est invalide.
         */
        fun parsageTexte(jsonStr: String): String = runCatching {
            (json.parseToJsonElement(jsonStr).jsonObject["text"] as? JsonPrimitive)
                ?.content
                ?: ""
        }.getOrDefault("")
    }

    override val etat: TranscripteurLocal.Etat by lazy { detecterEtat() }

    private var modele: Model? = null
    private var reconnaisseur: Recognizer? = null

    private fun detecterEtat(): TranscripteurLocal.Etat {
        val dansAssets = runCatching {
            context.assets.list("")?.contains(NOM_MODELE) == true
        }.getOrDefault(false)
        // Saveur `fdroid` : rien dans les assets, mais le modèle a pu être téléchargé
        // depuis Paramètres → Modèles et décompressé ici par DepotModeles.
        val dansFilesDir = File(context.filesDir, NOM_MODELE)
            .let { it.isDirectory && !it.listFiles().isNullOrEmpty() }
        val disponible = dansAssets || dansFilesDir
        if (!disponible) {
            Log.w(TAG, "Modèle Vosk absent — transcription serveur requise tant qu'il n'est pas téléchargé")
        }
        return if (disponible) TranscripteurLocal.Etat.DISPONIBLE else TranscripteurLocal.Etat.INDISPONIBLE
    }

    /** Retourne le [Recognizer] prêt à l'emploi, ou null si modèle indisponible. */
    private fun obtenirReconnaisseur(): Recognizer? {
        if (etat == TranscripteurLocal.Etat.INDISPONIBLE) return null
        if (reconnaisseur != null) return reconnaisseur
        return runCatching {
            val repModele = extraireModele()
            val m = Model(repModele.absolutePath).also { modele = it }
            Recognizer(m, FREQUENCE).also { reconnaisseur = it }
        }.onFailure { e ->
            Log.e(TAG, "Impossible d'initialiser le moteur Vosk", e)
        }.getOrNull()
    }

    /** Retourne le répertoire du modèle dans filesDir, après extraction si nécessaire. */
    private fun extraireModele(): File {
        val dest = File(context.filesDir, NOM_MODELE)
        if (dest.isDirectory && dest.list()?.isNotEmpty() == true) return dest
        Log.i(TAG, "Extraction du modèle Vosk vers ${dest.absolutePath} …")
        copierRepertoireAssets(NOM_MODELE, dest)
        Log.i(TAG, "Modèle Vosk extrait.")
        return dest
    }

    private fun copierRepertoireAssets(cheminAsset: String, destDir: File) {
        destDir.mkdirs()
        val enfants = context.assets.list(cheminAsset) ?: return
        for (enfant in enfants) {
            val sousAsset = "$cheminAsset/$enfant"
            val sousDest = File(destDir, enfant)
            val sousEnfants = context.assets.list(sousAsset) ?: emptyArray()
            if (sousEnfants.isEmpty()) {
                // fichier feuille
                context.assets.open(sousAsset).use { input ->
                    sousDest.outputStream().use { input.copyTo(it) }
                }
            } else {
                copierRepertoireAssets(sousAsset, sousDest)
            }
        }
    }

    override fun accepterEchantillons(data: ShortArray, longueur: Int): String? {
        val rec = obtenirReconnaisseur() ?: return null
        val octets = shortArrayVersOctets(data, longueur)
        val complet = rec.acceptWaveForm(octets, octets.size)
        return if (complet) parsageTexte(rec.result) else parsagePartiel(rec.partialResult)
    }

    override fun finaliser(): String {
        val rec = reconnaisseur ?: return ""
        return parsageTexte(rec.finalResult)
    }

    override fun reinitialiser() {
        reconnaisseur?.close()
        reconnaisseur = null
    }

    /** Convertit un tableau de Short PCM en ByteArray little-endian. */
    private fun shortArrayVersOctets(data: ShortArray, longueur: Int): ByteArray {
        val octets = ByteArray(longueur * 2)
        for (i in 0 until longueur) {
            val s = data[i].toInt()
            octets[i * 2] = (s and 0xFF).toByte()
            octets[i * 2 + 1] = (s ushr 8 and 0xFF).toByte()
        }
        return octets
    }
}