to-live-photo/Sources/LivePhotoCore/AIEnhancer/ODRManager.swift

//
//  ODRManager.swift
//  LivePhotoCore
//
//  On-Demand Resources manager for AI model download.
//

import Foundation
import os

// MARK: - Download State

/// Model download state
public enum ModelDownloadState: Sendable, Equatable {
    case notDownloaded
    case downloading(progress: Double)
    case downloaded
    case failed(String)
    
    public static func == (lhs: ModelDownloadState, rhs: ModelDownloadState) -> Bool {
        switch (lhs, rhs) {
        case (.notDownloaded, .notDownloaded): return true
        case (.downloaded, .downloaded): return true
        case let (.downloading(p1), .downloading(p2)): return p1 == p2
        case let (.failed(e1), .failed(e2)): return e1 == e2
        default: return false
        }
    }
}

// MARK: - ODR Manager

/// On-Demand Resources manager for AI model
public actor ODRManager {
    public static let shared = ODRManager()
    
    private static let modelTag = "ai-model"
    private static let modelName = "RealESRGAN_x4plus"
    
    private var resourceRequest: NSBundleResourceRequest?
    private var cachedModelURL: URL?
    private let logger = Logger(subsystem: "LivePhotoCore", category: "ODRManager")
    
    private init() {}
    
    // MARK: - Public API
    
    /// Check if model is available locally (either in ODR cache or bundle)
    public func isModelAvailable() async -> Bool {
        // First check if we have a cached URL
        if let url = cachedModelURL, FileManager.default.fileExists(atPath: url.path) {
            return true
        }
        
        // Check bundle (development/fallback)
        if getBundleModelURL() != nil {
            return true
        }
        
        // Check ODR conditionally (only available in app context)
        return await checkODRAvailability()
    }
    
    /// Get current download state
    public func getDownloadState() async -> ModelDownloadState {
        if await isModelAvailable() {
            return .downloaded
        }
        
        if resourceRequest != nil {
            return .downloading(progress: 0)
        }
        
        return .notDownloaded
    }
    
    /// Download model with progress callback
    /// - Parameter progress: Progress callback (0.0 to 1.0)
    public func downloadModel(progress: @escaping @Sendable (Double) -> Void) async throws {
        // Check if already available
        if await isModelAvailable() {
            logger.info("Model already available, skipping download")
            progress(1.0)
            return
        }
        
        logger.info("Starting ODR download for model: \(Self.modelTag)")
        
        // Create resource request
        let request = NSBundleResourceRequest(tags: [Self.modelTag])
        self.resourceRequest = request
        
        // Set up progress observation
        let observation = request.progress.observe(\.fractionCompleted) { progressObj, _ in
            Task { @MainActor in
                progress(progressObj.fractionCompleted)
            }
        }
        
        defer {
            observation.invalidate()
        }
        
        do {
            // Begin accessing resources
            try await request.beginAccessingResources()
            
            logger.info("ODR download completed successfully")
            
            // Find and cache the model URL
            if let url = findModelInBundle(request.bundle) {
                cachedModelURL = url
                logger.info("Model cached at: \(url.path)")
            }
            
            progress(1.0)
        } catch {
            logger.error("ODR download failed: \(error.localizedDescription)")
            self.resourceRequest = nil
            throw AIEnhanceError.modelLoadFailed("Download failed: \(error.localizedDescription)")
        }
    }
    
    /// Get model URL (after download or from bundle)
    public func getModelURL() -> URL? {
        // Return cached URL if available
        if let url = cachedModelURL {
            return url
        }
        
        // Check bundle fallback
        if let url = getBundleModelURL() {
            return url
        }
        
        // Try to find in ODR bundle
        if let request = resourceRequest, let url = findModelInBundle(request.bundle) {
            cachedModelURL = url
            return url
        }
        
        return nil
    }
    
    /// Release ODR resources when not in use
    public func releaseResources() {
        resourceRequest?.endAccessingResources()
        resourceRequest = nil
        cachedModelURL = nil
        logger.info("ODR resources released")
    }
    
    // MARK: - Private Helpers
    
    private func checkODRAvailability() async -> Bool {
        // Use conditionallyBeginAccessingResources to check without triggering download
        let request = NSBundleResourceRequest(tags: [Self.modelTag])

        return await withCheckedContinuation { continuation in
            request.conditionallyBeginAccessingResources { [request] available in
                // Capture request explicitly to prevent ARC from releasing it
                // before the callback fires
                _ = request
                if available {
                    self.logger.debug("ODR model is available locally")
                }
                continuation.resume(returning: available)
            }
        }
    }
    
    private func getBundleModelURL() -> URL? {
        // Try main bundle first
        if let url = Bundle.main.url(forResource: Self.modelName, withExtension: "mlmodelc") {
            return url
        }
        if let url = Bundle.main.url(forResource: Self.modelName, withExtension: "mlpackage") {
            return url
        }
        
        // Try SPM bundle (development)
        #if SWIFT_PACKAGE
        if let url = Bundle.module.url(forResource: Self.modelName, withExtension: "mlmodelc") {
            return url
        }
        if let url = Bundle.module.url(forResource: Self.modelName, withExtension: "mlpackage") {
            return url
        }
        #endif
        
        return nil
    }
    
    private func findModelInBundle(_ bundle: Bundle) -> URL? {
        if let url = bundle.url(forResource: Self.modelName, withExtension: "mlmodelc") {
            return url
        }
        if let url = bundle.url(forResource: Self.modelName, withExtension: "mlpackage") {
            return url
        }
        return nil
    }
}