import Foundation
import Speech
import ArgumentParser

/// Swift CLI 處理器：使用 Apple Speech Framework 進行語音辨識
/// 作為 Python ASR (faster-whisper) 的替代方案
///
/// 比較項目：
/// - Speech Framework 使用 Apple 內建模型（ANE 加速）
/// - 無需下載模型，系統內建
/// - 支援即時與批次辨識
/// - 語言支援依作業系統版本

/// 合併逐字 segments 成句子
/// 當相鄰 segment 間隔 < 0.5s 時合併，間隔 >= 0.5s 時視為句子邊界
func mergeWordSegments(_ segments: [[String: Any]]) -> [[String: Any]] {
    let gapThreshold: TimeInterval = 0.5
    var merged: [[String: Any]] = []
    var current: [String: Any]? = nil

    for seg in segments {
        guard let start = seg["start"] as? TimeInterval,
              let end = seg["end"] as? TimeInterval,
              let text = seg["text"] as? String,
              let conf = seg["confidence"] as? Float else {
            continue
        }

        if var cur = current {
            let curEnd = cur["end"] as? TimeInterval ?? 0
            let gap = start - curEnd

            if gap < gapThreshold {
                // 合併到當前句子
                let curText = cur["text"] as? String ?? ""
                let curConf = cur["confidence"] as? Float ?? 0
                let wordCount = cur["_wordCount"] as? Int ?? 0
                cur["text"] = curText + (curText.hasSuffix(" ") ? "" : " ") + text
                cur["end"] = end
                // 用 confidence 加權平均（依字數）
                let totalWords = wordCount + 1
                cur["confidence"] = (curConf * Float(wordCount) + conf) / Float(totalWords)
                cur["_wordCount"] = totalWords
                current = cur
            } else {
                // 句子邊界：儲存當前，開始新的
                cur.removeValue(forKey: "_wordCount")
                merged.append(cur)
                current = [
                    "start": start,
                    "end": end,
                    "text": text,
                    "confidence": conf,
                    "_wordCount": 1
                ]
            }
        } else {
            current = [
                "start": start,
                "end": end,
                "text": text,
                "confidence": conf,
                "_wordCount": 1
            ]
        }
    }

    if let cur = current {
        var finalCur = cur
        finalCur.removeValue(forKey: "_wordCount")
        merged.append(finalCur)
    }

    return merged
}

@main
struct ASRSwift: ParsableCommand {
    @Argument(help: "音訊/影片檔案路徑")
    var inputPath: String

    @Argument(help: "輸出 JSON 路徑")
    var outputPath: String

    @Option(name: .long, help: "UUID for Redis")
    var uuid: String = ""

    @Option(name: .long, help: "語言 (留空則自動嘗試支援的語種)")
    var language: String = ""

    /// 嘗試自動偵測語種
    func detectLanguage() -> String {
        if !language.isEmpty { return language }
        // 優先嘗試常用語種
        let candidates = ["zh-TW", "zh-Hans", "en-US", "ja-JP", "ko-KR"]
        for localeId in candidates {
            if let reco = SFSpeechRecognizer(locale: Locale(identifier: localeId)), reco.isAvailable {
                print("[ASR_Swift] Auto-detected language: \(localeId)")
                return localeId
            }
        }
        return "en-US"
    }

