import AVFoundation
import Foundation
import Speech

enum Mode: String { case file, live, pipe, listLocales, bench, check }

struct CLIArgs {
    var mode: Mode = .file
    var filePath: String? = nil
    var locale: String = "en-US"
    var verbose: Bool = false
    var includeVolatile: Bool = false
}

func parseArgs() -> CLIArgs {
    var args = CLIArgs()
    var i = 1
    let raw = CommandLine.arguments
    while i < raw.count {
        let a = raw[i]
        switch a {
        case "--live": args.mode = .live
        case "--pipe": args.mode = .pipe
        case "--list-locales": args.mode = .listLocales
        case "--bench": args.mode = .bench
        case "--check": args.mode = .check
        case "--locale": if i + 1 < raw.count { args.locale = raw[i+1]; i += 1 }
        case "--verbose", "-v": args.verbose = true
        case "--include-volatile": args.includeVolatile = true
        case "--help", "-h":
            fputs("""
            Usage:
              apple-speech-transcribe <audio-file> [--locale en-US] [-v] [--include-volatile]
              apple-speech-transcribe --bench <audio-file> [--locale en-US] [-v]
              apple-speech-transcribe --live [--locale en-US] [-v]
              apple-speech-transcribe --pipe [--locale en-US] [-v]   (stdin wav chunks: 4-byte BE len + wav payload)
              apple-speech-transcribe --check
              apple-speech-transcribe --list-locales

            Output (file mode): JSON lines per segment on stdout
            Output (bench):      single JSON object with full transcript + timing
            Output (live/pipe):  JSON lines streaming (isFinal + text)
            """, stderr)
            exit(0)
        default:
            if !a.hasPrefix("-") && args.filePath == nil && (args.mode == .file || args.mode == .bench) {
                args.filePath = a
            }
        }
        i += 1
    }
    return args
}

func logv(_ msg: String, verbose: Bool) { if verbose { fputs("[apple-speech] \(msg)\n", stderr) } }

func resolveLocale(_ id: String) async -> Locale {
    if let sup = await SpeechTranscriber.supportedLocale(equivalentTo: Locale(identifier: id)) { return sup }
    return Locale(identifier: id)
}

func makeTranscriber(locale: Locale, includeVolatile: Bool) -> SpeechTranscriber {
    let rep: Set<SpeechTranscriber.ReportingOption> = includeVolatile ? [.volatileResults, .alternativeTranscriptions] : [.alternativeTranscriptions]
    return SpeechTranscriber(locale: locale, transcriptionOptions: [], reportingOptions: rep, attributeOptions: [.audioTimeRange, .transcriptionConfidence])
}

func ensureAssets(_ transcriber: SpeechTranscriber, verbose: Bool) async -> Bool {
    let status = await AssetInventory.status(forModules: [transcriber])
    switch status {
    case .installed: logv("Assets installed", verbose: verbose); return true
    case .unsupported: fputs("Locale unsupported\n", stderr); return false
    case .supported:
        logv("Assets not installed, downloading...", verbose: true)
        do {
            if let req = try await AssetInventory.assetInstallationRequest(supporting: [transcriber]) { try await req.downloadAndInstall() }
            return true
        } catch { fputs("Asset download failed: \(error)\n", stderr); return false }
    case .downloading:
        logv("Assets downloading, waiting up to 30s...", verbose: true)
        for _ in 0..<30 {
            try? await Task.sleep(nanoseconds: 1_000_000_000)
            let s2 = await AssetInventory.status(forModules: [transcriber])
            if s2 == .installed { return true }
        }
        return false
    @unknown default: return true
    }
}

// MARK: - Check

