LocalTranscriber.swift

import AVFoundation
import Foundation
import Speech

protocol Transcribing: Sendable {
    func transcribe(
        localMP4: URL,
        languageIdentifier: String,
        status: @escaping @Sendable (ReviewStatus) async -> Void
    ) async throws -> String
}

struct TranscriptSegment: Sendable {
    let text: String
    let start: TimeInterval
}

func timestampedTranscript(_ segments: [TranscriptSegment]) -> String {
    segments.compactMap { segment in
        let text = segment.text.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !text.isEmpty else { return nil }
        let tenths = Int(max(0, segment.start.isFinite ? segment.start : 0) * 10)
        let seconds = tenths / 10
        let timestamp = seconds >= 3_600
            ? String(format: "%02d:%02d:%02d.%d", seconds / 3_600, seconds / 60 % 60, seconds % 60, tenths % 10)
            : String(format: "%02d:%02d.%d", seconds / 60, seconds % 60, tenths % 10)
        return "[\(timestamp)] \(text)"
    }.joined(separator: "\n")
}

enum LocalTranscriptionError: LocalizedError {
    case unavailable
    case unsupportedLanguage
    case invalidRecording
    case noSpeech

    var errorDescription: String? {
        switch self {
        case .unavailable: "On-device transcription is unavailable on this Mac."
        case .unsupportedLanguage: "This spoken language is unavailable for local transcription. Choose another in Settings, then retry."
        case .invalidRecording: "The recording has no readable audio. The original video is still saved."
        case .noSpeech: "No speech was detected. The original video is still saved."
        }
    }
}

struct LocalTranscriber: Transcribing {
    func transcribe(
        localMP4: URL,
        languageIdentifier: String,
        status: @escaping @Sendable (ReviewStatus) async -> Void
    ) async throws -> String {
        guard localMP4.isFileURL, FileManager.default.isReadableFile(atPath: localMP4.path) else {
            throw LocalTranscriptionError.invalidRecording
        }
        guard SpeechTranscriber.isAvailable else { throw LocalTranscriptionError.unavailable }
        guard let locale = await SpeechTranscriber.supportedLocale(equivalentTo: Locale(identifier: languageIdentifier)) else {
            throw LocalTranscriptionError.unsupportedLanguage
        }
        let transcriber = SpeechTranscriber(
            locale: locale,
            transcriptionOptions: [],
            reportingOptions: [],
            attributeOptions: [.audioTimeRange]
        )
        await status(.preparing)
        if let download = try await AssetInventory.assetInstallationRequest(supporting: [transcriber]) {
            try await download.downloadAndInstall()
        }
        try Task.checkCancellation()
        await status(.transcribing)

        // Extract locally so the file reader receives audio without the video track.
        let audioURL = FileManager.default.temporaryDirectory.appendingPathComponent("Repro-\(UUID()).m4a")
        defer { try? FileManager.default.removeItem(at: audioURL) }
        let asset = AVURLAsset(url: localMP4)
        guard try await !asset.loadTracks(withMediaType: .audio).isEmpty,
              let export = AVAssetExportSession(asset: asset, presetName: AVAssetExportPresetAppleM4A) else {
            throw LocalTranscriptionError.invalidRecording
        }
        try await export.export(to: audioURL, as: .m4a)
        let file = try AVAudioFile(forReading: audioURL)
        let analyzer = SpeechAnalyzer(modules: [transcriber])
        async let segments = collectResults(from: transcriber)
        do {
            if let lastSample = try await analyzer.analyzeSequence(from: file) {
                try await analyzer.finalizeAndFinish(through: lastSample)
            } else {
                await analyzer.cancelAndFinishNow()
                throw LocalTranscriptionError.noSpeech
            }
            let text = timestampedTranscript(try await segments)
            guard !text.isEmpty else { throw LocalTranscriptionError.noSpeech }
            return text
        } catch {
            await analyzer.cancelAndFinishNow()
            throw error
        }
    }

    private func collectResults(from transcriber: SpeechTranscriber) async throws -> [TranscriptSegment] {
        var segments: [TranscriptSegment] = []
        for try await result in transcriber.results {
            let text = result.text
            let runs = Array(text.runs)
            // Keep short timestamped spans for locating narration in the video.
            for index in stride(from: 0, to: runs.count, by: 4) {
                let chunk = runs[index..<min(index + 4, runs.count)]
                let start = runs[index].audioTimeRange?.start ?? result.range.start
                let words = chunk.map { String(text[$0.range].characters) }.joined()
                segments.append(TranscriptSegment(text: words, start: start.seconds))
            }
        }
        return segments
    }
}