// Copyright 2026 Everpaper. Licensed under the Apache License, Version 2.0 (see ../../../LICENSE).
//
// eg2-cli — command-line front end for the EmbeddingGemma2CoreAI library.
//
//   eg2-cli embed       <bundle> [--compute ane|gpu|cpu] [--kind document|search_query|...] [--dim 768] "text" ...
//   eg2-cli embed-image <bundle> [--compute ...] [--image-compute ane|gpu|cpu] [--exif] [--dim 768] image ...
//                       (one JSON line per image)
//   eg2-cli embed-audio <bundle> [--compute ...] [--audio-compute ane|gpu|cpu] [--long truncate|chunk_mean] [--dim 768] clip ...
//                       (one JSON line per clip)
//   eg2-cli smoke       <bundle> <reference.json> [<reference.json> ...] [--compute ...] [--image-compute ...] [--audio-compute ...] [--cleanup-temp]
//                       (text: tests/smoke-reference.json; images: tests/image-smoke-reference.json; audio: tests/audio-smoke-reference.json)
//   eg2-cli tokcheck    <bundle> <tokenizer-cases.json>
//   eg2-cli pixcheck    <preprocess-cases.json>       (decode + resize vs the Python reference, no model needed)
//   eg2-cli audiocheck  <audio-preprocess-cases.json> (decode + resample + log-mel vs the Python reference, no model needed)
//   eg2-cli dump        <image> <out-prefix> [--exif] (writes <prefix>.decoded.rgb / .resized.rgb / .json)
//   eg2-cli dump-audio  <clip> <out-prefix>          (writes <prefix>.wave16k.f32 / .features.f32 / .json)
//   eg2-cli latency     <bundle> [--compute ...] [--image-compute ...] [--audio-compute ...] [--trials 50] [--text-only]
//   eg2-cli prewarm     <bundle> [--compute ...] [--image-compute ...] [--audio-compute ...] [--text-only]
//   eg2-cli cleanup-temp          (removes this process's MPSGraph temp folders; mostly useful after the others)
//
// `embed` and `tokcheck` load the text asset only; `embed-image` also loads the vision asset and `embed-audio` the
// audio asset; `smoke` loads what its reference files need; `latency` and `prewarm` load every asset the bundle has
// unless --text-only. Every command accepts --cleanup-temp to delete the MPSGraph temp folders this process created,
// after the model is released (see EmbeddingGemma2.removeOwnMPSGraphTempFiles). `--compute` places the text asset
// (and the image and audio encoders unless `--image-compute` / `--audio-compute` are given).

import CryptoKit
import EmbeddingGemma2CoreAI
import Foundation

var args = Array(CommandLine.arguments.dropFirst())

@MainActor func usage() -> Never {
    print("""
    usage: eg2-cli <embed|embed-image|embed-audio|smoke|tokcheck|pixcheck|audiocheck|dump|dump-audio|latency|prewarm|cleanup-temp> ...
                   [--compute ane|gpu|cpu] [--image-compute ane|gpu|cpu] [--audio-compute ane|gpu|cpu] [--cleanup-temp]
    """)
    exit(2)
}

@MainActor func takeOption(_ name: String) -> String? {
    guard let i = args.firstIndex(of: name), i + 1 < args.count else { return nil }
    let v = args[i + 1]
    args.removeSubrange(i...(i + 1))
    return v
}

@MainActor func takeFlag(_ name: String) -> Bool {
    guard let i = args.firstIndex(of: name) else { return false }
    args.remove(at: i)
    return true
}

func fmt(_ x: Double, _ digits: Int = 2) -> String { String(format: "%.\(digits)f", x) }

func sha256(_ bytes: [UInt8]) -> String { SHA256.hash(data: Data(bytes)).map { String(format: "%02x", $0) }.joined() }

func cpuSeconds() -> Double {
    var u = rusage()
    getrusage(RUSAGE_SELF, &u)
    return Double(u.ru_utime.tv_sec) + Double(u.ru_utime.tv_usec) / 1e6 + Double(u.ru_stime.tv_sec) + Double(u.ru_stime.tv_usec) / 1e6
}

/// Peak physical footprint (MB) and the Neural Engine footprint (MB) of this process (proc_pid_rusage, RUSAGE_INFO_V6).
func footprints() -> (physMB: Double, neuralMB: Double) {
    var r = rusage_info_v6()
    let kr = withUnsafeMutablePointer(to: &r) { p -> Int32 in
        p.withMemoryRebound(to: rusage_info_t?.self, capacity: 1) { proc_pid_rusage(getpid(), RUSAGE_INFO_V6, $0) }
    }
    guard kr == 0 else { return (-1, -1) }
    return (Double(r.ri_lifetime_max_phys_footprint) / 1048576, Double(r.ri_neural_footprint) / 1048576)
}

