// HikoWasm.swift — Hikotest standalone loader (no SDK, no network).
// Ship release.wasm + manifest.json as bundle resources and call functions locally.
//
// Dependency (Swift Package Manager — pure Swift, no binary frameworks):
//   .package(url: "https://github.com/swiftwasm/WasmKit.git", from: "0.1.5")
//   product: "WasmKit"
//
// 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.

import Foundation
import CryptoKit
import WasmKit

final class HikoWasm {
    enum HikoError: Error {
        case badManifest
        case unsupportedABI(Int)
        case unknownFunction(String)
        case badArguments(String)
        case missingExport(String)
        case aborted
        case integrityFailed(String)
    }

    struct Param { let name: String; let type: String }
    struct Signature { let params: [Param]; let 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.
    enum VerifyMode { case off, warn, enforce }
    struct VerifyResult { let ok: Bool; let verified: Bool; let reason: String? }

    /// 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 static let publicSigningKeys: [String: String] = [
        "hk-2026-07": """
        -----BEGIN PUBLIC KEY-----
        MCowBQYDK2VwAyEAogmXuGBXvXMzYXnvBnBYgE8vpjJNTXe0Fcz6RDxoCD4=
        -----END PUBLIC KEY-----
        """,
    ]

    /// Fixed-format signing message — identical across all 3 SDKs + loaders.
    private static let signMessagePrefix = "hikotest.integrity.v1"

    private let instance: Instance
    private let memory: Memory
    private let signatures: [String: Signature]
    /// Remote config values embedded in the manifest (empty if none).
    let config: [String: Any]

    /// Load release.wasm + manifest.json from the app bundle.
    /// Pass `verify` (.warn/.enforce) to check integrity.json before instantiating.
    convenience init(
        bundle: Bundle = .main,
        wasmResource: String = "release",
        manifestResource: String = "manifest",
        verify: VerifyMode = .off,
        integrityResource: String = "integrity"
    ) throws {
        guard let wasmURL = bundle.url(forResource: wasmResource, withExtension: "wasm"),
              let manifestURL = bundle.url(forResource: manifestResource, withExtension: "json")
        else { throw HikoError.missingExport("release.wasm / manifest.json not found in bundle") }
        let integrityJSON: Data? = verify != .off
            ? bundle.url(forResource: integrityResource, withExtension: "json").flatMap { try? Data(contentsOf: $0) }
            : nil
        try self.init(
            wasmBytes: Data(contentsOf: wasmURL),
            manifestJSON: Data(contentsOf: manifestURL),
            verify: verify,
            integrityJSON: integrityJSON
        )
    }

    /// Instantiate from the raw bytes of release.wasm and manifest.json.
    /// `integrityJSON` = raw bytes of the detached integrity.json asset (verify only).
    init(wasmBytes: Data, manifestJSON: Data, verify: VerifyMode = .off, integrityJSON: Data? = nil) throws {
        if verify != .off {
            if let integrityJSON = integrityJSON {
                let result = HikoWasm.verifyIntegrity(integrityJSON, wasmBytes: wasmBytes, manifestBytes: manifestJSON)
                if !result.ok {
                    let msg = "Hikotest integrity check failed: \(result.reason ?? "")"
                    if verify == .enforce { throw HikoError.integrityFailed(msg) }
                    FileHandle.standardError.write(Data("[hiko] \(msg) (verify=warn → running anyway)\n".utf8))
                }
            } else if verify == .enforce {
                throw HikoError.integrityFailed("integrity.json required in enforce mode but not provided")
            }
        }

        guard let manifest = try JSONSerialization.jsonObject(with: manifestJSON) as? [String: Any],
              let functions = manifest["functions"] as? [[String: Any]]
        else { throw HikoError.badManifest }
        let abi = manifest["abi"] as? Int ?? -1
        guard abi == 1 else { throw HikoError.unsupportedABI(abi) }

        var signatures: [String: Signature] = [:]
        for f in functions {
            guard let name = f["name"] as? String,
                  let sig = f["signature"] as? [String: Any],
                  let params = sig["params"] as? [[String: Any]],
                  let returns = sig["returns"] as? String else { continue }
            signatures[name] = Signature(
                params: params.compactMap { p in
                    guard let n = p["name"] as? String, let t = p["type"] as? String else { return nil }
                    return Param(name: n, type: t)
                },
                returns: returns
            )
        }
        self.signatures = signatures
        self.config = (manifest["config"] as? [String: Any]) ?? [:]

        let store = Store(engine: Engine())
        var imports = Imports()
        // AssemblyScript modules import env.abort — instantiation fails without it.
        imports.define(
            module: "env", name: "abort",
            Function(store: store, parameters: [.i32, .i32, .i32, .i32], results: []) { _, _ in
                throw HikoError.aborted
            }
        )
        let module = try parseWasm(bytes: [UInt8](wasmBytes))
        instance = try module.instantiate(store: store, imports: imports)
        guard let memory = instance.exports[memory: "memory"] else {
            throw HikoError.missingExport("memory")
        }
        self.memory = memory
    }

