kazeia/kazeia-android/app/src/main/java/com/kazeia/llm/LlmLoader.kt

package com.kazeia.llm

import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer
import java.io.File
import java.nio.file.Paths

// Chargeur LLM agnostique au format. Détecte par magic-byte et instancie le bon moteur :
//   - GGUF  -> GgufLlmEngine  (libkazeia_engine / libllama, CPU-i8mm). TOUS les modèles,
//             dont l'hybride Qwen3.5-4B (GatedDeltaNet) qui ne s'exporte PAS en .pte.
//   - .pte  -> PteLlmEngine   (ExecuTorch + QNN, NPU V79, lib kazeia_pte). Denses UNIQUEMENT
//             (Qwen3-4B, Qwen3-8B, Qwen2.5-7B, Qwen3Guard-4B). Chemin in-app autorisé
//             (libQnn Maven, sans root). RAM ~÷2, prefill NPU rapide ; decode ~= CPU.
//
// Parc validé on-device (17/06, SM8750/V79, decode tok/s) :
//   Qwen3-4B .pte 15.7 | Qwen3-8B .pte 10.5 | Qwen3Guard-4B .pte 17.2 | Qwen2.5-7B .pte 8.0
//   Qwen3.5-4B GGUF 16.6 (hybride, CPU-i8mm)

interface LlmEngine {
    fun generate(sys: String, usr: String, max: Int = 96): String
    fun generateStream(sys: String, usr: String, max: Int = 96, onToken: (String) -> Boolean)
    fun lastStats(): GenStats
    fun reset()
    fun release()
}

object LlmLoader {
    private fun magic(path: String): ByteArray =
        ByteArray(8).also { b -> File(path).inputStream().use { it.read(b) } }

    private fun ByteArray.has(off: Int, tag: String): Boolean =
        tag.indices.all { i -> off + i < size && this[off + i] == tag[i].code.toByte() }

    // GGUF -> "GGUF" aux octets 0-3. ExecuTorch .pte -> "ET12" à l'octet 4.
    // tokenizerPath : requis pour .pte (tokeniseur HF externe). Défaut = tokenizer.json voisin du .pte.
    fun load(path: String, ctx: Int = 4096, nThreads: Int = 6, tokenizerPath: String? = null): LlmEngine {
        val m = magic(path)
        return when {
            m.has(0, "GGUF") -> GgufLlmEngine(path, ctx, nThreads)
            m.has(4, "ET12") -> {
                val tk = tokenizerPath ?: (File(path).parent ?: ".") + "/tokenizer.json"
                PteLlmEngine(path, tk, seqLen = ctx)
            }
            else -> error("Format LLM inconnu (ni GGUF ni ExecuTorch .pte) : $path")
        }
    }
}

// Adaptateur GGUF : délègue au moteur natif existant (libkazeia_engine / libllama, CPU-i8mm).
class GgufLlmEngine(model: String, ctx: Int = 4096, nThreads: Int = 6) : LlmEngine {
    private val e = EngineLlmEngine(model, ctx, nThreads)
    override fun generate(sys: String, usr: String, max: Int) = e.generate(sys, usr, max)
    override fun generateStream(sys: String, usr: String, max: Int, onToken: (String) -> Boolean) =
        e.generateStream(sys, usr, max, onToken)
    override fun lastStats() = e.lastStats()
    override fun reset() = e.reset()
    override fun release() = e.release()
}

// JNI ExecuTorch+QNN. Lib `kazeia_pte` = wrapper natif du runner ExecuTorch (rebuild QAIRT 2.42).
// jniLibs/arm64-v8a doit aussi contenir : libqnn_executorch_backend.so (2.42) +
// libQnnHtp/System/Prepare/HtpV79Stub/HtpNetRunExtensions.so (QAIRT 2.42) + libQnnHtpV79Skel.so.
// Set figé : /opt/Kazeia/pte_prep_work/ship_runtime/.
internal object ExecuTorchJni {
    external fun load(ptePath: String, tokenizerPath: String, seqLen: Int, evalMode: Int): Long
    external fun generateFromIds(handle: Long, promptIds: LongArray, maxNewTokens: Int): LongArray
    external fun lastStats(handle: Long): LongArray   // [prefillMs, decodeMs, nGenerated]
    external fun free(handle: Long)
    init { System.loadLibrary("kazeia_pte") }
}

// Adaptateur ExecuTorch+QNN (.pte) pour les denses sur NPU V79.
// La tokenisation est faite ICI (côté Kotlin) car le tokeniseur natif est cassé pour Qwen
// (regex lookahead RE2/PCRE2) : on applique le template ChatML, on encode en IDs, et on passe
// les IDs au natif (eval_mode 1 hybride). Le natif décode greedy jusqu'à EOS et renvoie les IDs générés.
class PteLlmEngine(ptePath: String, tokenizerJson: String, private val seqLen: Int = 4096) : LlmEngine {
    private val tok = HuggingFaceTokenizer.newInstance(Paths.get(tokenizerJson))
    private val h = ExecuTorchJni.load(ptePath, tokenizerJson, seqLen, /*evalMode=hybrid*/ 1)

    init { require(h != 0L) { "échec chargement .pte : $ptePath" } }

    // Template ChatML Qwen ; les balises sont dans le texte -> encode sans special tokens auto.
    private fun chatml(sys: String, usr: String): LongArray {
        val p = buildString {
            if (sys.isNotEmpty()) append("<|im_start|>system\n").append(sys).append("<|im_end|>\n")
            append("<|im_start|>user\n").append(usr).append("<|im_end|>\n<|im_start|>assistant\n")
        }
        return tok.encode(p, /*addSpecialTokens=*/false, /*withOverflowingTokens=*/false).ids
    }

    override fun generate(sys: String, usr: String, max: Int): String {
        val outIds = ExecuTorchJni.generateFromIds(h, chatml(sys, usr), max)
        return tok.decode(outIds, /*skipSpecialTokens=*/true).trim()
    }

    // Streaming réel = variante native à callback de token (TODO côté JNI).
    // En attendant : génération bloquante puis émission en un bloc.
    override fun generateStream(sys: String, usr: String, max: Int, onToken: (String) -> Boolean) {
        onToken(generate(sys, usr, max))
    }

    override fun lastStats(): GenStats =
        ExecuTorchJni.lastStats(h).let { GenStats(it[0], it[1], it[2]) }

    override fun reset() { /* runner sans état persistant entre appels */ }
    override fun release() { ExecuTorchJni.free(h); tok.close() }
}