import ASRKit
import Foundation
import SpeechCore

// SpeechKit regression harness (macOS).
// Usage:
//   asrkit-cli <iphone-asr-dir> [--stress]
//   asrkit-cli <iphone-asr-dir> --file path/to/audio.wav
//   asrkit-cli <iphone-asr-dir> --file path/to/audio.wav --repeat 10
// Verifies each stage against Python references, exercising the SAME
// public API integrators use (ASRTranscriber + AssetBundle).

setvbuf(stdout, nil, _IONBF, 0)

let args = Array(CommandLine.arguments.dropFirst())
func value(after flag: String) -> String? {
    guard let i = args.firstIndex(of: flag), i + 1 < args.count else { return nil }
    return args[i + 1]
}
let fileArg = value(after: "--file")
let repeatCount = value(after: "--repeat").flatMap(Int.init) ?? 1
let rootArg = args.first { arg in
    !arg.hasPrefix("--") && arg != fileArg && arg != value(after: "--repeat")
}
let root = URL(fileURLWithPath: rootArg ?? FileManager.default.currentDirectoryPath)
let refs = root.appendingPathComponent("ios_refs")
let bundleURL = root.appendingPathComponent("dist/ASRModels.bundle")

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

do {
    if let fileArg {
        print("loading transcriber ...")
        let transcriber = try ASRTranscriber(bundleURL: bundleURL)
        transcriber.warmUp()
        let url = URL(fileURLWithPath: fileArg)
        let samples = try AudioResampler.loadFile(url)
        func footprintMB() -> Double {
            var info = task_vm_info_data_t()
            var count = mach_msg_type_number_t(
                MemoryLayout<task_vm_info_data_t>.size / MemoryLayout<natural_t>.size)
            _ = withUnsafeMutablePointer(to: &info) {
                $0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) {
                    task_info(mach_task_self_, task_flavor_t(TASK_VM_INFO), $0, &count)
                }
            }
            return Double(info.phys_footprint) / 1_048_576
        }
        print(String(format: "[file] start footprint %.0f MB", footprintMB()))
        var result: ASRTranscriber.Result?
        let passes = max(repeatCount, 1)
        for i in 1...passes {
            let r = try transcriber.transcribe(samples)
            result = r
            if passes > 1 {
                print(String(format: "[file] pass %d/%d footprint %.0f MB transcript: %@",
                             i, passes, footprintMB(), r.text))
            }
        }
        guard let result else { exit(1) }
        let t = result.timings
        print("[file] transcript: \(result.text)")
        print("[file] tokenIDs: \(result.tokenIDs)")
        print("[file] hitStop: \(result.hitStop)")
        print(String(format: "[file] %.0fms: mel %.0f + tower %.0f + decode %.0f",
                     t.total * 1000, t.mel * 1000, t.tower * 1000, t.decode * 1000))
        exit(0)
    }

    let refAudio = try loadFloats(refs.appendingPathComponent("ref_audio.bin"))
    let refMel = try loadFloats(refs.appendingPathComponent("ref_mel.bin"))
    let refProjected = try loadFloats(refs.appendingPathComponent("ref_projected.bin"))
    let refTranscript = try String(contentsOf: refs.appendingPathComponent("ref_transcript.txt"))
        .trimmingCharacters(in: .whitespacesAndNewlines)

    print("loading transcriber (verifyAssets: true) ...")
    var options = ASRTranscriber.Options()
    options.verifyAssets = true
    let transcriber = try ASRTranscriber(bundleURL: bundleURL, options: options)
    transcriber.warmUp()
    print("[assets] bundle verified  PASS")

    // stage 1: mel parity
    let (melOut, encLen, bucketFrames) = transcriber.mel.extract(refAudio)
    var melDiff: Float = 0
    for t in 0..<bucketFrames {
        for m in 0..<128 {
            melDiff = max(melDiff, abs(melOut[m * bucketFrames + t] - refMel[m * 3000 + t]))
        }
    }
    print(String(format: "[mel] max_abs_diff=%.6f encLen=%d bucket=%d  %@",
                 melDiff, encLen, bucketFrames, melDiff < 1e-3 ? "PASS" : "FAIL"))

    // stage 2: tower parity
    let masks = try transcriber.decoder.generateMasks(encLen: encLen)
    let audioEmbeds = try transcriber.tower.embed(
        mel: melOut, bucketFrames: bucketFrames, attnMask390: masks.attn,
        validMask390: masks.valid, sampleCount: refAudio.count)
    let flat = audioEmbeds.flatMap { $0 }
    var dot = 0.0, na = 0.0, nb = 0.0
    for (x, y) in zip(flat, refProjected) {
        dot += Double(x) * Double(y)
        na += Double(x) * Double(x)
        nb += Double(y) * Double(y)
    }
    let cos = dot / (na.squareRoot() * nb.squareRoot() + 1e-12)
    print(String(format: "[tower] tokens=%d cosine=%.6f  %@",
                 audioEmbeds.count, cos, cos > 0.995 ? "PASS" : "FAIL"))

    // stage 3: e2e via public API
    let result = try transcriber.transcribe(refAudio)
    let t = result.timings
    print("[e2e] transcript: \(result.text)")
    print("[e2e] tokenIDs: \(result.tokenIDs)")
    print("[e2e] reference:  \(refTranscript)")
    print(String(format: "[e2e] %@  (%.0fms: mel %.0f + tower %.0f + decode %.0f)",
                 result.text == refTranscript ? "PASS" : "FAIL",
                 t.total * 1000, t.mel * 1000, t.tower * 1000, t.decode * 1000))

    // error contract checks
    do {
        _ = try transcriber.transcribe([Float](repeating: 0, count: 100))
        print("[errors] audioTooShort  FAIL (no throw)")
    } catch let e as SpeechError {
        print("[errors] audioTooShort -> \(e)  PASS")
    }
    let token = ASRTranscriber.CancellationToken()
    token.cancel()
    do {
        _ = try transcriber.transcribe(refAudio, cancellation: token)
        print("[errors] cancellation  FAIL (no throw)")
    } catch SpeechError.cancelled {
        print("[errors] cancellation -> cancelled  PASS")
    }

    if CommandLine.arguments.contains("--stress") {
        func footprintMB() -> Double {
            var info = task_vm_info_data_t()
            var count = mach_msg_type_number_t(
                MemoryLayout<task_vm_info_data_t>.size / MemoryLayout<natural_t>.size)
            _ = withUnsafeMutablePointer(to: &info) {
                $0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) {
                    task_info(mach_task_self_, task_flavor_t(TASK_VM_INFO), $0, &count)
                }
            }
            return Double(info.phys_footprint) / 1_048_576
        }
        print(String(format: "[stress] start footprint %.0f MB", footprintMB()))
        for i in 1...100 {
            _ = try transcriber.transcribe(refAudio)
            if i % 20 == 0 {
                print(String(format: "[stress] pass %3d footprint %.0f MB", i, footprintMB()))
            }
        }
        print(String(format: "[stress] end footprint %.0f MB", footprintMB()))
    }
} catch {
    print("ERROR: \(error)")
    exit(1)
}