notes/Notes/Services/TranscriptionSettings.swift

import Foundation
import Speech

/// The user's transcription-language configuration and the Speech asset state
/// behind it. `@Observable @MainActor` so the Settings UI binds to it directly.
///
/// There is no automatic language detection in `SpeechAnalyzer`, so the user
/// enables a small set of languages; a new voice note is sampled in each and the
/// highest-confidence one wins (see `TranscriptionService`). Enabling a language
/// *reserves* it with `AssetInventory` — which is both what keeps its model from
/// being purged and what the system caps at `AssetInventory.maximumReservedLocales`
/// — so "enabled" and "reserved" are kept in lock-step and the UI enforces the cap.
///
/// The supported set is the *union* of two on-device SpeechAnalyzer modules'
/// coverage: `SpeechTranscriber` (30 locales) and `DictationTranscriber` (54,
/// including ru-RU, uk-UA, pl-PL … that `SpeechTranscriber` lacks). Each locale
/// is routed to whichever engine supports it, preferring `SpeechTranscriber`
/// when both do (see `engine(for:)`). Reservation is locale-based and engine-
/// agnostic; only asset queries and analyzer construction pick a module type.
@Observable
@MainActor
final class TranscriptionSettings {
    /// Which SpeechAnalyzer module transcribes a given locale. `SpeechTranscriber`
    /// is preferred where available (it reports per-run confidence the sampling
    /// pass relies on); `DictationTranscriber` covers the locales it doesn't.
    enum EngineKind: Equatable {
        case speech
        case dictation
    }

    private static let enabledKey = "Notes.enabledTranscriptionLocales"
    private static let seededKey = "Notes.transcriptionLocalesSeeded"

    /// Per-install id of the transcription owner, reused from `InstallIdentity`
    /// — a voice note stamps this into `AudioInfo.transcriberDeviceID` so only the
    /// recording device transcribes it.
    static var installID: String { InstallIdentity.installID }

    /// Canonical BCP-47 ids of the enabled languages, persisted in UserDefaults.
    private(set) var enabledLocaleIDs: [String]

    /// Every locale either engine can transcribe — the union of `SpeechTranscriber`
    /// and `DictationTranscriber` coverage, de-duplicated by canonical id and
    /// loaded once by `prepare()`. Empty until then.
    private(set) var supportedLocales: [Locale] = []

    /// Canonical ids `SpeechTranscriber` supports, for engine resolution. A locale
    /// in this set is routed to `.speech`; anything else supported falls to
    /// `.dictation`. Populated by `prepare()`.
    private var speechLocaleIDs: Set<String> = []

    /// Canonical ids `DictationTranscriber` supports, for engine resolution and
    /// for distinguishing a genuinely unsupported locale from a dictation-only one.
    private var dictationLocaleIDs: Set<String> = []

    /// Runtime cap on how many languages can be enabled at once — read from the
    /// system, not hard-coded, because it can change.
    private(set) var maximumReserved = 0

    /// Per-language asset state, keyed by canonical BCP-47 id, for the row UI.
    private(set) var assetStates: [String: AssetState] = [:]

    /// Where a language's on-device model sits.
    enum AssetState: Equatable {
        case unsupported
        case notInstalled
        /// Downloading; the fraction is nil until progress is known.
        case downloading(Double?)
        case installed
    }

    private let defaults: UserDefaults

    init(defaults: UserDefaults = .standard) {
        self.defaults = defaults
        self.enabledLocaleIDs = defaults.stringArray(forKey: Self.enabledKey) ?? []
    }

    // MARK: - Canonical ids

    /// The stable key for a locale everywhere in the app: its BCP-47 identifier
    /// (`en-US`), so it round-trips through `Locale(identifier:)`.
    static func canonicalID(_ locale: Locale) -> String { locale.identifier(.bcp47) }

    /// The enabled languages as `Locale`s, in the order the user has them.
    var enabledLocales: [Locale] { enabledLocaleIDs.map(Locale.init(identifier:)) }

    /// Localized display name for a locale ("English (United States)"), falling
    /// back to the raw identifier.
    static func displayName(for locale: Locale) -> String {
        Locale.current.localizedString(forIdentifier: locale.identifier(.bcp47))
            ?? locale.identifier(.bcp47)
    }

