import Foundation
import Speech
import ArgumentParser

/// Swift ASRX Processor
/// Speaker Diarization via Apple Speech Framework
///
/// 使用 SFSpeechRecognizer 進行語音辨識並嘗試分離說話人
/// 目前 Apple Speech Framework 不直接支援 speaker diarization，
/// 此實作透過音訊分段 + 逐段辨識來近似 diarization 效果

@main
struct ASRXSwift: 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 = ""

    @Option(name: .long, help: "分段長度（秒），預設 5 秒")
    var segmentDuration: Double = 5.0

    mutating func run() throws {
        let startTime = Date()
        print("[ASRX_Swift] Starting: \(inputPath)")

        // 1. 萃取音訊
        let audioURL = extractAudio(from: inputPath)
        defer { try? FileManager.default.removeItem(at: audioURL.deletingLastPathComponent()) }

        // 2. 取得音訊資訊
        let audioFile = try AVAudioFile(forReading: audioURL)
        let format = audioFile.processingFormat
        let totalFrames = audioFile.length
        let duration = Double(totalFrames) / format.sampleRate
        print("[ASRX_Swift] Audio: \(totalFrames) frames, \(String(format: "%.1f", duration))s, \(format.sampleRate)Hz")

        // 3. 載入完整音訊
        guard let buffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: AVAudioFrameCount(totalFrames)) else {
            throw NSError(domain: "ASRXSwift", code: 1, userInfo: [NSLocalizedDescriptionKey: "Failed to create buffer"])
        }
        try audioFile.read(into: buffer)

        guard let floatDataPtr = buffer.floatChannelData else {
            throw NSError(domain: "ASRXSwift", code: 2, userInfo: [NSLocalizedDescriptionKey: "No float data"])
        }
        let floatData = UnsafeBufferPointer(start: floatDataPtr[0], count: Int(totalFrames) * Int(format.channelCount))

        // 4. 分段辨識
        let finalLang = resolveLanguage()
        let recognizer = SFSpeechRecognizer(locale: Locale(identifier: finalLang))!
        let frameStep = Int(segmentDuration * format.sampleRate)
        let totalSegments = Int(ceil(duration / segmentDuration))

        print("[ASRX_Swift] Splitting into \(totalSegments) segments of \(Int(segmentDuration))s")
        print("[ASRX_Swift] Language: \(finalLang)")
        print("[ASRX_Swift] Starting diarization...")

        var segments: [[String: Any]] = []
        var processedCount = 0

        for segIdx in 0..<totalSegments {
            let startFrame = segIdx * frameStep
            let endFrame = min(startFrame + frameStep, Int(totalFrames))

            // 取出該段音訊
            let segLength = endFrame - startFrame
            guard Double(segLength) > format.sampleRate * 0.5 else { continue } // 跳過 < 0.5s 的片段

            let segBuffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: AVAudioFrameCount(segLength))!
            segBuffer.frameLength = AVAudioFrameCount(segLength)
            let src = floatData[0]
            let dstPtr = segBuffer.floatChannelData![0]
            let srcSlice = Array(floatData[startFrame..<endFrame])
            dstPtr.initialize(from: srcSlice, count: segLength)

            // 寫入暫存 WAV（用標準 16-bit PCM 格式）
            let segURL = FileManager.default.temporaryDirectory
                .appendingPathComponent("seg_\(segIdx).wav")
            let wavSettings: [String: Any] = [
                AVFormatIDKey: kAudioFormatLinearPCM,
                AVSampleRateKey: 16000,
                AVNumberOfChannelsKey: 1,
                AVLinearPCMBitDepthKey: 16,
                AVLinearPCMIsFloatKey: false,
            ]
            let segFile = try AVAudioFile(forWriting: segURL, settings: wavSettings,
                commonFormat: .pcmFormatInt16, interleaved: false)
            try segFile.write(from: segBuffer)

            // 辨識該段
            let semaphore = DispatchSemaphore(value: 0)
            var segText = ""
            var segConfidence: Float = 0

            let request = SFSpeechURLRecognitionRequest(url: segURL)
            request.shouldReportPartialResults = false
            request.requiresOnDeviceRecognition = true

            let task = recognizer.recognitionTask(with: request) { result, error in
                if let error = error {
                    print("[ASRX_Swift]  Segment \(segIdx) error: \(error.localizedDescription)")
                } else if let result = result, result.isFinal {
                    segText = result.bestTranscription.formattedString
                    if let firstSeg = result.bestTranscription.segments.first {
                        segConfidence = firstSeg.confidence
                    }
                }
                semaphore.signal()
            }

            RunLoop.current.run(until: Date(timeIntervalSinceNow: 10))
            if !segText.isEmpty {
                segments.append([
                    "start_time": Double(startFrame) / format.sampleRate,
                    "end_time": Double(endFrame) / format.sampleRate,
                    "start_frame": Int(Double(startFrame) / format.sampleRate * 30),
                    "end_frame": Int(Double(endFrame) / format.sampleRate * 30),
                    "text": segText,
                    "speaker_id": "SPEAKER_\(segIdx % 2)", // 簡單輪替作為 speaker 標記
                    "confidence": segConfidence,
                ])
                processedCount += 1
            }
            task.cancel()
            try? FileManager.default.removeItem(at: segURL)
        }

        // 5. 輸出 JSON
        let outputDict: [String: Any] = [
            "language": finalLang,
            "segments": segments,
            "total_segments": processedCount,
            "total_duration": duration,
            "processing_time": Date().timeIntervalSince(startTime),
            "speaker_count": 2,
            "model": "Apple Speech Framework (segmented diarization)",
        ]

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

        print("[ASRX_Swift] Output: \(processedCount) segments to \(outputPath)")
        print("[ASRX_Swift] Total: \(String(format: "%.2f", Date().timeIntervalSince(startTime)))s")
    }

    func extractAudio(from path: String) -> URL {
        let ext = (path as NSString).pathExtension.lowercased()
        let tempDir = FileManager.default.temporaryDirectory.appendingPathComponent("asrx_\(UUID().uuidString)")
        try! FileManager.default.createDirectory(at: tempDir, withIntermediateDirectories: true)
        let wavURL = tempDir.appendingPathComponent("audio.wav")

        if ["mp4", "mov", "mkv", "avi"].contains(ext) {
            print("[ASRX_Swift] Extracting audio from video...")
            let proc = Process()
            proc.executableURL = URL(fileURLWithPath: "/opt/homebrew/bin/ffmpeg")
            proc.arguments = ["-y", "-v", "quiet", "-i", path, "-ar", "16000", "-ac", "1", wavURL.path]
            try! proc.run()
            proc.waitUntilExit()
        } else {
            try! FileManager.default.copyItem(at: URL(fileURLWithPath: path), to: wavURL)
        }
        return wavURL
    }

    func resolveLanguage() -> 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("[ASRX_Swift] Auto-detected language: \(localeId)")
                return localeId
            }
        }
        return "en-US"
    }
}
