notes/Notes/Services/TranscriptionService.swift

import AVFoundation
import Foundation
import IndieSync
import Speech
import SwiftData

/// On-device, after-the-fact transcription of voice notes. `@Observable
/// @MainActor` with a FIFO one-at-a-time queue: a `SpeechAnalyzer` is memory- and
/// model-heavy, so notes are transcribed strictly one after another. The analyzer
/// work is awaited (the analyzer is an actor), so it doesn't block the main actor
/// despite this class being main-actor confined.
///
/// Only this device's own pending notes are ever enqueued — ownership is decided
/// by `AudioInfo.transcriberDeviceID` upstream (`SyncEngine`), which this service
/// trusts. A failed job is terminal until the user asks to retry.
@Observable
@MainActor
final class TranscriptionService {
    private let syncEngine: SyncEngine
    private let settings: TranscriptionSettings
    private let modelContainer: ModelContainer

    /// Seconds of audio the sampling pass listens to when choosing a language.
    private static let sampleSeconds: Double = 15

    private struct Job {
        let noteID: ULID
        /// A fresh recording's temp file, passed on enqueue for the fast path; nil
        /// for re-transcription, which pulls the audio back down from the container.
        let localAudioURL: URL?
        /// A forced language (re-transcribe menu), else auto-select.
        let forcedLocaleID: String?
    }

    private var queue: [Job] = []
    /// Ids currently queued or in flight, so overlapping triggers (launch resume,
    /// watch ingest, a re-transcribe) can't enqueue the same note twice.
    private var inFlightIDs: Set<String> = []
    private var pumpTask: Task<Void, Never>?

    init(syncEngine: SyncEngine, settings: TranscriptionSettings, modelContainer: ModelContainer) {
        self.syncEngine = syncEngine
        self.settings = settings
        self.modelContainer = modelContainer
    }

    // MARK: - Queue

    /// Queue a note for transcription. `localAudioURL` is the fresh-recording fast
    /// path (skips the container round-trip); `forcedLocaleID` skips language
    /// auto-selection. Duplicate enqueues of an in-flight note are ignored.
    func enqueue(noteID: ULID, localAudioURL: URL? = nil, forcedLocaleID: String? = nil) {
        let key = noteID.stringValue
        guard !inFlightIDs.contains(key) else { return }
        inFlightIDs.insert(key)
        queue.append(Job(noteID: noteID, localAudioURL: localAudioURL, forcedLocaleID: forcedLocaleID))
        pump()
    }

