package com.k2fsa.sherpa.onnx

import android.util.Log

// ============================================================================
// LlmEngine — libsummarize-llm-jni.so(llama.cpp + 薄 JNI)的 Kotlin 封装。
//
// 与工程其它 native wrapper(OnlineRecognizer/OfflineRecognizer/Extractor)
// 同范式:ptr + require + finalize()/release()。但注意:
//  * 本库绝不能进录音/主线程路径 —— generate() 是长阻塞调用,设计为只在
//    摘要专用后台线程上使用,一次只能有一个 generate 在跑(JNI 内有锁)。
//  * 模型文件必须由用户 adb push,App 不下载;调用 nativeNewFromFile 前
//    必须在 Kotlin 侧确认文件存在(工程教训:缺文件时有的 native 会 exit(),
//    本 .so 已做 access() 双保险并抛异常,但仍按规矩先查)。
// ============================================================================

/** 摘要 LLM 的可用性探测。so 只编了 arm64-v8a,x86_64 模拟器上自动不可用。 */
object LlmRuntime {
    private const val LIB = "summarize-llm-jni"
    @Volatile private var tried = false
    @Volatile var nativeAvailable = false
        private set

    fun ensure(): Boolean {
        if (!tried) {
            synchronized(this) {
                if (!tried) {
                    tried = true
                    nativeAvailable = try {
                        System.loadLibrary(LIB)
                        true
                    } catch (t: Throwable) {
                        Log.w("sherpa-onnx", "summarize lib $LIB not loadable: $t")
                        false
                    }
                }
            }
        }
        return nativeAvailable
    }
}

class LlmEngine(
    modelPath: String,
    ctxSize: Int = 8192,       // 滚动笔记+单块输入远小于此(§3 资源账)
    nThreads: Int = 6,         // 8 Gen3 大核 4 个,给摘要留余量:6 线程够用
    kvTypeK: String = "q8_0",  // KV 用 q8_0 压显存/内存(§3)
    kvTypeV: String = "f16",   // V 保持 f16:K 量化到 q8_0 是 llama.cpp 推荐档
    chatTemplate: String = "gemma4",  // "gemma4" | "gemma3",须与 GGUF 内嵌模板一致
) {
    companion object {
        private const val TAG = "sherpa-onnx"

        // ---- Gemma 4 canonical template 的单轮(system+user)渲染 ----
        // 与 GGUF 内嵌模板(enable_thinking=false, 无 tools)逐字等价:
        //   bos <|turn>system\nSYS<turn|>\n <|turn>user\nUSR<turn|>\n <|turn>model\n
        // bos 由 tokenizer 的 add_bos_token 自动添加(native 侧 parse_special=true)。
        fun renderGemma4Prompt(system: String, user: String): String =
            buildString {
                append("<|turn>system\n").append(system.trim()).append("<turn|>\n")
                append("<|turn>user\n").append(user.trim()).append("<turn|>\n")
                append("<|turn>model\n")
            }

        // ---- Gemma 3 官方模板的单轮渲染(ggml-org/gemma-3-1b-it-GGUF)----
        // 官方 jinja:system 不作为独立 turn,而是拼进第一个 user turn;
        //   <start_of_turn>user\n{system}\n\n{user}<end_of_turn>\n<start_of_turn>model\n
        // EOG token 是 <end_of_turn>(gemma3 的 eot),依赖 native 按 eog 停止。
        fun renderGemma3Prompt(system: String, user: String): String =
            buildString {
                append("<start_of_turn>user\n")
                append(system.trim()).append("\n\n")
                append(user.trim())
                append("<end_of_turn>\n")
                append("<start_of_turn>model\n")
            }

        // ---- MiniCPM5(ChatML,生成 prompt 默认带 think 块)----
        // GGUF 模板 enable_thinking=false 时的 generation prompt 为:
        //   <|im_start|>assistant\n<think>\n\n</think>\n\n
        // 显式预填空 think → 模型直接输出正文,不浪费解码在思考上。
        // EOG 为 <|im_end|>(llama.cpp 官方 GGUF,元数据已配)。
        fun renderMinicpmPrompt(system: String, user: String): String =
            buildString {
                append("<|im_start|>system\n").append(system.trim()).append("<|im_end|>\n")
                append("<|im_start|>user\n").append(user.trim()).append("<|im_end|>\n")
                append("<|im_start|>assistant\n<think>\n\n</think>\n\n")
            }

        fun renderChatPrompt(system: String, user: String,
                             template: String = "gemma4"): String =
            when (template) {
                "gemma3" -> renderGemma3Prompt(system, user)
                "minicpm" -> renderMinicpmPrompt(system, user)
                else -> renderGemma4Prompt(system, user)
            }
    }

    private val chatTemplate: String = chatTemplate
    private var ptr: Long = 0L

    init {
        require(LlmRuntime.ensure()) { "摘要推理库不可用(设备 ABI 不支持或未打包)" }
        require(existsFile(modelPath)) { "找不到模型文件: $modelPath" }
        ptr = nativeNewFromFile(modelPath, ctxSize, nThreads, kvTypeK, kvTypeV)
        require(ptr != 0L) { "LLM 模型加载失败: $modelPath" }
    }

    // 与工程教训对齐的存在性检查:加载任何模型文件前必须先确认文件在
    private fun existsFile(path: String): Boolean =
        java.io.File(path).let { it.isFile && it.length() > 100L }

    /**
     * 阻塞式生成(仅供摘要后台线程调用)。返回模型输出文本(不含输入)。
     * temperature=0 为 Greedy。abort() 可从其它线程安全打断。
     * repeatPenalty>1 启用重复惩罚(翻译路径传 1.1 治 1B 复读循环;
     * 摘要路径保持默认 0=关闭,行为与已验证版本一致)。
     */
    fun generate(
        systemPrompt: String,
        userPrompt: String,
        maxNewTokens: Int = 900,
        temperature: Float = 0f,
        topK: Int = 64,
        topP: Float = 0.95f,
        repeatPenalty: Float = 0f,
    ): String {
        require(ptr != 0L) { "LlmEngine already released" }
        return nativeGenerate(
            ptr, renderChatPrompt(systemPrompt, userPrompt, chatTemplate),
            maxNewTokens, temperature, topK, topP, repeatPenalty
        )
    }

    fun abort() {
        if (ptr != 0L) nativeAbort(ptr)
    }

    /** 模型内嵌的 jinja 聊天模板(调试用,验证渲染一致性)。 */
    fun chatTemplate(): String? {
        if (ptr == 0L) return null
        return nativeGetStringMeta(ptr, "tokenizer.chat_template")
    }

    protected fun finalize() {
        if (ptr != 0L) {
            nativeFree(ptr)   // JNI 内部等待在途 generate 结束后再释放
            ptr = 0L
        }
    }

    fun release() = finalize()

    // ---- native(实现在 tools/llama-jni/src/summarize_llm_jni.cpp)----
    private external fun nativeNewFromFile(
        path: String, ctxSize: Int, nThreads: Int, typeK: String, typeV: String
    ): Long

    private external fun nativeGenerate(
        ptr: Long, prompt: String, maxNewTokens: Int,
        temperature: Float, topK: Int, topP: Float, repeatPenalty: Float
    ): String

    private external fun nativeAbort(ptr: Long)
    private external fun nativeFree(ptr: Long)
    private external fun nativeGetStringMeta(ptr: Long, key: String): String?
}