    // MARK: - Loading

    /// One-time load of the supported-locale list and the reservation cap, plus a
    /// first-run seed of the enabled set from the user's preferred languages. Idempotent
    /// and cheap to call again — it re-fetches the list but only seeds once.
    func prepare() async {
        maximumReserved = AssetInventory.maximumReservedLocales

        // Query both engines concurrently; the supported set is their union.
        async let speech = SpeechTranscriber.supportedLocales
        async let dictation = DictationTranscriber.supportedLocales
        let speechList = await speech
        let dictationList = await dictation

        speechLocaleIDs = Set(speechList.map(Self.canonicalID))
        dictationLocaleIDs = Set(dictationList.map(Self.canonicalID))

        // De-duplicate by canonical id, keeping the `SpeechTranscriber` locale
        // object where both engines cover it (harmless — same bcp47 either way).
        var byID: [String: Locale] = [:]
        for locale in dictationList { byID[Self.canonicalID(locale)] = locale }
        for locale in speechList { byID[Self.canonicalID(locale)] = locale }
        supportedLocales = byID.values
            .sorted { Self.displayName(for: $0) < Self.displayName(for: $1) }

        if !defaults.bool(forKey: Self.seededKey) {
            seedDefaults()
            defaults.set(true, forKey: Self.seededKey)
        }
        await refreshStates()
    }

    /// Seed the enabled set from `Locale.preferredLanguages ∩ supportedLocales`,
    /// matched by language code and capped at the reservation limit. Falls back to
    /// English (or the first supported locale) so a note recorded before the user
    /// visits Settings still has something to transcribe with.
    private func seedDefaults() {
        guard enabledLocaleIDs.isEmpty else { return }
        var seed: [String] = []
        for preferred in Locale.preferredLanguages.map(Locale.init(identifier:)) {
            guard let match = supportedLocales.first(where: {
                $0.language.languageCode == preferred.language.languageCode
            }) else { continue }
            let id = Self.canonicalID(match)
            if !seed.contains(id) { seed.append(id) }
        }
        if seed.isEmpty {
            let fallback = supportedLocales.first { $0.language.languageCode?.identifier == "en" }
                ?? supportedLocales.first
            if let fallback { seed = [Self.canonicalID(fallback)] }
        }
        enabledLocaleIDs = Array(seed.prefix(max(1, maximumReserved)))
        persist()
        // Reserve the seeded languages so their models are protected from purge.
        Task { await self.reserveEnabled() }
    }

    private func reserveEnabled() async {
        for locale in enabledLocales {
            _ = try? await AssetInventory.reserve(locale: locale)
        }
    }

    /// Refresh every supported locale's asset state (installed / downloadable /
    /// unsupported). Runs sequentially — `status(forModules:)` is cheap but there
    /// are dozens of locales, and there's no need to fan out.
    func refreshStates() async {
        for locale in supportedLocales {
            let id = Self.canonicalID(locale)
            // Preserve an in-flight download's fraction rather than stomping it.
            if case .downloading = assetStates[id] { continue }
            assetStates[id] = await currentState(for: locale)
        }
    }

    private func currentState(for locale: Locale) async -> AssetState {
        switch await AssetInventory.status(forModules: [module(for: locale)]) {
        case .unsupported: return .unsupported
        case .supported: return .notInstalled
        case .downloading: return .downloading(nil)
        case .installed: return .installed
        @unknown default: return .notInstalled
        }
    }

    // MARK: - Enable / disable

    var isAtCapacity: Bool { enabledLocaleIDs.count >= maximumReserved }

    func isEnabled(_ locale: Locale) -> Bool {
        enabledLocaleIDs.contains(Self.canonicalID(locale))
    }

    /// Enable a language: reserve it, mark it enabled, and download its model if
    /// needed (surfacing progress through `assetStates`). Throws `atCapacity` when
    /// already at the reservation cap.
    func enable(_ locale: Locale) async throws {
        let id = Self.canonicalID(locale)
        guard !enabledLocaleIDs.contains(id) else { return }
        guard !isAtCapacity else { throw SettingsError.atCapacity }

        _ = try await AssetInventory.reserve(locale: locale)
        enabledLocaleIDs.append(id)
        persist()
        try await installIfNeeded(locale)
    }