    /// The function names available in this bundle.
    var functions: [String] { Array(signatures.keys) }

    /// Call a function by name; arguments follow the manifest signature order.
    /// Pass Int / Double / Bool / String values; returns one of the same.
    @discardableResult
    func call(_ name: String, _ args: Any...) throws -> Any {
        guard let sig = signatures[name] else { throw HikoError.unknownFunction(name) }
        guard args.count == sig.params.count else {
            throw HikoError.badArguments("\(name) expects \(sig.params.count) argument(s), got \(args.count)")
        }
        let hasString = sig.returns == "string" || sig.params.contains { $0.type == "string" }

        var wasmArgs: [Value] = []
        for (i, p) in sig.params.enumerated() {
            switch p.type {
            case "string":
                // A string parameter becomes a UTF-8 (ptr, len) pair via hiko_alloc.
                let bytes = [UInt8](stringValue(args[i]).utf8)
                guard let alloc = instance.exports[function: "hiko_alloc"] else {
                    throw HikoError.missingExport("hiko_alloc")
                }
                let ptr = try alloc([.i32(UInt32(bytes.count))])[0].i32
                memory.withUnsafeMutableBufferPointer(offset: UInt(ptr), count: bytes.count) {
                    $0.copyBytes(from: bytes)
                }
                wasmArgs.append(.i32(ptr))
                wasmArgs.append(.i32(UInt32(bytes.count)))
            case "float":
                wasmArgs.append(.f64(try doubleValue(args[i], for: name).bitPattern))
            case "boolean":
                wasmArgs.append(.i32((args[i] as? Bool == true) ? 1 : 0))
            default: // int
                wasmArgs.append(.i32(UInt32(bitPattern: Int32(truncatingIfNeeded: try intValue(args[i], for: name)))))
            }
        }

        // Signatures with strings are called through the hiko_run_<name> wrapper export.
        let exportName = hasString ? "hiko_run_\(name)" : name
        guard let fn = instance.exports[function: exportName] else {
            throw HikoError.missingExport(exportName)
        }
        let raw = try fn(wasmArgs)[0]

        switch sig.returns {
        case "boolean":
            return raw.i32 != 0
        case "float":
            return Double(bitPattern: raw.f64)
        case "string":
            // String return = pointer to [u32 little-endian length][UTF-8 bytes].
            let data = memory.data
            let ptr = Int(raw.i32)
            let len = Int(data[ptr]) | Int(data[ptr + 1]) << 8 | Int(data[ptr + 2]) << 16 | Int(data[ptr + 3]) << 24
            return String(decoding: data[(ptr + 4)..<(ptr + 4 + len)], as: UTF8.self)
        default: // int
            return Int(Int32(bitPattern: raw.i32))
        }
    }

    private func stringValue(_ v: Any) -> String {
        v as? String ?? String(describing: v)
    }

