// Transcribe LibriSpeech clips with Apple's Speech framework and write the raw
// hypotheses to JSON, so they can be diffed against the macOS 26.5.1 transcripts
// published at https://lyonesse.app/data/speech-benchmark/
//
//   swiftc -O run.swift -o run
//   ./run manifest-test-clean.json 300 out.json            # SpeechTranscriber
//   ./run manifest-test-clean.json 300 out-dict.json dictation
//
// No Xcode project, no test host, no speech permission: SpeechTranscriber and
// DictationTranscriber do not require SFSpeechRecognizer authorization.
import Foundation
import AVFoundation
import Speech

struct Item: Codable { let id: String; let audio: String; let reference: String }
struct Manifest: Codable { let items: [Item] }
struct Row: Codable { let id: String; let hypothesis: String; let seconds: Double }
struct Output: Codable {
    let module: String
    let os: String
    let chip: String
    let memoryGB: Int
    let date: String
    let audioSeconds: Double
    let computeSeconds: Double
    let rows: [Row]
}

func sysctl(_ name: String) -> String {
    var size = 0
    sysctlbyname(name, nil, &size, nil, 0)
    var value = [CChar](repeating: 0, count: size)
    sysctlbyname(name, &value, &size, nil, 0)
    return String(cString: value)
}

let args = CommandLine.arguments
guard args.count >= 4 else {
    print("usage: run <manifest.json> <limit> <out.json> [dictation]")
    exit(2)
}
let manifest = try JSONDecoder().decode(Manifest.self, from: Data(contentsOf: URL(fileURLWithPath: args[1])))
let limit = Int(args[2]) ?? 300
let useDictation = args.count > 4 && args[4].lowercased().hasPrefix("dict")
let locale = Locale(identifier: "en_US")

// The exact configuration the published benchmark used: no transcription options,
// no reporting options, no attributes. Changing any of these changes the output,
// so leave them alone if the numbers are meant to be comparable.
func makeModule() -> any SpeechModule {
    if useDictation {
        return DictationTranscriber(locale: locale, preset: .longDictation)
    }
    return SpeechTranscriber(locale: locale, transcriptionOptions: [], reportingOptions: [], attributeOptions: [])
}

func text(of module: any SpeechModule) async throws -> String {
    var parts: [String] = []
    if let t = module as? SpeechTranscriber {
        for try await r in t.results where r.isFinal {
            parts.append(String(r.text.characters).trimmingCharacters(in: .whitespacesAndNewlines))
        }
    } else if let d = module as? DictationTranscriber {
        for try await r in d.results where r.isFinal {
            parts.append(String(r.text.characters).trimmingCharacters(in: .whitespacesAndNewlines))
        }
    }
    return parts.filter { !$0.isEmpty }.joined(separator: " ")
}

func transcribe(_ url: URL) async throws -> (String, Double) {
    let module = makeModule()
    let file = try AVAudioFile(forReading: url)
    let seconds = Double(file.length) / file.processingFormat.sampleRate
    let analyzer = try await SpeechAnalyzer(inputAudioFile: file, modules: [module], finishAfterFile: true)
    _ = analyzer
    return (try await text(of: module), seconds)
}

let probe = makeModule()
print("module:", useDictation ? "DictationTranscriber(.longDictation)" : "SpeechTranscriber")
print("asset status:", await AssetInventory.status(forModules: [probe]))
if let request = try await AssetInventory.assetInstallationRequest(supporting: [probe]) {
    print("installing model assets…")
    try await request.downloadAndInstall()
    print("installed")
}

let chip = sysctl("machdep.cpu.brand_string")
let memGB = Int(ProcessInfo.processInfo.physicalMemory / 1_073_741_824)
print("machine:", chip, "\(memGB)GB,", ProcessInfo.processInfo.operatingSystemVersionString)

var rows: [Row] = []
var audioTotal = 0.0
let started = Date()
for (index, item) in manifest.items.prefix(limit).enumerated() {
    let t0 = Date()
    do {
        let (hypothesis, seconds) = try await transcribe(URL(fileURLWithPath: item.audio))
        audioTotal += seconds
        rows.append(Row(id: item.id, hypothesis: hypothesis, seconds: Date().timeIntervalSince(t0)))
    } catch {
        // Record the failure rather than aborting: one bad clip should not cost the run.
        print("FAILED \(item.id): \(error)")
        rows.append(Row(id: item.id, hypothesis: "", seconds: Date().timeIntervalSince(t0)))
    }
    if (index + 1) % 50 == 0 { print("  \(index + 1)/\(min(limit, manifest.items.count))") }
}

let out = Output(
    module: useDictation ? "DictationTranscriber(.longDictation)" : "SpeechTranscriber",
    os: ProcessInfo.processInfo.operatingSystemVersionString,
    chip: chip,
    memoryGB: memGB,
    date: ISO8601DateFormatter().string(from: Date()),
    audioSeconds: audioTotal,
    computeSeconds: Date().timeIntervalSince(started),
    rows: rows)
let encoder = JSONEncoder()
encoder.outputFormatting = [.prettyPrinted, .sortedKeys]
try encoder.encode(out).write(to: URL(fileURLWithPath: args[3]))
let empties = rows.filter { $0.hypothesis.isEmpty }.count
print("wrote \(args[3]): \(rows.count) rows, \(empties) empty, RTF \(String(format: "%.4f", out.computeSeconds / max(out.audioSeconds, 1)))")