func runCheck() async {
    print("{\"event\":\"check_start\"}")
    let isAvailable = SpeechTranscriber.isAvailable
    print("{\"event\":\"availability\",\"isAvailable\":\(isAvailable)}")
    if !isAvailable { print("{\"event\":\"check_end\",\"ok\":false}"); return }
    let supported = await SpeechTranscriber.supportedLocales
    let installed = await SpeechTranscriber.installedLocales
    let dict: [String: Any] = [
        "event": "locales",
        "supported": supported.map { $0.identifier }.sorted(),
        "installed": installed.map { $0.identifier }.sorted(),
        "supported_count": supported.count,
        "installed_count": installed.count
    ]
    if let d = try? JSONSerialization.data(withJSONObject: dict), let s = String(data: d, encoding: .utf8) { print(s) }
    let t = SpeechTranscriber(locale: Locale(identifier: "en-US"), preset: .transcription)
    let st = await AssetInventory.status(forModules: [t])
    let str: String
    switch st { case .unsupported: str="unsupported"; case .supported: str="supported"; case .downloading: str="downloading"; case .installed: str="installed"; @unknown default: str="unknown" }
    print("{\"event\":\"asset_status\",\"locale\":\"en-US\",\"status\":\"\(str)\"}")
    print("{\"event\":\"check_end\",\"ok\":true}")
}

func runListLocales() async {
    if !SpeechTranscriber.isAvailable { fputs("Not available\n", stderr); exit(1) }
    for loc in (await SpeechTranscriber.supportedLocales).sorted(by: { $0.identifier < $1.identifier }) { print(loc.identifier) }
}

// MARK: - File transcription (fixed)

struct SegmentJSON: Codable {
    var text: String
    var isFinal: Bool
    var start: Double?
    var duration: Double?
    var alternatives: [String]
}

func runFile(filePath: String, localeId: String, includeVolatile: Bool, bench: Bool, verbose: Bool) async {
    guard FileManager.default.fileExists(atPath: filePath) else { fputs("File not found: \(filePath)\n", stderr); exit(1) }
    guard SpeechTranscriber.isAvailable else { fputs("Not available\n", stderr); exit(1) }

    let locale = await resolveLocale(localeId)
    logv("Locale resolved: \(locale.identifier) (req \(localeId))", verbose: verbose)
    let finalTranscriber: SpeechTranscriber = {
        if bench { return makeTranscriber(locale: locale, includeVolatile: false) }
        return makeTranscriber(locale: locale, includeVolatile: includeVolatile)
    }()

    guard await ensureAssets(finalTranscriber, verbose: verbose) else { exit(2) }

    let fileURL = URL(fileURLWithPath: filePath)
    let audioFile: AVAudioFile
    do { audioFile = try AVAudioFile(forReading: fileURL) } catch { fputs("Open failed: \(error)\n", stderr); exit(4) }
    let audioDuration = Double(audioFile.length) / audioFile.fileFormat.sampleRate
    logv("File: \(filePath) sr=\(audioFile.fileFormat.sampleRate) ch=\(audioFile.fileFormat.channelCount) dur=\(String(format: "%.2f", audioDuration))s", verbose: verbose)

    // Use init(inputAudioFile:finishAfterFile:true) which is the path that emits volatile correctly
    // and automatically handles format.
    let analyzer: SpeechAnalyzer
    do {
        analyzer = try await SpeechAnalyzer(inputAudioFile: audioFile, modules: [finalTranscriber], finishAfterFile: true)
    } catch { fputs("Analyzer init failed: \(error)\n", stderr); exit(5) }

    var segments: [SegmentJSON] = []
    var volatileCount = 0
    var finalCount = 0
    let t0 = Date()

    // Consume results; loop ends naturally when analyzer finishes file
    do {
        for try await result in finalTranscriber.results {
            let text = String(result.text.characters)
            if text.isEmpty { continue }
            let isF = result.isFinal
            if isF { finalCount += 1 } else { volatileCount += 1 }

            let seg = SegmentJSON(
                text: text,
                isFinal: isF,
                start: result.range.start.seconds,
                duration: result.range.duration.seconds,
                alternatives: result.alternatives.map { String($0.characters) }
            )

            if bench {
                if isF { segments.append(seg) }
            } else {
                if !includeVolatile && !isF { continue }
                segments.append(seg)
                if let data = try? JSONEncoder().encode(seg), let line = String(data: data, encoding: .utf8) {
                    print(line); fflush(stdout)
                }
            }
        }
    } catch {
        fputs("Results error: \(error)\n", stderr)
    }

    // Ensure analyzer is done (it should be since results ended)
    _ = analyzer // keep alive

    let elapsed = Date().timeIntervalSince(t0)
    let rtf = audioDuration / max(elapsed, 0.001)

    if bench {
        let fullText = segments.filter { $0.isFinal }.map { $0.text }.joined(separator: " ")
        let out: [String: Any] = [
            "file": filePath,
            "locale": locale.identifier,
            "requested_locale": localeId,
            "audio_duration_sec": audioDuration,
            "processing_sec": elapsed,
            "rtf": rtf,
            "volatile_count": volatileCount,
            "final_count": finalCount,
            "text": fullText,
            "segments": segments.map { ["text": $0.text, "start": $0.start ?? 0, "duration": $0.duration ?? 0] }
        ]
        if let data = try? JSONSerialization.data(withJSONObject: out, options: [.prettyPrinted, .sortedKeys]),
           let str = String(data: data, encoding: .utf8) { print(str) }
    } else {
        if verbose {
            fputs("Done: \(finalCount) final / \(volatileCount) volatile, RTF=\(String(format: "%.1f", rtf))x elapsed=\(String(format: "%.2f", elapsed))s\n", stderr)
        }
    }
}