    private func doubleValue(_ v: Any, for name: String) throws -> Double {
        if let d = v as? Double { return d }
        if let i = v as? Int { return Double(i) }
        throw HikoError.badArguments("\(name): expected a numeric argument, got \(v)")
    }

    private func intValue(_ v: Any, for name: String) throws -> Int {
        if let i = v as? Int { return i }
        throw HikoError.badArguments("\(name): expected an Int argument, got \(v)")
    }

    // ─── Integrity helpers (mirror of integrity-core.mjs, client side, CryptoKit) ──
    /// `ok`: bytes are consistent with the integrity statement. `verified`: an Ed25519
    /// signature was actually checked (Layer 2).
    static func verifyIntegrity(_ integrityJSON: Data, wasmBytes: Data, manifestBytes: Data) -> VerifyResult {
        guard let i = try? JSONSerialization.jsonObject(with: integrityJSON) as? [String: Any],
              (i["integrity"] as? Int) == 1 else {
            return VerifyResult(ok: false, verified: false, reason: "integrity.json missing or unsupported version")
        }
        guard let wasmSha = (i["wasm"] as? [String: Any])?["sha256"] as? String,
              let manifestSha = (i["manifest"] as? [String: Any])?["sha256"] as? String else {
            return VerifyResult(ok: false, verified: false, reason: "integrity.json malformed")
        }
        if sha256Hex(wasmBytes) != wasmSha {
            return VerifyResult(ok: false, verified: false, reason: "wasm.sha256 mismatch (transit corruption / tamper)")
        }
        if sha256Hex(manifestBytes) != manifestSha {
            return VerifyResult(ok: false, verified: false, reason: "manifest.sha256 mismatch (tamper)")
        }
        guard let signature = i["signature"] as? String, !signature.isEmpty else {
            return VerifyResult(ok: true, verified: false, reason: "unsigned (Layer 1 — hash only)")
        }
        let keyId = i["keyId"] as? String ?? ""
        guard let pem = publicSigningKeys[keyId] else {
            return VerifyResult(ok: false, verified: false, reason: "unknown/invalid keyId: \(keyId)")
        }
        let message = buildSigningMessage(
            bundleVersion: i["bundleVersion"] as? String ?? "", wasmSha: wasmSha, manifestSha: manifestSha, keyId: keyId,
        )
        if !verifyEd25519(message: message, signatureB64: signature, spkiPEM: pem) {
            return VerifyResult(ok: false, verified: false, reason: "Ed25519 signature failed (forged source?)")
        }
        return VerifyResult(ok: true, verified: true, reason: nil)
    }

    private static func sha256Hex(_ data: Data) -> String {
        SHA256.hash(data: data).map { String(format: "%02x", $0) }.joined()
    }

    private static func buildSigningMessage(bundleVersion: String, wasmSha: String, manifestSha: String, keyId: String) -> String {
        [
            signMessagePrefix,
            "bundleVersion=\(bundleVersion)",
            "wasm.sha256=\(wasmSha)",
            "manifest.sha256=\(manifestSha)",
            "keyId=\(keyId)",
        ].joined(separator: "\n")
    }

    private static func verifyEd25519(message: String, signatureB64: String, spkiPEM: String) -> Bool {
        let b64 = spkiPEM
            .replacingOccurrences(of: "-----BEGIN PUBLIC KEY-----", with: "")
            .replacingOccurrences(of: "-----END PUBLIC KEY-----", with: "")
            .components(separatedBy: .whitespacesAndNewlines).joined()
        // Ed25519 SPKI DER = 12-byte header + 32-byte raw key. CryptoKit wants the raw key.
        guard let der = Data(base64Encoded: b64), der.count == 44,
              let sig = Data(base64Encoded: signatureB64),
              let key = try? Curve25519.Signing.PublicKey(rawRepresentation: der.suffix(32)) else {
            return false
        }
        return key.isValidSignature(sig, for: Data(message.utf8))
    }
}