    /// Disable a language: drop it from the enabled set and release its
    /// reservation (letting the system reclaim the model later). Best-effort.
    func disable(_ locale: Locale) async {
        let id = Self.canonicalID(locale)
        enabledLocaleIDs.removeAll { $0 == id }
        persist()
        _ = await AssetInventory.release(reservedLocale: locale)
        assetStates[id] = await currentState(for: locale)
    }

    private func installIfNeeded(_ locale: Locale) async throws {
        let id = Self.canonicalID(locale)
        let transcriber = module(for: locale)
        guard await AssetInventory.status(forModules: [transcriber]) != .installed else {
            assetStates[id] = .installed
            return
        }
        guard let request = try await AssetInventory.assetInstallationRequest(supporting: [transcriber]) else {
            assetStates[id] = await currentState(for: locale)
            return
        }
        let progress = request.progress
        assetStates[id] = .downloading(0)
        // Poll the Progress so the @Observable dict updates the row; Progress KVO
        // doesn't drive SwiftUI on its own.
        let poll = Task { [weak self] in
            while !Task.isCancelled {
                try? await Task.sleep(for: .milliseconds(300))
                guard let self else { return }
                self.assetStates[id] = .downloading(progress.fractionCompleted)
                if progress.isFinished { return }
            }
        }
        defer { poll.cancel() }
        try await request.downloadAndInstall()
        assetStates[id] = await currentState(for: locale)
    }

    /// The subset of enabled languages whose models are installed — the only ones
    /// `TranscriptionService` can actually run.
    func enabledInstalledLocales() async -> [Locale] {
        var installed: [Locale] = []
        for locale in enabledLocales
        where await AssetInventory.status(forModules: [module(for: locale)]) == .installed {
            installed.append(locale)
        }
        return installed
    }

    // MARK: - Engine routing

    /// Which module transcribes `locale`: `.dictation` only for a locale the
    /// dictation engine covers and `SpeechTranscriber` does not; `.speech`
    /// otherwise. This both prefers `SpeechTranscriber` when both cover a locale
    /// and defaults to it before `prepare()` has populated the sets (empty sets ⇒
    /// `.speech` everywhere, the historical behavior) — so a transcription that
    /// races startup still routes sanely rather than sending English to dictation.
    /// A truly unsupported locale reports `.unsupported` under either engine.
    func engine(for locale: Locale) -> EngineKind {
        let id = Self.canonicalID(locale)
        if dictationLocaleIDs.contains(id), !speechLocaleIDs.contains(id) { return .dictation }
        return .speech
    }

    /// The analyzer module for a locale, built with its resolved engine. Used for
    /// asset queries (confidence omitted, which doesn't affect asset status) and,
    /// with `requestingConfidence: true`, for the transcription passes.
    func module(for locale: Locale, requestingConfidence: Bool = false) -> any SpeechModule {
        Self.module(for: locale, engine: engine(for: locale), requestingConfidence: requestingConfidence)
    }

    /// A locale's analyzer module for a given engine. Pure, so `TranscriptionService`
    /// can build the same module the settings screen queries. `requestingConfidence`
    /// asks the engine for the `.transcriptionConfidence` attribute — honored where
    /// the engine emits it (`SpeechTranscriber` always; `DictationTranscriber` may
    /// leave runs unattributed, which `LanguagePicker` degrades gracefully).
    static func module(
        for locale: Locale,
        engine: EngineKind,
        requestingConfidence: Bool = false
    ) -> any SpeechModule {
        switch engine {
        case .speech:
            return SpeechTranscriber(
                locale: locale,
                transcriptionOptions: [],
                reportingOptions: [],
                attributeOptions: requestingConfidence ? [.transcriptionConfidence] : []
            )
        case .dictation:
            return DictationTranscriber(
                locale: locale,
                contentHints: [],
                transcriptionOptions: [],
                reportingOptions: [],
                attributeOptions: requestingConfidence ? [.transcriptionConfidence] : []
            )
        }
    }

    private func persist() {
        defaults.set(enabledLocaleIDs, forKey: Self.enabledKey)
    }

    enum SettingsError: Error {
        case atCapacity
    }
}