package com.lila.wakeup.kws

import android.content.Context
import android.util.Log
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

/**
 * sherpa-onnx KWS 引擎封装。
 *
 * 该类是对 sherpa-onnx 官方 [KeywordSpotter] Kotlin Wrapper 的薄封装，
 * 隐藏底层 API 细节，对外暴露简洁的 [process] 接口。
 *
 * 模型路径：见 [Config.MODEL_DIR]，从 APK assets 加载。
 *
 * ⚠️ 版本兼容：本骨架基于 sherpa-onnx v1.13.0+ 的 Kotlin API。
 * 若克隆的 sherpa-onnx 仓库 API 有差异（如类名变更），仅需调整本文件。
 *
 * 业务结果数据类。
 *
 * @param keyword 命中的唤醒词文本（如 "你好Lila"），未命中为空字符串
 * @param confidence 置信度（[0,1]）
 */
data class KwsResult(
    val keyword: String,
    val confidence: Float
) {
    val isHit: Boolean get() = keyword.isNotEmpty()
}

class KwsEngine(private val context: Context) {

    companion object {
        private const val TAG = "KwsService.Engine"
    }

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

    /**
     * 加载模型 + keywords.txt，初始化推理引擎。
     * 必须在 [process] 之前调用一次。
     */
    fun init() {
        val cfg = KeywordSpotterConfig(
            featConfig = com.k2fsa.sherpa.onnx.FeatureConfig(
                sampleRate = Config.SAMPLE_RATE,
                featureDim = 80
            ),
            modelConfig = OnlineModelConfig(
                transducer = OnlineTransducerModelConfig(
                    encoder = Config.MODEL_ENCODER,
                    decoder = Config.MODEL_DECODER,
                    joiner = Config.MODEL_JOINER
                ),
                tokens = Config.MODEL_TOKENS,
                modelType = "zipformer"
            ),
            keywordsFile = Config.MODEL_KEYWORDS,
            keywordsThreshold = Config.KWS_THRESHOLD_PRIMARY,
            keywordsScore = 1.5f,
            numTrailingBlanks = 2
        )

        spotter = KeywordSpotter(assetManager = context.assets, config = cfg)
        stream = spotter?.createStream()
        Log.i(TAG, "[Engine] init done, keywordsFile=${Config.MODEL_KEYWORDS}")
    }

    /**
     * 喂一帧 PCM 数据进引擎，返回是否命中。
     *
     * @param pcm 16-bit PCM short[] 数据，长度任意（一般 10ms 帧 = 160 samples）
     * @return 命中则 [KwsResult.isHit] 为 true
     */
    fun process(pcm: ShortArray): KwsResult {
        val sp = spotter ?: return KwsResult("", 0f)
        val st = stream ?: return KwsResult("", 0f)

        // sherpa-onnx 输入要求 [-1, 1] 范围 float
        val floatBuf = FloatArray(pcm.size) { pcm[it] / 32768f }
        st.acceptWaveform(floatBuf, Config.SAMPLE_RATE)

        while (sp.isReady(st)) {
            sp.decode(st)
        }

        val result = sp.getResult(st)
        val keyword = result.keyword

        return if (keyword.isNotEmpty()) {
            // 命中后必须 reset，否则下次不会再触发
            sp.reset(st)
            Log.i(TAG, "[Engine] hit keyword=$keyword")
            KwsResult(keyword, 1f)   // sherpa-onnx 不直接给 confidence，通过 keyword_score 间接控制
        } else {
            KwsResult("", 0f)
        }
    }

    /**
     * 释放 native 资源。在 Service.onDestroy 中调用。
     */
    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")
    }
}