    /// On launch, re-enqueue every note this device owns that is still pending —
    /// recovering work that a crash or kill interrupted mid-transcription. Reads
    /// the cache (already rebuilt by the sync engine's connect).
    func resumePendingWork() {
        let context = modelContainer.mainContext
        let installID = TranscriptionSettings.installID
        let pendingRaw = Note.TranscriptionStatus.pending.rawValue
        let descriptor = FetchDescriptor<NoteEntity>(predicate: #Predicate {
            $0.hasAudio
                && $0.transcriptionStatusRaw == pendingRaw
                && $0.transcriberDeviceID == installID
        })
        let rows = (try? context.fetch(descriptor)) ?? []
        for row in rows {
            if let id = ULID(string: row.id) { enqueue(noteID: id) }
        }
    }

    private func pump() {
        guard pumpTask == nil else { return }
        pumpTask = Task { [weak self] in
            while let self, let job = self.dequeue() {
                await self.process(job)
                self.inFlightIDs.remove(job.noteID.stringValue)
            }
            self?.pumpTask = nil
        }
    }

    private func dequeue() -> Job? {
        queue.isEmpty ? nil : queue.removeFirst()
    }

    // MARK: - Processing

    /// Transcribe one job, retrying once before marking the note failed. A failed
    /// note is left for the user to retry explicitly (re-transcribe / Retry).
    private func process(_ job: Job) async {
        do {
            try await transcribe(job)
        } catch {
            do {
                try await transcribe(job)
            } catch {
                print("[Transcription] \(job.noteID.stringValue) failed: \(error)")
                try? await syncEngine.markTranscriptionFailed(id: job.noteID)
            }
        }
    }

    private func transcribe(_ job: Job) async throws {
        let audio = try await resolveAudioFile(for: job)
        defer { cleanUp(audio: audio, job: job) }

        let locale = try await selectLocale(fileURL: audio.url, forcedLocaleID: job.forcedLocaleID)
        let transcript = try await runTranscription(fileURL: audio.url, locale: locale)
        let text = String(transcript.characters).trimmingCharacters(in: .whitespacesAndNewlines)
        try await syncEngine.applyTranscription(
            id: job.noteID,
            text: text,
            localeID: TranscriptionSettings.canonicalID(locale),
            // Real mean confidence when the engine attributed any runs; nil for
            // a confidence-less engine so it's never mistaken for a measured 0.
            confidence: LanguagePicker.hasConfidence(of: transcript)
                ? LanguagePicker.meanConfidence(of: transcript) : nil,
            languageWasForced: job.forcedLocaleID != nil
        )
    }

    // MARK: - Audio acquisition

    private struct AudioFile {
        let url: URL
        /// Whether we wrote a temp copy that must be cleaned up after the job.
        let isTemporary: Bool
    }

    /// The fresh recording if it's still on disk, else a temp copy pulled back
    /// from the container (bounded wait for an evicted file).
    private func resolveAudioFile(for job: Job) async throws -> AudioFile {
        if let local = job.localAudioURL, FileManager.default.fileExists(atPath: local.path) {
            return AudioFile(url: local, isTemporary: false)
        }
        let data = try await syncEngine.loadAudioDataWaiting(for: job.noteID)
        let url = FileManager.default.temporaryDirectory
            .appending(path: "transcribe-\(job.noteID.stringValue).m4a")
        try data.write(to: url)
        return AudioFile(url: url, isTemporary: true)
    }

    private func cleanUp(audio: AudioFile, job: Job) {
        let fm = FileManager.default
        if audio.isTemporary { try? fm.removeItem(at: audio.url) }
        // A fresh recording's temp file is no longer needed once the container
        // holds the bytes; drop it to keep the temp dir tidy.
        if let local = job.localAudioURL,
           local.path.hasPrefix(fm.temporaryDirectory.path) {
            try? fm.removeItem(at: local)
        }
    }

    // MARK: - Language selection

    private func selectLocale(fileURL: URL, forcedLocaleID: String?) async throws -> Locale {
        if let forcedLocaleID {
            return Locale(identifier: forcedLocaleID)
        }
        let candidates = await settings.enabledInstalledLocales()
        guard !candidates.isEmpty else { throw TranscriptionError.noInstalledLanguage }
        if candidates.count == 1 { return candidates[0] }

        // Sampling pass: transcribe a short clip in each language and pick the
        // most confident. Run sequentially — one analyzer at a time keeps the
        // memory ceiling low.
        let sample = try makeSampleClip(from: fileURL, seconds: Self.sampleSeconds)
        defer { if sample.isTemporary { try? FileManager.default.removeItem(at: sample.url) } }

        var scored: [LanguageCandidate] = []
        for locale in candidates {
            let transcript = (try? await runTranscription(fileURL: sample.url, locale: locale)) ?? AttributedString()
            // `rankingConfidence`, not `meanConfidence`: a `DictationTranscriber`
            // locale may leave runs unattributed, so an unscored-but-non-empty
            // transcript gets a neutral score and competes on length rather than
            // losing outright to a confidence-reporting `SpeechTranscriber` locale.
            scored.append(LanguageCandidate(
                localeID: TranscriptionSettings.canonicalID(locale),
                meanConfidence: LanguagePicker.rankingConfidence(of: transcript),
                transcriptLength: transcript.characters.count
            ))
        }
        guard let winner = LanguagePicker.pick(from: scored) else { return candidates[0] }
        return Locale(identifier: winner.localeID)
    }

    /// The first `seconds` of the recording as its own file, so the sampling pass
    /// isn't run over a whole long note. Returns the original file (no copy) when
    /// the recording is already shorter than the sample window.
    private func makeSampleClip(from fileURL: URL, seconds: Double) throws -> AudioFile {
        let source = try AVAudioFile(forReading: fileURL)
        let format = source.processingFormat
        let wantFrames = format.sampleRate * seconds
        guard Double(source.length) > wantFrames else {
            return AudioFile(url: fileURL, isTemporary: false)
        }
        let frameCount = AVAudioFrameCount(wantFrames)
        guard frameCount > 0, let buffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: frameCount) else {
            throw TranscriptionError.audioUnreadable
        }
        try source.read(into: buffer, frameCount: frameCount)
        let clipURL = FileManager.default.temporaryDirectory
            .appending(path: "sample-\(UUID().uuidString).caf")
        let clip = try AVAudioFile(forWriting: clipURL, settings: format.settings)
        try clip.write(from: buffer)
        return AudioFile(url: clipURL, isTemporary: true)
    }

    // MARK: - Analyzer

    /// Transcribe a whole file in one locale and return the finalized attributed
    /// transcript (carrying per-run confidence where the engine emits it). Routes
    /// to whichever engine supports the locale — `SpeechTranscriber` or, for its
    /// dictation-only locales (ru-RU, uk-UA, pl-PL …), `DictationTranscriber` —
    /// via `TranscriptionSettings.engine(for:)`. Both request confidence; the two
    /// `Result` types share `isFinal`/`text` but only concretely, so the engine
    /// branch supplies the module and its result stream to a shared `analyze`.
    private func runTranscription(fileURL: URL, locale: Locale) async throws -> AttributedString {
        switch settings.engine(for: locale) {
        case .speech:
            let transcriber = SpeechTranscriber(
                locale: locale,
                transcriptionOptions: [],
                reportingOptions: [],
                attributeOptions: [.transcriptionConfidence]
            )
            return try await analyze(module: transcriber, results: transcriber.results, fileURL: fileURL) { $0.text }
        case .dictation:
            let transcriber = DictationTranscriber(
                locale: locale,
                contentHints: [],
                transcriptionOptions: [],
                reportingOptions: [],
                attributeOptions: [.transcriptionConfidence]
            )
            return try await analyze(module: transcriber, results: transcriber.results, fileURL: fileURL) { $0.text }
        }
    }

    /// Drive `SpeechAnalyzer` over the whole file with a single module and fold its
    /// finalized results into one transcript. Generic over the module's result
    /// stream so one pipeline serves both engines (their `Result`s share `isFinal`
    /// via `SpeechModuleResult` but expose `text` only concretely — hence `textOf`).
    /// Follows the documented pattern: start a collector over `results`, feed the
    /// file, finalize.
    private func analyze<Results>(
        module: any SpeechModule,
        results: Results,
        fileURL: URL,
        textOf: @escaping @Sendable (Results.Element) -> AttributedString
    ) async throws -> AttributedString
    where Results: AsyncSequence & Sendable,
          Results.Element: SpeechModuleResult & Sendable,
          Results.Failure == any Error {
        let analyzer = SpeechAnalyzer(
            modules: [module],
            options: SpeechAnalyzer.Options(priority: .utility, modelRetention: .whileInUse)
        )
        let file = try AVAudioFile(forReading: fileURL)

        // Collect finalized results concurrently; `results` terminates when the
        // analyzer finishes, ending the loop.
        let collector = Task {
            var full = AttributedString()
            for try await result in results where result.isFinal {
                full.append(textOf(result))
            }
            return full
        }

        do {
            if let lastSample = try await analyzer.analyzeSequence(from: file) {
                try await analyzer.finalizeAndFinish(through: lastSample)
            } else {
                await analyzer.cancelAndFinishNow()
            }
        } catch {
            collector.cancel()
            throw error
        }
        return try await collector.value
    }

    enum TranscriptionError: Error {
        case audioUnreadable
        case noInstalledLanguage
    }
}