func percentile(_ xs: [Double], _ p: Double) -> Double {
    let s = xs.sorted()
    return s[min(s.count - 1, Int(Double(s.count) * p))]
}

func floatsData(_ v: [Float]) -> Data { v.withUnsafeBufferPointer { Data(buffer: $0) } }

func readFloats(_ url: URL) throws -> [Float] {
    let d = try Data(contentsOf: url)
    return d.withUnsafeBytes { Array($0.bindMemory(to: Float.self)) }
}

guard !args.isEmpty else { usage() }
let cmd = args.removeFirst()
let cleanup = takeFlag("--cleanup-temp")
@MainActor func computeOption(_ name: String) -> ComputePreference? {
    guard let raw = takeOption(name) else { return nil }
    guard let c = ComputePreference(rawValue: raw) else {
        print("\(name) must be ane, gpu, cpu or default")
        exit(2)
    }
    return c
}
let compute = computeOption("--compute") ?? .neuralEngine
let imageCompute = computeOption("--image-compute")   // nil = same as --compute
let audioCompute = computeOption("--audio-compute")   // nil = same as --compute
let exif = takeFlag("--exif")
let textOnly = takeFlag("--text-only")
let longPolicyRaw = takeOption("--long") ?? "truncate"
guard let longPolicy = LongAudioPolicy(rawValue: longPolicyRaw) else {
    print("--long must be truncate or chunk_mean")
    exit(2)
}

@MainActor func finish(_ code: Int32) -> Never {
    if cleanup {
        let r = EmbeddingGemma2.removeOwnMPSGraphTempFiles()
        print("cleanup-temp: removed \(r.folders) MPSGraph temp folder(s), \(fmt(Double(r.bytes) / 1048576, 1)) MB")
    }
    exit(code)
}

if cmd == "cleanup-temp" { finish(0) }

struct PixCase: Decodable {
    let file: String; let width: Int; let height: Int; let resized_width: Int; let resized_height: Int; let n_soft: Int
    let decoded_sha256: String; let resized_sha256: String; let lossy: Bool
}

/// One entry of tests/audio-preprocess-cases.json: the Python front end's result on a shipped clip, with the
/// 16 kHz waveform and the log-mel features stored next to it as raw float32 files.
struct AudioCase: Decodable {
    let file: String; let sample_rate: Int; let channels: Int; let n_samples16k: Int; let n_frames: Int; let n_tokens: Int
    let bucket: Int; let wave_file: String; let features_file: String
    let wave_tolerance: Float; let feature_tolerance: Float; let note: String?
}

// Smoke reference files: text (v1.0 format, "cases"), images ("images" + "texts") or audio ("audio" + "texts").
struct TextCase: Decodable { let text: String; let kind: String; let embedding: [Float] }
struct ImageCase: Decodable { let file: String; let query: String; let embedding: [Float] }
struct AudioRefCase: Decodable { let file: String; let query: String?; let embedding: [Float]; let threshold_cosine: Double?; let fp32_rank1: Bool? }
struct RefFile: Decodable {
    let threshold_cosine: Double
    let threshold_cosine_image: Double?
    let threshold_cosine_audio: Double?
    let cases: [TextCase]?
    let images: [ImageCase]?
    let audio: [AudioRefCase]?
    let texts: [TextCase]?
}

