// Copyright (c) 2026, CIIL Anuvani-LID. Apache-2.0.
//
// A starting point for an Android app: one class over the Kotlin API of the patched sherpa-onnx (LidPipeline for
// whole recordings, AudioTagging for one stretch of speech). Copy it into your app and change the package. It is
// compiled against android.jar and run in ART on an emulator by tests/android (docs/03-android.md).
package com.example.anuvanilid

import android.os.Build
import android.system.Os
import com.k2fsa.sherpa.onnx.AudioTagging
import com.k2fsa.sherpa.onnx.AudioTaggingConfig
import com.k2fsa.sherpa.onnx.AudioTaggingModelConfig
import com.k2fsa.sherpa.onnx.LidPipeline
import com.k2fsa.sherpa.onnx.LidPipelineConfig
import org.json.JSONObject
import java.io.File

/** The model file and execution provider to run (docs/01-formats-and-devices.md). */
enum class Accelerator {
    /** anuvani_lid.int8.onnx on the CPU: every phone; the default. */
    CPU_INT8,

    /** anuvani_lid.fp32.onnx on the CPU: the reference numbers, 2.7 times the file size of int8. */
    CPU_FP32,

    /** The 15 s accelerator format on the Qualcomm Hexagon NPU (Snapdragon with Hexagon v69 or newer; needs the
     *  QNN plugin in the app: docs/03-android.md). */
    QNN,

    /** The 15 s NNAPI encoder through NNAPI (Android 8.1 to 14; deprecated from Android 15). */
    NNAPI,
}

data class LanguageScore(val language: String, val probability: Double)

data class Segment(val startS: Double, val endS: Double, val language: String, val probability: Double)

/** The result for a recording: the anuvani-lid/1 JSON of the CIIL Anuvani-LID server, parsed. */
data class LidResult(
    /** The recording's language (null when there is no speech). */
    val language: String?,
    /** exp(duration-weighted mean log-probability) of that language. */
    val confidence: Double,
    val topK: List<LanguageScore>,
    /** The language with the most speech time. */
    val dominantByTime: String?,
    val timeShare: Map<String, Double>,
    /** One entry per classified piece of speech, in order. */
    val segments: List<Segment>,
    val speechSeconds: Double,
    val json: String,
)