// MARK: - Live (AnalyzerInput)

actor AudioBufferQueue {
    private var buffers: [AVAudioPCMBuffer] = []
    private var conts: [CheckedContinuation<AVAudioPCMBuffer?, Never>] = []
    private var finished = false
    func enqueue(_ b: AVAudioPCMBuffer) {
        if let c = conts.first { conts.removeFirst(); c.resume(returning: b) } else { buffers.append(b) }
    }
    func dequeue() async -> AVAudioPCMBuffer? {
        if !buffers.isEmpty { return buffers.removeFirst() }
        if finished { return nil }
        return await withCheckedContinuation { conts.append($0) }
    }
    func finish() { finished = true; for c in conts { c.resume(returning: nil) }; conts.removeAll() }
}

struct LiveInputSequence: AsyncSequence, Sendable {
    let queue: AudioBufferQueue
    struct AsyncIterator: AsyncIteratorProtocol {
        let queue: AudioBufferQueue
        mutating func next() async -> AnalyzerInput? {
            guard let buf = await queue.dequeue() else { return nil }
            return AnalyzerInput(buffer: buf)
        }
    }
    func makeAsyncIterator() -> AsyncIterator { AsyncIterator(queue: queue) }
    typealias Element = AnalyzerInput
}

func runLive(localeId: String, verbose: Bool) async {
    guard SpeechTranscriber.isAvailable else { fputs("Not available\n", stderr); exit(1) }
    let locale = await resolveLocale(localeId)
    logv("Live locale \(locale.identifier) (req \(localeId))", verbose: verbose)
    let transcriber = SpeechTranscriber(locale: locale, transcriptionOptions: [], reportingOptions: [.volatileResults, .alternativeTranscriptions], attributeOptions: [.audioTimeRange, .transcriptionConfidence])
    guard await ensureAssets(transcriber, verbose: verbose) else { exit(2) }

    let engine = AVAudioEngine()
    let inputNode = engine.inputNode
    let inputFormat = inputNode.inputFormat(forBus: 0)
    logv("Mic format: \(inputFormat)", verbose: verbose)

    let queue = AudioBufferQueue()
    let inputSequence = LiveInputSequence(queue: queue)
    let analyzer = SpeechAnalyzer(inputSequence: inputSequence, modules: [transcriber])
    let targetFormat = await SpeechAnalyzer.bestAvailableAudioFormat(compatibleWith: [transcriber], considering: inputFormat)
    logv("Target format: \(String(describing: targetFormat))", verbose: verbose)

    let converter: AVAudioConverter? = {
        guard let tf = targetFormat, tf != inputFormat else { return nil }
        return AVAudioConverter(from: inputFormat, to: tf)
    }()

    var shouldStop = false
    let sig = DispatchSource.makeSignalSource(signal: SIGINT, queue: .main)
    signal(SIGINT, SIG_IGN)
    sig.setEventHandler { fputs("\nStopping...\n", stderr); shouldStop = true; Task { await queue.finish() } }
    sig.resume()

    let analyzerTask = Task { do { try await analyzer.start(inputSequence: inputSequence) } catch { fputs("Analyzer err: \(error)\n", stderr) } }
    let resultsTask = Task {
        do {
            for try await result in transcriber.results {
                let dict: [String: Any] = [
                    "text": String(result.text.characters),
                    "isFinal": result.isFinal,
                    "start": result.range.start.seconds,
                    "duration": result.range.duration.seconds,
                    "alternatives": result.alternatives.map { String($0.characters) }
                ]
                if let d = try? JSONSerialization.data(withJSONObject: dict), let s = String(data: d, encoding: .utf8) { print(s); fflush(stdout) }
            }
        } catch { fputs("Results err: \(error)\n", stderr) }
    }

    if let tf = targetFormat, let conv = converter {
        inputNode.installTap(onBus: 0, bufferSize: 4096, format: inputFormat) { buffer, _ in
            if shouldStop { return }
            let cap = AVAudioFrameCount(Double(buffer.frameLength) * tf.sampleRate / inputFormat.sampleRate + 10)
            guard let out = AVAudioPCMBuffer(pcmFormat: tf, frameCapacity: cap) else { return }
            var err: NSError?
            let block: AVAudioConverterInputBlock = { _, outStatus in outStatus.pointee = .haveData; return buffer }
            conv.convert(to: out, error: &err, withInputFrom: block)
            if err == nil && out.frameLength > 0 { Task { await queue.enqueue(out) } }
        }
    } else {
        inputNode.installTap(onBus: 0, bufferSize: 4096, format: inputFormat) { buffer, _ in
            if shouldStop { return }
            guard let copy = AVAudioPCMBuffer(pcmFormat: buffer.format, frameCapacity: buffer.frameCapacity) else { return }
            copy.frameLength = buffer.frameLength
            if let src = buffer.floatChannelData, let dst = copy.floatChannelData {
                for ch in 0..<Int(buffer.format.channelCount) { memcpy(dst[ch], src[ch], Int(buffer.frameLength) * MemoryLayout<Float>.size) }
            }
            Task { await queue.enqueue(copy) }
        }
    }

    engine.prepare()
    do { try engine.start() } catch { fputs("Audio engine start failed: \(error)\n", stderr); exit(6) }
    logv("Listening... Ctrl-C to stop", verbose: true)
    while !shouldStop { try? await Task.sleep(nanoseconds: 100_000_000) }
    inputNode.removeTap(onBus: 0)
    engine.stop()
    await queue.finish()
    analyzerTask.cancel()
    resultsTask.cancel()
}