    mutating func run() throws {
        let startTime = Date()

        print("[ASR_Swift] Starting: \(inputPath)")
        print("[ASR_Swift] Language: \(language)")
        print("[ASR_Swift] Output: \(outputPath)")

        // 1. 萃取音訊（若為影片檔）
        let audioURL: URL
        let ext = (inputPath as NSString).pathExtension.lowercased()
        let tempDir = FileManager.default.temporaryDirectory
            .appendingPathComponent("asr_swift_\(UUID().uuidString)")

        if ["mp4", "mov", "mkv", "avi"].contains(ext) {
            // 需要 ffmpeg 萃取音訊
            let wavPath = tempDir.appendingPathComponent("audio.wav").path
            print("[ASR_Swift] Extracting audio from video...")
            let proc = Process()
            proc.executableURL = URL(fileURLWithPath: "/opt/homebrew/bin/ffmpeg")
            proc.arguments = ["-y", "-v", "quiet", "-i", inputPath,
                              "-ar", "16000", "-ac", "1", wavPath]
            try proc.run()
            proc.waitUntilExit()
            guard FileManager.default.fileExists(atPath: wavPath) else {
                print("[ASR_Swift] Error: ffmpeg failed to extract audio")
                throw NSError(domain: "ASRSwift", code: 1, userInfo: nil)
            }
            audioURL = URL(fileURLWithPath: wavPath)
            print("[ASR_Swift] Audio extracted: \(wavPath)")
        } else {
            audioURL = URL(fileURLWithPath: inputPath)
        }

        // 2. 語音辨識
        print("[ASR_Swift] Starting recognition...")

        // 請求授權
        let authGroup = DispatchGroup()
        authGroup.enter()
        var authStatus: SFSpeechRecognizerAuthorizationStatus = .notDetermined
        SFSpeechRecognizer.requestAuthorization { status in
            authStatus = status
            authGroup.leave()
        }
        authGroup.wait()
        guard authStatus == .authorized else {
            print("[ASR_Swift] Speech recognition not authorized: \(authStatus.rawValue)")
            return
        }
        print("[ASR_Swift] Speech recognition authorized")

        let finalLang = detectLanguage()
        let recognizer = SFSpeechRecognizer(locale: Locale(identifier: finalLang))
        guard let recognizer = recognizer, recognizer.isAvailable else {
            print("[ASR_Swift] Error: Speech recognizer not available for \(language)")
            // 回傳空結果
            let emptyResult: [String: Any] = [
                "language": language,
                "segments": [],
                "processing_time": Date().timeIntervalSince(startTime),
                "model": "Apple Speech Framework",
                "error": "Recognizer not available"
            ]
            let jsonData = try JSONSerialization.data(withJSONObject: emptyResult, options: [.prettyPrinted])
            try jsonData.write(to: URL(fileURLWithPath: outputPath))
            return
        }

        let request = SFSpeechURLRecognitionRequest(url: audioURL)
        request.shouldReportPartialResults = false
        request.taskHint = .dictation

        var allSegments: [[String: Any]] = []
        let semaphore = DispatchSemaphore(value: 0)

        let task = recognizer.recognitionTask(with: request) { result, error in
            if let error = error {
                print("[ASR_Swift] Recognition error: \(error.localizedDescription)")
                semaphore.signal()
                return
            }

            if let result = result, result.isFinal {
                let duration = Date().timeIntervalSince(startTime)
                print("[ASR_Swift] Recognition completed in \(String(format: "%.2f", duration))s")

                // 將辨識結果轉為 segment 格式
                for segment in result.bestTranscription.segments {
                    let seg: [String: Any] = [
                        "start": segment.timestamp,
                        "end": segment.timestamp + segment.duration,
                        "text": segment.substring,
                        "speaker_id": nil as String?,
                        "confidence": segment.confidence
                    ]
                    allSegments.append(seg)
                }

                // 合併逐字 segments 成句子
                if !allSegments.isEmpty {
                    let beforeCount = allSegments.count
                    allSegments = mergeWordSegments(allSegments)
                    print("[ASR_Swift] Merged segments: \(beforeCount) → \(allSegments.count)")
                }

                // 若無 segment，用整個文字建立一個
                if allSegments.isEmpty {
                    let fullText = result.bestTranscription.formattedString
                    let seg: [String: Any] = [
                        "start": 0.0,
                        "end": Date().timeIntervalSince(startTime),
                        "text": fullText,
                        "speaker_id": nil as String?,
                        "confidence": 1.0
                    ]
                    allSegments.append(seg)
                }

                semaphore.signal()
            }
        }

        // RunLoop 是必要的 — Speech Framework 需要 main runloop 觸發 callback
        // 使用 semaphore 等待完成，避免固定 timeout
        while semaphore.wait(timeout: .now()) == .timedOut {
            RunLoop.current.run(mode: .default, before: Date(timeIntervalSinceNow: 0.1))
        }
        task.cancel()

        // 3. 輸出 JSON
        let outputDict: [String: Any] = [
            "language": language,
            "segments": allSegments,
            "processing_time": Date().timeIntervalSince(startTime),
            "model": "Apple Speech Framework (ANE accelerated)",
            "total_segments": allSegments.count
        ]

        let jsonData = try JSONSerialization.data(withJSONObject: outputDict, options: [.prettyPrinted])
        try jsonData.write(to: URL(fileURLWithPath: outputPath))

        print("[ASR_Swift] Saved \(allSegments.count) segments to \(outputPath)")
        print("[ASR_Swift] Total time: \(String(format: "%.2f", Date().timeIntervalSince(startTime)))s")

        // 清理暫存
        if tempDir != audioURL.deletingLastPathComponent() {
            try? FileManager.default.removeItem(at: tempDir)
        }
    }
}
