// Transcribe a WAV file with macOS 26's SpeechAnalyzer.
//
// The accurate recogniser on this OS — SpeechTranscriber, driven by
// SpeechAnalyzer — is a Swift-only API built on actors and AsyncSequence, so
// pyobjc cannot reach it. This is the smallest thing that can: read a file,
// print JSON, exit. bot.py runs it as a subprocess.
//
// Why bother, versus the SFSpeechRecognizer path in apple_stt.py:
//   - Roughly four times more accurate on published benchmarks.
//   - No session length limit, which is what forces the transcript-stitching
//     workaround in apple_stt.py.
//   - Still supports vocabulary biasing, via AnalysisContext on the analyzer
//     rather than on the request.
//   - Returns ranked alternatives, which measurably help LLM post-correction.
//
// Usage:
//   speech-helper <file.wav> [--locale en-US] [--terms a,b,c] [--alternatives]
//   speech-helper --check [--locale en-US]

import AVFoundation
import Foundation
import Speech

struct Options {
    var path: String?
    var locale = "en-US"
    var terms: [String] = []
    var alternatives = false
    var check = false
    var dictation = false
}

func parseArguments() -> Options {
    var options = Options()
    var arguments = Array(CommandLine.arguments.dropFirst())
    while let argument = arguments.first {
        arguments.removeFirst()
        switch argument {
        case "--locale":
            options.locale = arguments.isEmpty ? options.locale : arguments.removeFirst()
        case "--terms":
            let raw = arguments.isEmpty ? "" : arguments.removeFirst()
            options.terms = raw.split(separator: ",").map {
                $0.trimmingCharacters(in: .whitespaces)
            }.filter { !$0.isEmpty }
        case "--alternatives":
            options.alternatives = true
        case "--dictation":
            options.dictation = true
        case "--check":
            options.check = true
        default:
            if options.path == nil { options.path = argument }
        }
    }
    return options
}

func emit(_ payload: [String: Any]) {
    let data = try! JSONSerialization.data(withJSONObject: payload, options: [.sortedKeys])
    FileHandle.standardOutput.write(data)
    FileHandle.standardOutput.write("\n".data(using: .utf8)!)
}

func fail(_ message: String) -> Never {
    emit(["error": message])
    exit(1)
}

/// Download the on-device model if this locale hasn't been used before.
///
/// The first run for a locale has to fetch assets; later runs return
/// immediately. Without this the analyzer fails rather than installing them.
func ensureModel(for module: any SpeechModule) async throws {
    if let request = try await AssetInventory.assetInstallationRequest(supporting: [module]) {
        try await request.downloadAndInstall()
    }
}

func transcribe(_ options: Options) async throws {
    guard let path = options.path else { fail("no audio file given") }
    let url = URL(fileURLWithPath: path)
    guard FileManager.default.fileExists(atPath: path) else { fail("no such file: \(path)") }

    let locale = Locale(identifier: options.locale)

    // SpeechTranscriber has the better acoustic model but ignores
    // AnalysisContext.contextualStrings; DictationTranscriber is the older
    // model and does consume them. Which trade wins depends on the speech.
    let module: any SpeechModule = options.dictation
        ? DictationTranscriber(locale: locale, preset: .shortDictation)
        : SpeechTranscriber(
            locale: locale,
            preset: options.alternatives ? .transcriptionWithAlternatives : .transcription
        )
    try await ensureModel(for: module)

    // Vocabulary biasing lives on the analyzer's context here, not on the
    // recognition request as it did in the old API — so the accurate model and
    // term biasing can be used together.
    let context = AnalysisContext()
    if !options.terms.isEmpty {
        context.contextualStrings = [.general: options.terms]
    }

    let file = try AVAudioFile(forReading: url)
    let analyzer = try await SpeechAnalyzer(
        inputAudioFile: file,
        modules: [module],
        analysisContext: context,
        finishAfterFile: true
    )

    // Results arrive per utterance and are concatenated; unlike the old API
    // there is no restart to stitch around.
    var pieces: [String] = []
    var alternatives: [String] = []
    if let dictationModule = module as? DictationTranscriber {
        for try await result in dictationModule.results {
            let text = String(result.text.characters)
            if !text.isEmpty { pieces.append(text) }
        }
    } else if let speechModule = module as? SpeechTranscriber {
        for try await result in speechModule.results {
            let text = String(result.text.characters)
            if !text.isEmpty { pieces.append(text) }
            if options.alternatives {
                alternatives.append(contentsOf: result.alternatives.map { String($0.characters) })
            }
        }
    }
    try await analyzer.finalizeAndFinishThroughEndOfInput()

    var payload: [String: Any] = [
        "text": pieces.joined(separator: " ").trimmingCharacters(in: .whitespaces)
    ]
    if options.alternatives { payload["alternatives"] = alternatives }
    emit(payload)
}

/// Report whether this machine can run the model, without transcribing.
func check(_ options: Options) async throws {
    let locale = Locale(identifier: options.locale)
    let supported = await SpeechTranscriber.supportedLocales.contains {
        $0.identifier(.bcp47) == locale.identifier(.bcp47)
    }
    let transcriber = SpeechTranscriber(locale: locale, preset: .transcription)
    let status = await AssetInventory.status(forModules: [transcriber])
    emit([
        "available": supported,
        "locale": options.locale,
        "assets": String(describing: status),
    ])
}

@main
struct SpeechHelper {
    static func main() async {
        let options = parseArguments()
        do {
            if options.check {
                try await check(options)
            } else {
                try await transcribe(options)
            }
        } catch {
            fail(String(describing: error))
        }
    }
}