@main
struct AppleSpeechCLI {
    static func main() async {
        let cli = parseArgs()
        switch cli.mode {
        case .check: await runCheck()
        case .listLocales: await runListLocales()
        case .file:
            guard let fp = cli.filePath else {
                fputs("Usage: apple-speech-transcribe <file> [--locale en-US] [-v] [--include-volatile]\n", stderr); exit(1)
            }
            await runFile(filePath: fp, localeId: cli.locale, includeVolatile: cli.includeVolatile, bench: false, verbose: cli.verbose)
        case .bench:
            guard let fp = cli.filePath else { fputs("Usage: --bench <file> [--locale]\n", stderr); exit(1) }
            await runFile(filePath: fp, localeId: cli.locale, includeVolatile: false, bench: true, verbose: cli.verbose)
        case .live: await runLive(localeId: cli.locale, verbose: cli.verbose)
        case .pipe: await runPipe(localeId: cli.locale, verbose: cli.verbose)
        }
    }
}

// MARK: - Pipe (stdin wav chunks) for draft streaming from Python VAD

actor WavChunkQueue {
    private var chunks: [URL] = []
    private var conts: [CheckedContinuation<URL?, Never>] = []
    private var finished = false
    func enqueue(_ url: URL) {
        if let c = conts.first { conts.removeFirst(); c.resume(returning: url) } else { chunks.append(url) }
    }
    func dequeue() async -> URL? {
        if !chunks.isEmpty { return chunks.removeFirst() }
        if finished { return nil }
        return await withCheckedContinuation { conts.append($0) }
    }
    func finish() { finished = true; for c in conts { c.resume(returning: nil) }; conts.removeAll() }
}

