attendi/Packages/Features/Sources/Recording/Services/AssetPreinstalling.swift

import Foundation
import OSLog
import Speech

/// The preinstalling service that downloads speech model assets through the system's asset inventory — the sole owner of the
/// asset inventory knowledge in the feature.
///
/// The service offers two installation paths over the same core. ``preinstall(for:)`` installs the assets for the supported equivalent
/// of a given locale ahead of its first transcription, logging and emitting its outcomes through ``events`` rather than throwing.
/// ``install(for:onDownloadStart:)`` installs the assets for an already supported locale and throws on failure, for callers that
/// need the assets right away — like the transcribing service, which uses it as its fallback when a transcription starts.
///
/// Either path releases the locale reservations of any other locales beforehand: the feature transcribes a single locale at a time,
/// and the system only permits a limited number of reservations per app — exceeding it would make the installation request throw.
public final class AssetPreinstalling: Preinstalling {

    // MARK: Properties

    /// The stream of events the service emits while preinstalling.
    public let events: AsyncStream<PreinstallingEvent>

    /// The continuation that feeds ``events``.
    private let continuation: AsyncStream<PreinstallingEvent>.Continuation

    // MARK: Initializers

    /// Creates an asset preinstalling service.
    public init() {
        (events, continuation) = AsyncStream.makeStream(of: PreinstallingEvent.self)
    }

    deinit {
        continuation.finish()
    }

    // MARK: Methods

    /// Installs the speech model assets for the given supported locale, waiting for any download to finish.
    ///
    /// The locale reservations of any other locales are released beforehand. Unlike ``preinstall(for:)``, failures are thrown rather
    /// than emitted, so a caller that needs the assets right away — like the transcribing service — can react to them directly.
    ///
    /// - Parameters:
    ///   - locale: The supported locale to install the speech model assets for.
    ///   - onDownloadStart: The closure invoked when the assets are missing and their download starts. Defaults to `nil`.
    /// - Throws: Any error thrown while reserving, downloading, or installing the assets, including a `CancellationError`
    /// when the surrounding task is cancelled mid-download.
    public func install(
        for locale: Locale,
        onDownloadStart: (@Sendable () -> Void)? = nil
    ) async throws {
        for reserved in await AssetInventory.reservedLocales where reserved != locale {
            _ = await AssetInventory.release(reservedLocale: reserved)
        }

        let transcriber = SpeechTranscriber(
            locale: locale,
            preset: .transcription
        )

        guard let request = try await AssetInventory.assetInstallationRequest(
            supporting: [transcriber]
        ) else {
            return
        }

        onDownloadStart?()

        try await withTaskCancellationHandler {
            try await request.downloadAndInstall()
        } onCancel: { [progress = request.progress] in
            progress.cancel()
        }
    }

    /// Preinstalls the speech model assets for the supported equivalent of the given locale, emitting the started, cancelled, and
    /// failed downloads through ``events`` — a download that finishes, or assets that are already installed, emit nothing. The
    /// emitted events carry the given locale, not its supported equivalent, so a caller can match them against its requests.
    ///
    /// - Parameter locale: The locale to preinstall the speech model assets for.
    public func preinstall(
        for locale: Locale
    ) async {
        guard let supported = await supportedLocale(equivalentTo: locale) else {
            return
        }

        do {
            try await install(for: supported) { [continuation] in
                continuation.yield(.started(locale))
            }
        } catch {
            if error is CancellationError || Task.isCancelled {
                continuation.yield(.cancelled(locale))
            } else {
                continuation.yield(.failed(locale))

                logger.error("The speech model assets for the \"\(supported.identifier, privacy: .public)\" locale failed to preinstall: \(String(describing: error), privacy: .public)")
            }
        }
    }

    /// Returns the supported locale equivalent to the given locale.
    ///
    /// - Parameter locale: The locale to find a supported equivalent for.
    /// - Returns: The supported equivalent, or `nil` when the transcriber supports none.
    public func supportedLocale(
        equivalentTo locale: Locale
    ) async -> Locale? {
        await SpeechTranscriber.supportedLocale(equivalentTo: locale)
    }

    /// Returns the locales the transcriber supports.
    ///
    /// - Returns: The supported locales, including ones whose assets are not installed yet.
    public func supportedLocales() async -> [Locale] {
        await SpeechTranscriber.supportedLocales
    }

}

// MARK: - Constants

/// The logger that records the failures of the asset preinstalling service.
private let logger = Logger(
    subsystem: "Features.Recording",
    category: "AssetPreinstalling"
)