package com.k2fsa.sherpa.onnx

import android.content.res.AssetManager
import android.util.Log
import kotlin.math.sqrt

// ---------------------------------------------------------------------------
// Native speaker-embedding extractor (wraps libsherpa-onnx-jni).
// Ported from android/SherpaOnnxSpeakerIdentification so both apps stay
// compatible with the same JNI symbols.
// ---------------------------------------------------------------------------
class SpeakerEmbeddingExtractor(
    assetManager: AssetManager? = null,
    config: SpeakerEmbeddingExtractorConfig,
) {
    private var ptr: Long

    init {
        ptr = if (assetManager != null) {
            newFromAsset(assetManager, config)
        } else {
            newFromFile(config)
        }
        require(ptr != 0L) {
            "Invalid SpeakerEmbeddingExtractorConfig: failed to create native SpeakerEmbeddingExtractor"
        }
    }

    protected fun finalize() {
        if (ptr != 0L) {
            delete(ptr)
            ptr = 0
        }
    }

    fun release() = finalize()

    fun createStream(): OnlineStream {
        val p = createStream(ptr)
        return OnlineStream(p)
    }

    fun isReady(stream: OnlineStream) = isReady(ptr, stream.ptr)
    fun compute(stream: OnlineStream) = compute(ptr, stream.ptr)
    fun dim() = dim(ptr)

    private external fun newFromAsset(
        assetManager: AssetManager,
        config: SpeakerEmbeddingExtractorConfig,
    ): Long

    private external fun newFromFile(config: SpeakerEmbeddingExtractorConfig): Long
    private external fun delete(ptr: Long)
    private external fun createStream(ptr: Long): Long
    private external fun isReady(ptr: Long, streamPtr: Long): Boolean
    private external fun compute(ptr: Long, streamPtr: Long): FloatArray
    private external fun dim(ptr: Long): Int

    companion object {
        init {
            System.loadLibrary("sherpa-onnx-jni")
        }
    }
}

// ---------------------------------------------------------------------------
// SpeakerRegistry: lightweight in-Kotlin speaker store + cosine matching.
//
// We avoid the native SpeakerEmbeddingManager on purpose: it cannot rename a
// speaker once enrolled, which we need for the "auto-label then assign a real
// name" flow. Matching here prefers named entries over auto (S0/S1/...) ones.
// ---------------------------------------------------------------------------
class SpeakerRegistry {
    private val named = ArrayList<Pair<String, FloatArray>>()
    private val auto = ArrayList<Pair<String, FloatArray>>()

    val namedNames: List<String>
        get() = named.map { it.first }

    // The rolling window checker runs on its own thread while the endpoint
    // identification runs on the recording thread: serialize everything.
    @Synchronized
    fun clear() {
        named.clear()
        auto.clear()
    }

    // Returns the best label for embedding `e`, or null if no speaker is close
    // enough (below threshold). Registered names win ties against auto labels.
    @Synchronized
    fun search(e: FloatArray, threshold: Float): String? {
        match(named, e, threshold)?.let { return it }
        match(auto, e, threshold)?.let { return it }
        return null
    }

    @Synchronized
    fun addAuto(name: String, e: FloatArray) {
        auto.add(Pair(name, l2normalize(e)))
    }

    // Enroll (or overwrite) a real name for a given embedding.
    @Synchronized
    fun addNamed(name: String, e: FloatArray) {
        val v = l2normalize(e)
        named.removeAll { it.first == name }
        // Drop any auto entry whose embedding now equals this named one so the
        // same voice is not also reported as an anonymous Sx id.
        named.add(Pair(name, v))
    }

    /**
     * Best cosine similarity of `e` against all enrolled speakers (0..1 scale,
     * may be slightly noisy). Used by the rolling checker to tell "no idea
     * whose voice this is" from "definitely the current speaker".
     */
    @Synchronized
    fun bestScore(e: FloatArray): Float {
        val q = l2normalize(e)
        var best = -1.0f
        for ((_, emb) in named) best = maxOf(best, dot(q, emb))
        for ((_, emb) in auto) best = maxOf(best, dot(q, emb))
        return best
    }

    /**
     * 最近邻兜底(收紧版,P1):按"名字级"聚合所有嵌入,返回最优名与
     * 名级最高/次高分。同人多嵌入只算一个竞争者,优势差判定才有意义;
     * 调用方用 score≥minScore 且 (score-secondScore)≥margin 才归入,
     * 模糊时判新人(宁多勿错)。
     */
    class Rank(val name: String?, val score: Float, val secondScore: Float)

    @Synchronized
    fun nearestRank(e: FloatArray): Rank {
        val q = l2normalize(e)
        val byName = LinkedHashMap<String, Float>()
        for ((name, emb) in named) {
            val s = dot(q, emb)
            if (s > (byName[name] ?: -2f)) byName[name] = s
        }
        for ((name, emb) in auto) {
            val s = dot(q, emb)
            if (s > (byName[name] ?: -2f)) byName[name] = s
        }
        var best: String? = null
        var bestScore = -2f
        var secondScore = -2f
        for ((name, s) in byName) {
            if (s > bestScore) { secondScore = bestScore; bestScore = s; best = name }
            else if (s > secondScore) secondScore = s
        }
        return Rank(best, bestScore, secondScore)
    }

    /**
     * 点击标签命名用:把内部键 S{n}(自动注册、未命名)升级为真名,
     * 或把已有真名改成另一个真名。同名允许多条嵌入并存 —— 之前误判成
     * 两个说话人、后被命名为同一名字的两组声纹,自此 search 都返回同一名。
     */
    @Synchronized
    fun renameToName(oldKey: String, name: String): Boolean {
        val ai = auto.indexOfFirst { it.first == oldKey }
        if (ai >= 0) {
            named.add(Pair(name, auto[ai].second))
            auto.removeAt(ai)
            return true
        }
        var changed = false
        for (i in named.indices) {
            if (named[i].first == oldKey) {
                named[i] = Pair(name, named[i].second)
                changed = true
            }
        }
        return changed
    }