func runPipe(localeId: String, verbose: Bool) async {
    guard SpeechTranscriber.isAvailable else { fputs("Not available\n", stderr); exit(1) }
    let locale = await resolveLocale(localeId)
    let warm = SpeechTranscriber(locale: locale, transcriptionOptions: [], reportingOptions: [.volatileResults], attributeOptions: [])
    guard await ensureAssets(warm, verbose: verbose) else { exit(2) }
    logv("Pipe ready locale=\(locale.identifier)", verbose: true)

    let stdinH = FileHandle.standardInput
    var leftover = Data()
    var chunkIdx = 0

    func readExact(_ n: Int) -> Data? {
        var out = Data()
        out.reserveCapacity(n)
        // first consume leftover
        if leftover.count >= n {
            let d = leftover.prefix(n); leftover = leftover.dropFirst(n); return Data(d)
        }
        if leftover.count > 0 { out.append(leftover); leftover = Data() }
        while out.count < n {
            let d = stdinH.readData(ofLength: n - out.count)
            if d.isEmpty {
                // EOF? try one more small sleep then retry once if we have partial
                if out.count == 0 { return nil }
                // partial -> return nil -> broken stream
                return nil
            }
            out.append(d)
        }
        return out
    }

    while true {
        guard let lenData = readExact(4) else { break }
        let length = lenData.withUnsafeBytes { $0.load(as: UInt32.self).bigEndian }
        if length == 0 { logv("Pipe EOF", verbose: true); break }
        if length > 20_000_000 { fputs("Pipe: chunk too large \(length)\n", stderr); break }
        guard let chunkData = readExact(Int(length)) else {
            fputs("Pipe: truncated, expected \(length)\n", stderr); break
        }
        chunkIdx += 1
        let tmpURL = URL(fileURLWithPath: NSTemporaryDirectory()).appendingPathComponent("apple-pipe-\(ProcessInfo.processInfo.processIdentifier)-\(chunkIdx).wav")
        do { try chunkData.write(to: tmpURL) } catch { fputs("Pipe write err: \(error)\n", stderr); continue }
        guard let audioFile = try? AVAudioFile(forReading: tmpURL) else {
            try? FileManager.default.removeItem(at: tmpURL)
            fputs("Pipe open fail \(chunkIdx)\n", stderr); continue
        }
        let t = SpeechTranscriber(locale: locale, transcriptionOptions: [], reportingOptions: [.volatileResults], attributeOptions: [.audioTimeRange])
        guard let analyzer = try? await SpeechAnalyzer(inputAudioFile: audioFile, modules: [t], finishAfterFile: true) else {
            try? FileManager.default.removeItem(at: tmpURL); continue
        }
        do {
            for try await res in t.results {
                let txt = String(res.text.characters).trimmingCharacters(in: .whitespacesAndNewlines)
                if txt.isEmpty { continue }
                let d: [String: Any] = [
                    "event": res.isFinal ? "final" : "draft",
                    "text": txt, "isFinal": res.isFinal,
                    "chunk": chunkIdx, "start": res.range.start.seconds,
                    "duration": res.range.duration.seconds
                ]
                if let jd = try? JSONSerialization.data(withJSONObject: d), let s = String(data: jd, encoding: .utf8) {
                    print(s); fflush(stdout)
                }
            }
        } catch { fputs("Pipe results err \(chunkIdx): \(error)\n", stderr) }
        _ = analyzer
        try? FileManager.default.removeItem(at: tmpURL)
    }
    logv("Pipe done \(chunkIdx) chunks", verbose: true)
}