// HikoWasm.kt — Hikotest standalone loader (no SDK, no network).
// Ship release.wasm + manifest.json in your app's assets/ and call functions locally.
//
// Dependency (pure JVM, no NDK/JNI):
//   implementation("com.dylibso.chicory:runtime:1.5.3")
//
// Contract: Hikotest WASM ABI v1 (see the standalone guide for details).
// Note: hiko_alloc never frees; in long-lived processes with heavy string
// traffic, re-create the HikoWasm instance periodically if memory growth matters.

package com.example.hiko // change to your package

import android.content.Context
import com.dylibso.chicory.runtime.HostFunction
import com.dylibso.chicory.runtime.ImportValues
import com.dylibso.chicory.runtime.Instance
import com.dylibso.chicory.wasm.Parser
import com.dylibso.chicory.wasm.types.FunctionType
import com.dylibso.chicory.wasm.types.ValType
import org.json.JSONObject
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.security.KeyFactory
import java.security.MessageDigest
import java.security.Signature
import java.security.spec.X509EncodedKeySpec
import java.util.Base64

class HikoWasm private constructor(
    private val instance: Instance,
    private val signatures: Map<String, Signature>,
    /** Remote config values embedded in the manifest (empty if none). */
    val config: Map<String, Any>,
) {
    data class Param(val name: String, val type: String)
    data class Signature(val params: List<Param>, val returns: String)

    // ─── Integrity verification (docs/WASM_INTEGRITY.md §6) ──────────────────
    // Optional, opt-in. Verify the detached integrity.json (SHA-256 of release.wasm
    // + manifest.json, plus an Ed25519 signature over both) before instantiating.
    // Default OFF → byte-behaves as before. WARN = verify if present, log + run on
    // mismatch, skip silently if absent. ENFORCE = reject if absent or verify fails.
    // NOTE: Ed25519 via java.security requires Android API 33+ (or a Conscrypt/
    // BouncyCastle provider on older devices); on plain JVM it needs Java 15+.
    enum class VerifyMode { OFF, WARN, ENFORCE }
    data class VerifyResult(val ok: Boolean, val verified: Boolean, val reason: String?)

    companion object {
        /** Public key map — keyId → SPKI PEM. Mirror of src/lib/wasm/signing-keys.mjs.
         *  PUBLIC keys only; safe to ship in the app. Rotation: add the new keyId here. */
        private val PUBLIC_SIGNING_KEYS = mapOf(
            "hk-2026-07" to """
                -----BEGIN PUBLIC KEY-----
                MCowBQYDK2VwAyEAogmXuGBXvXMzYXnvBnBYgE8vpjJNTXe0Fcz6RDxoCD4=
                -----END PUBLIC KEY-----
            """.trimIndent(),
        )

        /** Fixed-format signing message — identical across all 3 SDKs + loaders. */
        private const val SIGN_MESSAGE_PREFIX = "hikotest.integrity.v1"

        /** Load release.wasm + manifest.json from the app's assets/ directory.
         *  Pass `verify` (WARN/ENFORCE) to check integrity.json before instantiating. */
        fun fromAssets(
            context: Context,
            wasmAsset: String = "release.wasm",
            manifestAsset: String = "manifest.json",
            verify: VerifyMode = VerifyMode.OFF,
            integrityAsset: String = "integrity.json",
        ): HikoWasm {
            val integrityJson = if (verify != VerifyMode.OFF) {
                runCatching { context.assets.open(integrityAsset).bufferedReader().readText() }.getOrNull()
            } else null
            return load(
                context.assets.open(wasmAsset).readBytes(),
                context.assets.open(manifestAsset).bufferedReader().readText(),
                verify,
                integrityJson,
            )
        }

        /** Instantiate from the raw bytes of release.wasm and the manifest.json text.
         *  `integrityJson` = raw text of the detached integrity.json asset (verify only). */
        fun load(
            wasmBytes: ByteArray,
            manifestJson: String,
            verify: VerifyMode = VerifyMode.OFF,
            integrityJson: String? = null,
        ): HikoWasm {
            if (verify != VerifyMode.OFF) {
                if (integrityJson == null) {
                    // §6: WARN skips a missing integrity.json silently; ENFORCE rejects.
                    require(verify != VerifyMode.ENFORCE) {
                        "Hikotest: integrity.json required in enforce mode but not provided"
                    }
                } else {
                    val result = verifyIntegrity(
                        JSONObject(integrityJson), wasmBytes, manifestJson.toByteArray(Charsets.UTF_8),
                    )
                    if (!result.ok) {
                        val msg = "Hikotest integrity check failed: ${result.reason}"
                        if (verify == VerifyMode.ENFORCE) error(msg)
                        System.err.println("[hiko] $msg (verify=warn → running anyway)")
                    }
                }
            }

            val manifest = JSONObject(manifestJson)
            require(manifest.optInt("abi") == 1) { "Unsupported Hikotest ABI: ${manifest.optInt("abi")}" }

            val signatures = buildMap {
                val fns = manifest.getJSONArray("functions")
                for (i in 0 until fns.length()) {
                    val f = fns.getJSONObject(i)
                    val sig = f.getJSONObject("signature")
                    val params = sig.getJSONArray("params")
                    put(
                        f.getString("name"),
                        Signature(
                            params = (0 until params.length()).map {
                                val p = params.getJSONObject(it)
                                Param(p.getString("name"), p.getString("type"))
                            },
                            returns = sig.getString("returns"),
                        ),
                    )
                }
            }
            val config = buildMap {
                val cfg = manifest.optJSONObject("config") ?: return@buildMap
                for (key in cfg.keys()) put(key, cfg.get(key))
            }

            // AssemblyScript modules import env.abort — instantiation fails without it.
            val abort = HostFunction(
                "env",
                "abort",
                FunctionType.of(listOf(ValType.I32, ValType.I32, ValType.I32, ValType.I32), listOf()),
            ) { _, _ -> throw IllegalStateException("wasm abort") }

            val instance = Instance.builder(Parser.parse(wasmBytes))
                .withImportValues(ImportValues.builder().addFunction(abort).build())
                .build()
            return HikoWasm(instance, signatures, config)
        }

        /** Mirror of integrity-core.mjs verifyIntegrity (client side, JCA). */
        fun verifyIntegrity(integrity: JSONObject, wasmBytes: ByteArray, manifestBytes: ByteArray): VerifyResult {
            if (integrity.optInt("integrity") != 1) {
                return VerifyResult(false, false, "integrity.json missing or unsupported version")
            }
            if (sha256Hex(wasmBytes) != integrity.getJSONObject("wasm").getString("sha256")) {
                return VerifyResult(false, false, "wasm.sha256 mismatch (transit corruption / tamper)")
            }
            if (sha256Hex(manifestBytes) != integrity.getJSONObject("manifest").getString("sha256")) {
                return VerifyResult(false, false, "manifest.sha256 mismatch (tamper)")
            }
            if (integrity.isNull("signature")) {
                return VerifyResult(true, false, "unsigned (Layer 1 — hash only)")
            }
            val keyId = integrity.optString("keyId")
            val pem = PUBLIC_SIGNING_KEYS[keyId]
                ?: return VerifyResult(false, false, "unknown/invalid keyId: $keyId")
            if (!verifyEd25519(buildSigningMessage(integrity), integrity.getString("signature"), pem)) {
                return VerifyResult(false, false, "Ed25519 signature failed (forged source?)")
            }
            return VerifyResult(true, true, null)
        }

        private fun sha256Hex(bytes: ByteArray): String =
            MessageDigest.getInstance("SHA-256").digest(bytes).joinToString("") { "%02x".format(it) }

        private fun buildSigningMessage(i: JSONObject): String = listOf(
            SIGN_MESSAGE_PREFIX,
            "bundleVersion=${i.optString("bundleVersion")}",
            "wasm.sha256=${i.getJSONObject("wasm").getString("sha256")}",
            "manifest.sha256=${i.getJSONObject("manifest").getString("sha256")}",
            "keyId=${i.optString("keyId")}",
        ).joinToString("\n")

        private fun verifyEd25519(message: String, signatureB64: String, spkiPem: String): Boolean = try {
            val der = Base64.getDecoder().decode(
                spkiPem.replace(Regex("-----[A-Z ]+-----"), "").replace(Regex("\\s"), ""),
            )
            val pub = KeyFactory.getInstance("Ed25519").generatePublic(X509EncodedKeySpec(der))
            Signature.getInstance("Ed25519").run {
                initVerify(pub)
                update(message.toByteArray(Charsets.UTF_8))
                verify(Base64.getDecoder().decode(signatureB64))
            }
        } catch (e: Exception) {
            false
        }
    }

    /** The function names available in this bundle. */
    fun functions(): Set<String> = signatures.keys

    /** Call a function by name; arguments follow the manifest signature order. */
    fun call(name: String, vararg args: Any): Any {
        val sig = signatures[name] ?: error("Function not in manifest: $name")
        require(args.size == sig.params.size) {
            "$name expects ${sig.params.size} argument(s), got ${args.size}"
        }
        val hasString = sig.returns == "string" || sig.params.any { it.type == "string" }

        // Chicory's calling convention is long[]: floats travel as raw f64 bits.
        // A string parameter becomes a UTF-8 (ptr, len) pair via hiko_alloc.
        val wasmArgs = ArrayList<Long>(sig.params.size + 1)
        sig.params.forEachIndexed { i, p ->
            when (p.type) {
                "string" -> {
                    val bytes = args[i].toString().toByteArray(Charsets.UTF_8)
                    val ptr = instance.export("hiko_alloc").apply(bytes.size.toLong())[0]
                    instance.memory().write(ptr.toInt(), bytes)
                    wasmArgs.add(ptr)
                    wasmArgs.add(bytes.size.toLong())
                }
                "float" -> wasmArgs.add((args[i] as Number).toDouble().toRawBits())
                "boolean" -> wasmArgs.add(if (args[i] == true) 1L else 0L)
                else -> wasmArgs.add((args[i] as Number).toLong()) // int
            }
        }

        // Signatures with strings are called through the hiko_run_<name> wrapper export.
        val export = instance.export(if (hasString) "hiko_run_$name" else name)
        val raw = export.apply(*wasmArgs.toLongArray())[0]

        return when (sig.returns) {
            "boolean" -> raw != 0L
            "float" -> Double.fromBits(raw)
            "string" -> {
                // String return = pointer to [u32 little-endian length][UTF-8 bytes].
                val len = ByteBuffer.wrap(instance.memory().readBytes(raw.toInt(), 4))
                    .order(ByteOrder.LITTLE_ENDIAN).int
                String(instance.memory().readBytes(raw.toInt() + 4, len), Charsets.UTF_8)
            }
            else -> raw.toInt() // int
        }
    }
}