    private fun match(
        list: List<Pair<String, FloatArray>>,
        e: FloatArray,
        threshold: Float,
    ): String? {
        if (list.isEmpty()) return null
        val q = l2normalize(e)
        var bestName: String? = null
        var bestScore = -1.0f
        for ((name, emb) in list) {
            val s = dot(q, emb)
            if (s > bestScore) {
                bestScore = s
                bestName = name
            }
        }
        return if (bestScore >= threshold) bestName else null
    }

    private fun l2normalize(v: FloatArray): FloatArray {
        var sum = 0.0f
        for (x in v) sum += x * x
        if (sum <= 0f) return v.copyOf()
        val inv = 1.0f / sqrt(sum)
        return FloatArray(v.size) { v[it] * inv }
    }

    private fun dot(a: FloatArray, b: FloatArray): Float {
        var s = 0.0f
        val n = minOf(a.size, b.size)
        for (i in 0 until n) s += a[i] * b[i]
        return s
    }
}

// ---------------------------------------------------------------------------
// Singleton holder shared by the UI thread, recording thread and the rolling
// speaker-check thread.
//
// 双模型方案(E):准确模型(eres2netv2)负责 endpoint 权威判定与注册;
// 轻量模型(eres2net base,可选)只负责滚动换人检测(推理快 ~4 倍)。
// ⚠ 两个模型的 embedding 不互通:各自配一份 registry,注册时双写镜像。
// ---------------------------------------------------------------------------
object SpeakerRecognition {
    @Volatile private var _extractor: SpeakerEmbeddingExtractor? = null
    @Volatile private var _fastExtractor: SpeakerEmbeddingExtractor? = null
    val extractor: SpeakerEmbeddingExtractor
        get() = _extractor!!

    /** accurate-space registry (endpoint verdict / naming source of truth) */
    val registry = SpeakerRegistry()
    /** fast-space mirror (rolling-check search only; may stay empty) */
    val fastRegistry = SpeakerRegistry()

    var autoCounter: Int = 0
        private set

    private val accLock = Any()
    private val fastLock = Any()

    // autoCounter is touched from worker threads while resetSpeakers runs on
    // the UI thread — keep them serialized.
    @Synchronized
    fun nextAutoLabel(): String = "S${autoCounter++}"

    @Synchronized
    fun resetSpeakers() {
        registry.clear()
        fastRegistry.clear()
        autoCounter = 0
    }

    // assetManager = null -> `model` is treated as an absolute file system path
    fun initExtractor(assetManager: AssetManager?, model: String) {
        synchronized(accLock) {
            if (_extractor != null) return
            Log.i(TAG, "Initializing speaker embedding extractor with $model")
            _extractor = SpeakerEmbeddingExtractor(
                assetManager = assetManager,
                config = SpeakerEmbeddingExtractorConfig(
                    model = model,
                    numThreads = 2,
                    debug = false,
                    provider = "cpu",
                )
            )
        }
    }

    // Optional lightweight model for the rolling speaker-change detector.
    fun initFastExtractor(assetManager: AssetManager?, model: String) {
        synchronized(fastLock) {
            if (_fastExtractor != null) return
            Log.i(TAG, "Initializing FAST speaker embedding extractor with $model")
            _fastExtractor = SpeakerEmbeddingExtractor(
                assetManager = assetManager,
                config = SpeakerEmbeddingExtractorConfig(
                    model = model,
                    numThreads = 1,
                    debug = false,
                    provider = "cpu",
                )
            )
        }
    }

    fun fastReady(): Boolean = _fastExtractor != null

    // Compute one embedding with the accurate extractor (endpoint verdicts).
    fun computeEmbedding(samples: FloatArray, sampleRate: Int): FloatArray? =
        synchronized(accLock) {
            val ext = _extractor ?: return null
            runCompute(ext, samples, sampleRate)
        }

    // Rolling-check embedding: fast space when available, else accurate space
    // (in the latter case both registries see the same space anyway).
    fun computeEmbeddingFast(samples: FloatArray, sampleRate: Int): FloatArray? {
        val fast = _fastExtractor
        return if (fast == null) computeEmbedding(samples, sampleRate)
        else synchronized(fastLock) { runCompute(fast, samples, sampleRate) }
    }

    private fun runCompute(
        ext: SpeakerEmbeddingExtractor,
        samples: FloatArray,
        sampleRate: Int,
    ): FloatArray? {
        if (samples.isEmpty()) return null
        val stream = ext.createStream()
        stream.acceptWaveform(samples = samples, sampleRate = sampleRate)
        stream.inputFinished()
        val emb = if (ext.isReady(stream)) ext.compute(stream) else null
        stream.release()
        return emb
    }

    // Cosine similarity between two embeddings (L2-normalized internally),
    // in [-1, 1]. Pure function, safe from any thread.
    fun similarity(a: FloatArray, b: FloatArray): Float {
        val n = minOf(a.size, b.size)
        if (n == 0) return 0f
        var dot = 0.0
        var na = 0.0
        var nb = 0.0
        for (i in 0 until n) {
            dot += a[i].toDouble() * b[i]
            na += a[i].toDouble() * a[i]
            nb += b[i].toDouble() * b[i]
        }
        if (na <= 0.0 || nb <= 0.0) return 0f
        return (dot / (kotlin.math.sqrt(na) * kotlin.math.sqrt(nb))).toFloat()
    }

    private const val TAG = "sherpa-onnx-speaker"
}