package com.lila.wakeup.kws

import android.content.Context
import android.util.Log
import com.k2fsa.sherpa.onnx.FeatureConfig
import com.k2fsa.sherpa.onnx.KeywordSpotter
import com.k2fsa.sherpa.onnx.KeywordSpotterConfig
import com.k2fsa.sherpa.onnx.OnlineModelConfig
import com.k2fsa.sherpa.onnx.OnlineStream
import com.k2fsa.sherpa.onnx.OnlineTransducerModelConfig
import com.lila.wakeup.Config
import java.io.File
import java.io.FileOutputStream
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.util.concurrent.atomic.AtomicInteger

/**
 * 流式 Keyword Spotter 唤醒引擎(v3 架构)。
 *
 * 工作流:
 *  1. AudioCapture 喂 100ms PCM 帧
 *  2. KeywordSpotter (sherpa-onnx Zipformer-2 流式 KWS) 持续接收音频
 *  3. 模型实时匹配 keywords.txt 里的拼音规则
 *  4. 命中即触发 onWakeup 回调
 *
 * 优势:
 *  - 流式实时识别,延迟 100~200ms(对比 ASR 切段 1~2 秒)
 *  - 模型 5MB(SenseVoice 228MB),启动快、内存低
 *  - 专为短唤醒词训练,识别率高
 *  - 不再依赖 VAD / 能量门 / DC 去除
 *  - keywords.txt 多变体覆盖("你好丽啦/咪啦/丽拉..."),修改不用改代码
 *
 * keywords.txt 格式:
 *   <拼音 token1> <token2> ... @<显示名>
 *   例: n ǐ h ǎo l ì l ā @你好Lila
 *
 * @param onWakeup 命中回调(在 ASR 推理线程调用)
 */
class AsrEngine(
    private val context: Context,
    var onWakeup: ((keyword: String) -> Unit)? = null,
) {
    companion object {
        private const val TAG = "KwsService.Engine"
    }

    private var spotter: KeywordSpotter? = null
    private var stream: OnlineStream? = null

    private val dumpCounter = AtomicInteger(0)

    fun init() {
        Log.i(TAG, "[Engine] init start")

        val config = KeywordSpotterConfig(
            featConfig = FeatureConfig(sampleRate = Config.SAMPLE_RATE, featureDim = 80),
            modelConfig = OnlineModelConfig(
                transducer = OnlineTransducerModelConfig(
                    encoder = Config.KWS_ENCODER,
                    decoder = Config.KWS_DECODER,
                    joiner = Config.KWS_JOINER,
                ),
                tokens = Config.KWS_TOKENS,
                modelType = "zipformer2",
                numThreads = Config.KWS_NUM_THREADS,
                provider = "cpu",
            ),
            keywordsFile = Config.KWS_KEYWORDS_FILE,
            keywordsScore = Config.KWS_SCORE,
            keywordsThreshold = Config.KWS_THRESHOLD,
            numTrailingBlanks = 2,
        )
        spotter = KeywordSpotter(assetManager = context.assets, config = config)
        stream = spotter!!.createStream()

        Log.i(TAG, "[Engine] init done (KWS Zipformer + keywords.txt loaded)")
    }

    /**
     * 喂一帧 PCM 数据。流式识别,命中立即回调。
     */
    fun feedAudio(pcm: ShortArray) {
        val s = stream ?: return
        val sp = spotter ?: return

        // PCM short → float [-1, 1]
        val samples = FloatArray(pcm.size) { pcm[it] / 32768f }
        s.acceptWaveform(samples, Config.SAMPLE_RATE)

        // 持续 decode 直到没有新数据
        while (sp.isReady(s)) {
            sp.decode(s)
        }

        // 检查命中
        val result = sp.getResult(s)
        if (result.keyword.isNotEmpty()) {
            Log.i(TAG, "[KWS] HIT keyword=\"${result.keyword}\" tokens=${result.tokens.joinToString(",")}")
            onWakeup?.invoke(result.keyword)

            // 命中后 reset stream,等待下一次唤醒
            sp.reset(s)

            // 命中时 dump 最近 PCM 到 WAV(诊断用)
            if (Config.DUMP_HIT_WAV) {
                dumpHitWav(pcm)
            }
        }
    }

    /**
     * 重置识别状态。在 PAUSE/RESUME 切换时调用。
     */
    fun reset() {
        val s = stream ?: return
        val sp = spotter ?: return
        sp.reset(s)
        Log.i(TAG, "[Engine] reset (stream cleared)")
    }

    /**
     * 释放 native 资源。
     */
    fun release() {
        try {
            stream?.release()
            spotter?.release()
        } catch (e: Exception) {
            Log.e(TAG, "release failed: ${e.message}")
        }
        stream = null
        spotter = null
        Log.i(TAG, "[Engine] released")
    }

    // ============================================================
    // WAV dump(诊断用)
    // ============================================================

    /**
     * 命中时 dump 最近一帧 PCM 到 WAV。
     *
     * 路径: /sdcard/Android/data/com.lila.wakeup/files/lila_kws/hit_NNN.wav
     * 拉回: adb pull /sdcard/Android/data/com.lila.wakeup/files/lila_kws/ ./
     */
    private fun dumpHitWav(pcm: ShortArray) {
        try {
            val baseDir = context.getExternalFilesDir(null) ?: context.filesDir
            val dir = File(baseDir, "lila_kws")
            if (!dir.exists()) dir.mkdirs()

            val idx = dumpCounter.incrementAndGet()
            val file = File(dir, "hit_${idx.toString().padStart(3, '0')}.wav")

            val dataBytes = pcm.size * 2
            val totalBytes = 36 + dataBytes
            val byteRate = Config.SAMPLE_RATE * 2
            val header = ByteBuffer.allocate(44).order(ByteOrder.LITTLE_ENDIAN).apply {
                put("RIFF".toByteArray())
                putInt(totalBytes)
                put("WAVE".toByteArray())
                put("fmt ".toByteArray())
                putInt(16)
                putShort(1)
                putShort(1)
                putInt(Config.SAMPLE_RATE)
                putInt(byteRate)
                putShort(2)
                putShort(16)
                put("data".toByteArray())
                putInt(dataBytes)
            }.array()

            FileOutputStream(file).use { fos ->
                fos.write(header)
                val data = ByteBuffer.allocate(dataBytes).order(ByteOrder.LITTLE_ENDIAN)
                for (s in pcm) data.putShort(s)
                fos.write(data.array())
            }
            Log.i(TAG, "[Dump] hit WAV saved: ${file.absolutePath}")
        } catch (e: Exception) {
            Log.e(TAG, "[Dump] failed: ${e.message}", e)
        }
    }
}
