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.
@Observable
@MainActor
final class TranscriptionSettings {
    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 `SpeechTranscriber` can transcribe (36+), loaded once by
    /// `prepare()`. Empty until then.
    private(set) var supportedLocales: [Locale] = []

    /// 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
        supportedLocales = await SpeechTranscriber.supportedLocales
            .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: [Self.transcriber(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 = Self.transcriber(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: [Self.transcriber(for: locale)]) == .installed {
            installed.append(locale)
        }
        return installed
    }

    // MARK: - Helpers

    /// A minimal transcriber for a locale, used purely for asset queries. The
    /// attribute options don't affect asset status, so they're omitted here.
    static func transcriber(for locale: Locale) -> SpeechTranscriber {
        SpeechTranscriber(
            locale: locale,
            transcriptionOptions: [],
            reportingOptions: [],
            attributeOptions: []
        )
    }

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

    enum SettingsError: Error {
        case atCapacity
    }
}