do {
    switch cmd {
    case "pixcheck":
        guard let casesPath = args.first else { usage() }
        let url = URL(fileURLWithPath: casesPath)
        let cases = try JSONDecoder().decode([PixCase].self, from: Data(contentsOf: url))
        let base = url.deletingLastPathComponent()
        var exact = 0, lossyDiff = 0, fail = 0
        let t0 = Date()
        for c in cases {
            let img = try ImageDecoder.decodeRGB(url: base.appendingPathComponent(c.file), applyEXIFOrientation: false)
            let r = try ImagePreprocessor.resized(img)
            let inputs = ImagePreprocessor.patchify(resizedRGB: r.rgb, width: r.width, height: r.height)
            let dims = img.width == c.width && img.height == c.height && r.width == c.resized_width && r.height == c.resized_height && inputs.softTokenCount == c.n_soft
            let dOK = sha256(img.rgb) == c.decoded_sha256
            let rOK = sha256(r.rgb) == c.resized_sha256
            if dims && dOK && rOK {
                exact += 1
                print("EXACT    \(c.file) \(img.width)x\(img.height) -> \(r.width)x\(r.height), n=\(inputs.softTokenCount)")
            } else if dims && c.lossy && !dOK {
                lossyDiff += 1
                print("DIFFERS  \(c.file) (lossy format: decoder differences are expected; resized \(rOK ? "identical" : "differs"))")
            } else {
                fail += 1
                print("FAIL     \(c.file) dims \(dims) decoded \(dOK) resized \(rOK) (\(img.width)x\(img.height) -> \(r.width)x\(r.height), n=\(inputs.softTokenCount))")
            }
        }
        print("PIXCHECK \(exact)/\(cases.count) bit-identical to the Python reference, \(lossyDiff) lossy-format decoder differences, \(fail) failures, \(fmt(Date().timeIntervalSince(t0), 2)) s")
        finish(fail == 0 ? 0 : 1)

    case "audiocheck":
        guard let casesPath = args.first else { usage() }
        let url = URL(fileURLWithPath: casesPath)
        let cases = try JSONDecoder().decode([AudioCase].self, from: Data(contentsOf: url))
        let base = url.deletingLastPathComponent()
        var pass = 0, fail = 0, waveExact = 0
        var worstWave: Float = 0, worstFeat: Float = 0
        let t0 = Date()
        for c in cases {
            let t1 = Date()
            let d = try AudioDecoder.decode(url: base.appendingPathComponent(c.file))
            let wav = AudioResampler.resampleTo16k(AudioFrontend.toMono(d), rate: d.sampleRate)
            let inputs = try AudioFrontend.prepare(wav16k: wav)
            let ms = Date().timeIntervalSince(t1) * 1000
            let refWave = try readFloats(base.appendingPathComponent(c.wave_file))
            let refFeat = try readFloats(base.appendingPathComponent(c.features_file))
            let counts = d.sampleRate == c.sample_rate && d.channelCount == c.channels && wav.count == c.n_samples16k
                && inputs.frameCount == c.n_frames && inputs.softTokenCount == c.n_tokens && inputs.bucket == c.bucket
            var wDiff: Float = 0
            if wav.count == refWave.count { for i in 0..<wav.count { wDiff = max(wDiff, abs(wav[i] - refWave[i])) } } else { wDiff = .infinity }
            var fDiff: Float = 0
            let n = inputs.frameCount * AudioInputs.nMels
            if refFeat.count == n {
                let (feats, _) = AudioFrontend.logMel(wav)
                for i in 0..<n { fDiff = max(fDiff, abs(feats[i] - refFeat[i])) }
            } else { fDiff = .infinity }
            let ok = counts && wDiff <= c.wave_tolerance && fDiff <= c.feature_tolerance
            if ok { pass += 1 } else { fail += 1 }
            if wDiff == 0 { waveExact += 1 }
            worstWave = max(worstWave, wDiff == .infinity ? worstWave : wDiff)
            worstFeat = max(worstFeat, fDiff == .infinity ? worstFeat : fDiff)
            print("\(ok ? "PASS" : "FAIL")  \(c.file)  \(d.sampleRate) Hz x\(d.channelCount) -> \(wav.count) samples, \(inputs.frameCount) frames, \(inputs.softTokenCount) tokens (bucket \(inputs.bucket)); "
                  + "wave max diff \(wDiff == 0 ? "0 (exact)" : String(format: "%.2e", wDiff)) (tol \(String(format: "%.0e", c.wave_tolerance))), "
                  + "features max diff \(String(format: "%.2e", fDiff)) (tol \(String(format: "%.0e", c.feature_tolerance))), \(fmt(ms, 0)) ms\(counts ? "" : "  COUNTS DIFFER")")
        }
        print("AUDIOCHECK \(pass)/\(cases.count) within tolerance of the Python reference (\(waveExact) waveforms bit-identical; worst wave diff \(String(format: "%.2e", worstWave)), worst feature diff \(String(format: "%.2e", worstFeat))), \(fail) failures, \(fmt(Date().timeIntervalSince(t0), 2)) s")
        finish(fail == 0 ? 0 : 1)

    case "dump":
        guard args.count >= 2 else { usage() }
        let img = try ImageDecoder.decodeRGB(url: URL(fileURLWithPath: args[0]), applyEXIFOrientation: exif)
        let r = try ImagePreprocessor.resized(img)
        let inputs = ImagePreprocessor.patchify(resizedRGB: r.rgb, width: r.width, height: r.height)
        try Data(img.rgb).write(to: URL(fileURLWithPath: args[1] + ".decoded.rgb"))
        try Data(r.rgb).write(to: URL(fileURLWithPath: args[1] + ".resized.rgb"))
        let info: [String: Any] = ["width": img.width, "height": img.height, "resized_width": r.width, "resized_height": r.height,
                                   "n_soft": inputs.softTokenCount, "decoded_sha256": sha256(img.rgb), "resized_sha256": sha256(r.rgb)]
        try JSONSerialization.data(withJSONObject: info).write(to: URL(fileURLWithPath: args[1] + ".json"))
        print(String(decoding: try JSONSerialization.data(withJSONObject: info), as: UTF8.self))
        finish(0)

    case "dump-audio":
        guard args.count >= 2 else { usage() }
        let d = try AudioDecoder.decode(url: URL(fileURLWithPath: args[0]))
        let wav = AudioResampler.resampleTo16k(AudioFrontend.toMono(d), rate: d.sampleRate)
        let (feats, frames) = AudioFrontend.logMel(wav)
        try floatsData(wav).write(to: URL(fileURLWithPath: args[1] + ".wave16k.f32"))
        try floatsData(feats).write(to: URL(fileURLWithPath: args[1] + ".features.f32"))
        let info: [String: Any] = ["sample_rate": d.sampleRate, "channels": d.channelCount, "decoded_frames": d.frameCount,
                                   "n_samples16k": wav.count, "n_frames": frames, "n_tokens": AudioFrontend.softTokenCount(frames: frames)]
        try JSONSerialization.data(withJSONObject: info).write(to: URL(fileURLWithPath: args[1] + ".json"))
        print(String(decoding: try JSONSerialization.data(withJSONObject: info), as: UTF8.self))
        finish(0)

    default:
        break
    }

    guard let bundlePath = args.first else { usage() }
    args.removeFirst()
    let bundle = URL(fileURLWithPath: bundlePath)

    switch cmd {
    case "embed":
        let kind = takeOption("--kind") ?? "document"
        let dim = Int(takeOption("--dim") ?? "768") ?? 768
        let eg2 = try await EmbeddingGemma2(bundle: bundle, compute: compute)
        for text in args {
            let v = try await eg2.embed(text, kind: kind == "none" ? nil : kind, dim: dim)
            let json: [String: Any] = ["text": text, "kind": kind, "dim": dim, "embedding": v.map { Double($0) }]
            print(String(decoding: try JSONSerialization.data(withJSONObject: json), as: UTF8.self))
        }
        _ = consume eg2
        finish(0)

    case "embed-image":
        let dim = Int(takeOption("--dim") ?? "768") ?? 768
        let eg2 = try await EmbeddingGemma2(bundle: bundle, compute: compute, loadVision: true, imageCompute: imageCompute)
        for path in args {
            let t0 = Date()
            let inputs = try ImagePreprocessor.prepare(url: URL(fileURLWithPath: path), applyEXIFOrientation: exif)
            let t1 = Date()
            let v = try await eg2.embedImage(inputs, dim: dim)
            let json: [String: Any] = ["image": path, "dim": dim, "n_soft": inputs.softTokenCount, "resized": [inputs.resizedWidth, inputs.resizedHeight],
                                       "exif_orientation_applied": exif, "preprocess_ms": Int(t1.timeIntervalSince(t0) * 1000),
                                       "model_ms": Int(Date().timeIntervalSince(t1) * 1000), "embedding": v.map { Double($0) }]
            print(String(decoding: try JSONSerialization.data(withJSONObject: json), as: UTF8.self))
        }
        _ = consume eg2
        finish(0)

    case "embed-audio":
        let dim = Int(takeOption("--dim") ?? "768") ?? 768
        let eg2 = try await EmbeddingGemma2(bundle: bundle, compute: compute, loadAudio: true, audioCompute: audioCompute)
        for path in args {
            let t0 = Date()
            let d = try AudioDecoder.decode(url: URL(fileURLWithPath: path))
            let wav = AudioResampler.resampleTo16k(AudioFrontend.toMono(d), rate: d.sampleRate)
            let t1 = Date()
            let v = try await eg2.embedAudio(samples16k: wav, dim: dim, longAudio: longPolicy)
            let frames = AudioFrontend.validFrameCount(samples: wav.count)
            let json: [String: Any] = ["audio": path, "dim": dim, "sample_rate": d.sampleRate, "channels": d.channelCount,
                                       "seconds": Double(wav.count) / 16000.0, "n_frames": frames, "n_soft": AudioFrontend.softTokenCount(frames: frames),
                                       "truncated_to_30s": wav.count > AudioFrontend.maxSamples && longPolicy == .truncate,
                                       "long_audio": longPolicy.rawValue, "decode_ms": Int(t1.timeIntervalSince(t0) * 1000),
                                       "model_ms": Int(Date().timeIntervalSince(t1) * 1000), "embedding": v.map { Double($0) }]
            print(String(decoding: try JSONSerialization.data(withJSONObject: json), as: UTF8.self))
        }
        _ = consume eg2
        finish(0)

    case "smoke":
        guard !args.isEmpty else { usage() }
        let refs = try args.map { path -> (URL, RefFile) in
            let url = URL(fileURLWithPath: path)
            return (url, try JSONDecoder().decode(RefFile.self, from: Data(contentsOf: url)))
        }
        let needsVision = refs.contains { !($0.1.images ?? []).isEmpty }
        let needsAudio = refs.contains { !($0.1.audio ?? []).isEmpty }
        let eg2 = try await EmbeddingGemma2(bundle: bundle, compute: compute, loadVision: needsVision, imageCompute: imageCompute,
                                            loadAudio: needsAudio, audioCompute: audioCompute)
        print("loaded in \(fmt(eg2.loadSeconds, 1)) s (tokenizer \(fmt(eg2.tokenizerLoadSeconds, 2)) s), compute=\(compute.rawValue), "
              + "image-compute=\(eg2.hasVision ? eg2.imageCompute.rawValue : "-"), audio-compute=\(eg2.hasAudio ? eg2.audioCompute.rawValue : "-"), "
              + "buckets=\(eg2.buckets), vision \(eg2.hasVision ? "loaded" : "not loaded"), audio \(eg2.hasAudio ? "loaded" : "not loaded")")
        var ok = true
        for (refURL, ref) in refs {
            print("-- \(refURL.lastPathComponent)")
            // text cases (v1.0 format)
            for c in ref.cases ?? [] {
                let v = try await eg2.embed(c.text, kind: c.kind)
                let cos = cosineSimilarity(v, c.embedding)
                let pass = cos >= ref.threshold_cosine && v.allSatisfy { $0.isFinite }
                ok = ok && pass
                print("\(pass ? "PASS" : "FAIL") cosine \(fmt(cos, 6)) (threshold \(ref.threshold_cosine))  \(c.kind): \(c.text.prefix(60))")
            }
            if let first = ref.cases?.dropFirst().first {
                // Matryoshka prefixes must stay close to the truncated + re-normalized reference
                for dim in [512, 256, 128] {
                    let v = try await eg2.embed(first.text, kind: first.kind, dim: dim)
                    let cos = cosineSimilarity(v, Array(first.embedding.prefix(dim)))
                    let pass = cos >= ref.threshold_cosine
                    ok = ok && pass
                    print("\(pass ? "PASS" : "FAIL") cosine \(fmt(cos, 6)) at dim \(dim)")
                }
            }
            // images, their queries and documents
            let images = ref.images ?? []
            var imgVecs: [[Float]] = []
            let imgThreshold = ref.threshold_cosine_image ?? ref.threshold_cosine
            for c in images {
                let t0 = Date()
                let v = try await eg2.embedImage(url: refURL.deletingLastPathComponent().appendingPathComponent(c.file))
                let cos = cosineSimilarity(v, c.embedding)
                let pass = cos >= imgThreshold && v.allSatisfy { $0.isFinite }
                ok = ok && pass
                imgVecs.append(v)
                print("\(pass ? "PASS" : "FAIL") cosine \(fmt(cos, 6)) (threshold \(imgThreshold))  image \(c.file)  \(Int(Date().timeIntervalSince(t0) * 1000)) ms")
            }
            // audio clips (per-clip thresholds: resampled and lossy files are allowed more)
            let audio = ref.audio ?? []
            var audVecs: [[Float]] = []
            for c in audio {
                let t0 = Date()
                let v = try await eg2.embedAudio(url: refURL.deletingLastPathComponent().appendingPathComponent(c.file))
                let cos = cosineSimilarity(v, c.embedding)
                let thr = c.threshold_cosine ?? ref.threshold_cosine_audio ?? ref.threshold_cosine
                let pass = cos >= thr && v.allSatisfy { $0.isFinite }
                ok = ok && pass
                audVecs.append(v)
                print("\(pass ? "PASS" : "FAIL") cosine \(fmt(cos, 6)) (threshold \(thr))  audio \(c.file)  \(Int(Date().timeIntervalSince(t0) * 1000)) ms")
            }
            var queryVecs: [(String, [Float])] = []
            for c in ref.texts ?? [] {
                let v = try await eg2.embed(c.text, kind: c.kind)
                let cos = cosineSimilarity(v, c.embedding)
                let pass = cos >= ref.threshold_cosine && v.allSatisfy { $0.isFinite }
                ok = ok && pass
                if c.kind == "search_query" { queryVecs.append((c.text, v)) }
                print("\(pass ? "PASS" : "FAIL") cosine \(fmt(cos, 6)) (threshold \(ref.threshold_cosine))  \(c.kind): \(c.text.prefix(60))")
            }
            // text -> image: each query must rank its own image first
            if !images.isEmpty {
                var hits = 0
                for (q, qv) in queryVecs {
                    let scores = imgVecs.map { cosineSimilarity(qv, $0) }
                    let best = scores.indices.max(by: { scores[$0] < scores[$1] })!
                    let hit = images.firstIndex(where: { $0.query == q }) == best
                    hits += hit ? 1 : 0
                    print("\(hit ? "PASS" : "FAIL") query \"\(q)\" -> \(images[best].file) (\(fmt(scores[best], 3)))")
                }
                ok = ok && hits == queryVecs.count
                print("text -> image: \(hits)/\(queryVecs.count) queries rank their own image first")
            }
            // text -> audio: each query must rank its own clip first (among the clips that have a query), checked only
            // where the fp32 reference itself ranks it first (fp32_rank1; synthesized sounds are a hard retrieval case)
            if !audio.isEmpty {
                let withQuery = audio.indices.filter { audio[$0].query != nil }
                var hits = 0, total = 0
                for (q, qv) in queryVecs {
                    guard let want = audio.firstIndex(where: { $0.query == q }) else { continue }
                    let scores = withQuery.map { cosineSimilarity(qv, audVecs[$0]) }
                    let best = withQuery[scores.indices.max(by: { scores[$0] < scores[$1] })!]
                    let hit = best == want
                    if audio[want].fp32_rank1 == false {
                        print("SKIP query \"\(q)\" -> \(audio[best].file) (\(fmt(scores.max()!, 3))) (fp32 does not rank this clip first either)")
                        continue
                    }
                    total += 1
                    hits += hit ? 1 : 0
                    print("\(hit ? "PASS" : "FAIL") query \"\(q)\" -> \(audio[best].file) (\(fmt(scores.max()!, 3)))")
                }
                ok = ok && hits == total
                print("text -> audio: \(hits)/\(total) queries rank their own clip first (as fp32 does)")
            }
        }
        let fp = footprints()
        print("peak footprint \(fmt(fp.physMB, 0)) MB, neural footprint \(fmt(fp.neuralMB, 0)) MB")
        print(ok ? "SMOKE PASS" : "SMOKE FAIL")
        _ = consume eg2
        finish(ok ? 0 : 1)

    case "tokcheck":
        guard let casesPath = args.first else { usage() }
        struct Case: Decodable { let cat: String; let kind: String; let text: String; let ids: [Int32] }
        let cases = try JSONDecoder().decode([Case].self, from: Data(contentsOf: URL(fileURLWithPath: casesPath)))
        let config = try JSONDecoder().decode(BundleConfig.self, from: Data(contentsOf: bundle.appendingPathComponent("config.json")))
        let t0 = Date()
        let tok = try GemmaTokenizer(contentsOf: bundle.appendingPathComponent("tokenizer/tokenizer.json"))
        print("tokenizer loaded in \(fmt(Date().timeIntervalSince(t0), 2)) s, vocab \(tok.vocabSize)")
        let prefixes = ["q": "task: search result | query: ", "d": "title: none | text: ", "raw": ""]
        var same = 0, totalIDs = 0, shown = 0
        var perCat: [String: (Int, Int)] = [:]
        let t1 = Date()
        let cpu0 = cpuSeconds()
        for c in cases {
            let ids = tok.encodeWithSpecials(prefixes[c.kind]! + c.text)
            let ok = ids == c.ids
            totalIDs += c.ids.count
            var e = perCat[c.cat] ?? (0, 0)
            e.1 += 1
            if ok {
                same += 1; e.0 += 1
            } else if shown < 10 {
                shown += 1
                print("DIFFERENT [\(c.cat)] \(c.text.prefix(60).debugDescription): got \(ids.prefix(12)) expected \(c.ids.prefix(12))")
            }
            perCat[c.cat] = e
        }
        let secs = Date().timeIntervalSince(t1), cpu = cpuSeconds() - cpu0
        for (cat, e) in perCat.sorted(by: { $0.key < $1.key }) { print("  \(cat): \(e.0)/\(e.1)") }
        print("TOKCHECK \(same)/\(cases.count) texts identical to Hugging Face tokenizers (\(totalIDs) ids) in \(fmt(secs, 2)) s wall / \(fmt(cpu, 2)) s CPU")
        print("config bos/eos/pad \(config.bos_token_id)/\(config.eos_token_id)/\(config.pad_token_id) == tokenizer \(tok.bosID)/\(tok.eosID)/\(tok.padID)")
        finish(same == cases.count ? 0 : 1)

    case "latency":
        let trials = Int(takeOption("--trials") ?? "50") ?? 50
        let hasAudioAsset = (try? JSONDecoder().decode(BundleConfig.self, from: Data(contentsOf: bundle.appendingPathComponent("config.json"))))?.audio != nil
        let eg2 = try await EmbeddingGemma2(bundle: bundle, compute: compute, loadVision: !textOnly, imageCompute: imageCompute,
                                            loadAudio: !textOnly && hasAudioAsset, audioCompute: audioCompute)
        print("loaded in \(fmt(eg2.loadSeconds, 1)) s, compute=\(compute.rawValue), image-compute=\(eg2.hasVision ? eg2.imageCompute.rawValue : "-"), "
              + "audio-compute=\(eg2.hasAudio ? eg2.audioCompute.rawValue : "-"), vision \(eg2.hasVision ? "loaded" : "not loaded"), audio \(eg2.hasAudio ? "loaded" : "not loaded")")
        let warm = try await eg2.prewarm()
        print("prewarm \(fmt(warm, 2)) s")
        if eg2.hasVision {
            // a synthetic 1008x624 image (273 soft tokens, the most common shape)
            var px = [UInt8](repeating: 0, count: 1008 * 624 * 3)
            for i in px.indices { px[i] = UInt8((i * 7919) % 251) }
            let inputs = ImagePreprocessor.patchify(resizedRGB: px, width: 1008, height: 624)
            var tv: [Double] = [], tt: [Double] = [], te: [Double] = []
            let imageTrials = max(5, trials / 2)
            let cpu0 = cpuSeconds()
            for _ in 0..<imageTrials {
                let t0 = Date()
                let soft = try await eg2.softTokens(inputs)
                let t1 = Date()
                _ = try await eg2.embed(sequence: eg2.imageSequence(softTokens: soft))
                let t2 = Date()
                tv.append(t1.timeIntervalSince(t0) * 1000); tt.append(t2.timeIntervalSince(t1) * 1000); te.append(t2.timeIntervalSince(t0) * 1000)
            }
            let cpu = (cpuSeconds() - cpu0) / Double(imageTrials) * 1000
            print("image (273 soft tokens, \(imageTrials) trials): vision_2520 p50 \(fmt(percentile(tv, 0.5), 1)) ms (min \(fmt(tv.min()!, 1))), "
                  + "mm_288 p50 \(fmt(percentile(tt, 0.5), 1)) ms, end-to-end p50 \(fmt(percentile(te, 0.5), 1)) ms, CPU \(fmt(cpu, 1)) ms/image (excl. decode+resize)")
            var big = [UInt8](repeating: 0, count: 2560 * 1600 * 3)
            for i in big.indices { big[i] = UInt8((i * 31) % 253) }
            let tp0 = Date()
            for _ in 0..<5 { _ = try ImagePreprocessor.prepare(DecodedImage(width: 2560, height: 1600, rgb: big)) }
            print("preprocess (resize 2560x1600 -> 1008x624 + patches): \(fmt(Date().timeIntervalSince(tp0) / 5 * 1000, 1)) ms")
        }
        if eg2.hasAudio {
            // synthetic clips of 3 s, 10 s and 30 s (a sum of tones and noise), per frame bucket
            let audioTrials = max(5, trials / 2)
            print("audio   seconds  bucket  tokens  audio_<Tb> p50 ms   mm_<S> p50 ms   end-to-end p50 ms   CPU ms/clip   front end ms")
            for seconds in [3.0, 10.0, 30.0] {
                let n = Int(seconds * 16000)
                var wav = [Float](repeating: 0, count: n)
                for i in 0..<n {
                    let t = Float(i) / 16000
                    wav[i] = 0.3 * sin(2 * .pi * 220 * t) + 0.2 * sin(2 * .pi * 1375 * t + 0.5 * sin(2 * .pi * 3 * t)) + 0.05 * Float((i * 7919) % 1000) / 1000 - 0.025
                }
                let tf0 = Date()
                let inputs = try AudioFrontend.prepare(wav16k: wav)
                let frontMs = Date().timeIntervalSince(tf0) * 1000
                var ta: [Double] = [], tt: [Double] = [], te: [Double] = []
                for _ in 0..<3 { _ = try await eg2.embedAudio(inputs) }
                let cpu0 = cpuSeconds()
                for _ in 0..<audioTrials {
                    let t0 = Date()
                    let soft = try await eg2.audioSoftTokens(inputs)
                    let t1 = Date()
                    _ = try await eg2.embed(sequence: eg2.audioSequence(softTokens: soft))
                    let t2 = Date()
                    ta.append(t1.timeIntervalSince(t0) * 1000); tt.append(t2.timeIntervalSince(t1) * 1000); te.append(t2.timeIntervalSince(t0) * 1000)
                }
                let cpu = (cpuSeconds() - cpu0) / Double(audioTrials) * 1000
                print(String(format: "        %7.1f  %6d  %6d  %17.1f  %14.1f  %18.1f  %12.1f  %12.1f", seconds, inputs.bucket, inputs.softTokenCount,
                             percentile(ta, 0.5), percentile(tt, 0.5), percentile(te, 0.5), cpu, frontMs))
            }
            // the front end on a 48 kHz stereo 30 s clip (decode excluded): resample + log-mel
            let n48 = 30 * 48000
            var l = [Float](repeating: 0, count: n48), r = [Float](repeating: 0, count: n48)
            for i in 0..<n48 { l[i] = sin(Float(i) * 0.01); r[i] = sin(Float(i) * 0.013) }
            let tp0 = Date()
            let mono = AudioFrontend.toMono(DecodedAudio(sampleRate: 48000, channelCount: 2, channels: [l, r]))
            let w16 = AudioResampler.resampleTo16k(mono, rate: 48000)
            let tp1 = Date()
            _ = try AudioFrontend.prepare(wav16k: w16)
            print("audio front end, 30 s stereo 48 kHz: mono + resample \(fmt(tp1.timeIntervalSince(tp0) * 1000, 0)) ms, log-mel + inputs \(fmt(Date().timeIntervalSince(tp1) * 1000, 0)) ms")
        }
        print("bucket   p50 ms   min ms   p90 ms   CPU ms/call   tok/s")
        for S in eg2.buckets {
            var ids = [Int32](repeating: 0, count: S)
            ids[0] = eg2.config.bos_token_id; ids[S - 1] = eg2.config.eos_token_id
            for i in 1..<(S - 1) { ids[i] = Int32(1000 + (i * 7919) % 200000) }
            for _ in 0..<3 { _ = try await eg2.embed(ids: ids, bucket: S) }
            var ts: [Double] = []
            let cpu0 = cpuSeconds()
            for _ in 0..<trials {
                let t = Date()
                _ = try await eg2.embed(ids: ids, bucket: S)
                ts.append(Date().timeIntervalSince(t) * 1000)
            }
            let cpu = (cpuSeconds() - cpu0) / Double(trials) * 1000
            let p50 = percentile(ts, 0.5)
            print(String(format: "%5d  %7.2f  %7.2f  %7.2f  %11.2f  %6.0f", S, p50, ts.min()!, percentile(ts, 0.9), cpu, Double(S) / p50 * 1000))
        }
        let fp = footprints()
        print("peak footprint \(fmt(fp.physMB, 0)) MB, neural footprint \(fmt(fp.neuralMB, 0)) MB (\(fp.neuralMB > 0 ? "Neural Engine in use" : "Neural Engine NOT used"))")
        _ = consume eg2
        finish(0)

    case "prewarm":
        let t0 = Date()
        let hasAudioAsset = (try? JSONDecoder().decode(BundleConfig.self, from: Data(contentsOf: bundle.appendingPathComponent("config.json"))))?.audio != nil
        let eg2 = try await EmbeddingGemma2(bundle: bundle, compute: compute, loadVision: !textOnly, imageCompute: imageCompute,
                                            loadAudio: !textOnly && hasAudioAsset, audioCompute: audioCompute)
        let warm = try await eg2.prewarm()
        print("load \(fmt(eg2.loadSeconds, 1)) s + prewarm \(fmt(warm, 2)) s = ready in \(fmt(Date().timeIntervalSince(t0), 1)) s "
              + "(vision \(eg2.hasVision ? "loaded" : "not loaded"), audio \(eg2.hasAudio ? "loaded" : "not loaded"))")
        _ = consume eg2
        finish(0)

    default:
        usage()
    }
} catch {
    print("error: \(error)")
    finish(1)
}