class LanguageIdentifier(
    /** Directory holding the model files (stored as files, not compressed assets: docs/03-android.md). */
    private val modelDir: File,
    /** A writable directory of the app (e.g. context.filesDir): provider config files, the NPU-compiled graph. */
    private val workDir: File,
    val accelerator: Accelerator = Accelerator.CPU_INT8,
    private val threads: Int = 4,
    /** For QNN: the app's nativeLibraryDir (context.applicationInfo.nativeLibraryDir), where the plugin's
     *  libraries are; the NPU loads its libraries from there. */
    private val nativeLibraryDir: String? = null,
) : AutoCloseable {

    private val modelConfig: AudioTaggingModelConfig by lazy {
        val (file, provider) = when (accelerator) {
            Accelerator.CPU_INT8 -> "anuvani_lid.int8.onnx" to "cpu"
            Accelerator.CPU_FP32 -> "anuvani_lid.fp32.onnx" to "cpu"
            Accelerator.NNAPI -> "anuvani_lid.npu-15s-nnapi.encoder.onnx" to providerConfig("nnapi", "NNAPI.fp16=1\n")
            Accelerator.QNN -> "anuvani_lid.npu-15s.encoder.onnx" to providerConfig("qnn", qnnConfig())
        }
        AudioTaggingModelConfig(lid = File(modelDir, file).path, numThreads = threads, provider = provider)
    }

    // created on first use: an app that only identifies recordings, or only classifies segments, loads one model
    private val pipelineDelegate = lazy {
        LidPipeline(
            config = LidPipelineConfig(
                model = modelConfig,
                vadModel = File(modelDir, "silero_vad_6.2.1.onnx").path,
            ),
        )
    }
    private val taggerDelegate = lazy { AudioTagging(config = AudioTaggingConfig(model = modelConfig, topK = 23)) }

    /** The language(s) of a recording of any length: Silero VAD, speech pieces, one decision per piece, the
     *  duration-weighted decision; the procedure and results of the CIIL Anuvani-LID server. samples: mono in
     *  [-1, 1] at sampleRate (resampled to 16 kHz if different). Blocking; call it off the main thread. */
    fun identify(samples: FloatArray, sampleRate: Int): LidResult =
        parse(pipelineDelegate.value.process(samples, sampleRate))

    /** Probabilities of the languages for one stretch of speech (e.g. a VAD segment), largest first. Audio shorter
     *  than 0.1 s gives an empty list. Blocking. */
    fun classify(samples: FloatArray, sampleRate: Int, topK: Int = 23): List<LanguageScore> {
        val tagger = taggerDelegate.value
        val stream = tagger.createStream()
        stream.acceptWaveform(samples, sampleRate)
        val events = tagger.compute(stream, topK)
        stream.release()
        return events.map { LanguageScore(it.name, it.prob.toDouble()) }
    }

    override fun close() {
        if (pipelineDelegate.isInitialized()) pipelineDelegate.value.release()
        if (taggerDelegate.isInitialized()) taggerDelegate.value.release()
    }

    private fun providerConfig(name: String, text: String): String {
        val f = File(workDir, "anuvani-lid-$name.cfg")
        f.writeText(text)
        return "$name:${f.path}"
    }

    private fun qnnConfig(): String {
        nativeLibraryDir?.let {
            // the NPU side of QNN (FastRPC) finds the libQnnHtpV*Skel.so libraries through ADSP_LIBRARY_PATH
            Os.setenv("ADSP_LIBRARY_PATH", "$it;/system/lib/rfsa/adsp;/system/vendor/lib/rfsa/adsp;/dsp", true)
        }
        val cache = File(workDir, "qnn").apply { mkdirs() }
        val sb = StringBuilder()
        sb.append("QNN.library=libonnxruntime_providers_qnn.so\n")
        sb.append("QNN.backend_type=htp\n")
        sb.append("QNN.htp_performance_mode=burst\n")
        sb.append("QNN.cache_dir=${cache.path}\n")
        qnnSocModel()?.let {
            sb.append("QNN.soc_model=$it\n")
            sb.append("QNN.enable_htp_fp16_precision=1\n")
        }
        return sb.toString()
    }

    companion object {
        /** Qualcomm chips the QNN plugin knows, with their soc_model IDs and Hexagon generations (onnxruntime-qnn,
         *  onnxruntime/core/providers/qnn/soc_utility/soc_table.h). */
        val QNN_CHIPS: Map<String, Pair<Int, Int>> = mapOf(
            "SM8350" to (30 to 68), "SM8325" to (34 to 68), "SM8450" to (36 to 69), "SM8475" to (42 to 69),
            "SM8550" to (43 to 73), "SM8650" to (57 to 75), "SM8635" to (68 to 73), "SM8750" to (69 to 79),
            "SM7675" to (70 to 73), "SM8850" to (87 to 81),
        )

        /** The chip name Android reports (Android 12+), e.g. "SM8650"; "" before Android 12. */
        fun socModel(): String = if (Build.VERSION.SDK_INT >= 31) Build.SOC_MODEL ?: "" else ""

        /** soc_model of this phone's chip for the QNN plugin, or null if the table does not list it. */
        fun qnnSocModel(): Int? = QNN_CHIPS[socModel().uppercase()]?.first

        /** True on Snapdragon chips whose NPU builds the accelerator format in fp16: Hexagon v69 (Snapdragon 8 Gen 1)
         *  or newer. On Qualcomm's emulated NPU, v68 (Snapdragon 888) cannot compile it; if QNN is chosen there
         *  anyway, the model runs on the CPU (docs/01). */
        fun qnnFloatCapable(): Boolean = (QNN_CHIPS[socModel().uppercase()]?.second ?: 0) >= 69

        /** Parse the anuvani-lid/1 JSON of LidPipeline.process(). */
        fun parse(json: String): LidResult {
            val j = JSONObject(json)
            val overall = j.getJSONObject("overall")
            val topK = overall.getJSONArray("top_k").let { a ->
                (0 until a.length()).map { LanguageScore(a.getJSONObject(it).getString("label"), a.getJSONObject(it).getDouble("prob")) }
            }
            val share = j.getJSONObject("language_time_share").let { o -> o.keys().asSequence().associateWith { o.getDouble(it) } }
            val segments = j.getJSONArray("segments").let { a ->
                (0 until a.length()).map {
                    val s = a.getJSONObject(it)
                    Segment(s.getDouble("start_s"), s.getDouble("end_s"), s.getString("label"), s.getDouble("prob"))
                }
            }
            return LidResult(
                language = if (overall.isNull("label")) null else overall.getString("label"),
                confidence = overall.getDouble("confidence"),
                topK = topK,
                dominantByTime = if (j.isNull("dominant_by_time")) null else j.getString("dominant_by_time"),
                timeShare = share,
                segments = segments,
                speechSeconds = j.getJSONObject("stats").getDouble("speech_seconds"),
                json = json,
            )
        }
    }
}