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

//
//  AIEnhancer.swift
//  LivePhotoCore
//
//  AI super-resolution enhancement using Real-ESRGAN Core ML model.
//

import CoreGraphics
import CoreML
import Foundation
import os

// MARK: - Configuration

/// AI enhancement configuration
public struct AIEnhanceConfig: Codable, Sendable, Hashable {
    /// Enable AI super-resolution
    public var enabled: Bool

    public init(enabled: Bool = false) {
        self.enabled = enabled
    }

    /// Disabled configuration
    public static let disabled = AIEnhanceConfig(enabled: false)

    /// Standard configuration
    public static let standard = AIEnhanceConfig(enabled: true)
}

// MARK: - Result

/// AI enhancement result
public struct AIEnhanceResult: Sendable {
    /// Enhanced image
    public let enhancedImage: CGImage

    /// Original image size
    public let originalSize: CGSize

    /// Enhanced image size
    public let enhancedSize: CGSize

    /// Processing time in milliseconds
    public let processingTimeMs: Double
}

// MARK: - Errors

/// AI enhancement error types
public enum AIEnhanceError: Error, Sendable, LocalizedError {
    case modelNotFound
    case modelLoadFailed(String)
    case inputImageInvalid
    case inferenceError(String)
    case memoryPressure
    case cancelled
    case deviceNotSupported

    public var errorDescription: String? {
        switch self {
        case .modelNotFound:
            return "AI model file not found in bundle"
        case let .modelLoadFailed(reason):
            return "Failed to load AI model: \(reason)"
        case .inputImageInvalid:
            return "Input image is invalid or cannot be processed"
        case let .inferenceError(reason):
            return "AI inference failed: \(reason)"
        case .memoryPressure:
            return "Not enough memory for AI processing"
        case .cancelled:
            return "AI enhancement was cancelled"
        case .deviceNotSupported:
            return "Device does not support AI enhancement"
        }
    }
}

// MARK: - Progress

/// Progress callback for AI enhancement
/// - Parameter progress: Value from 0.0 to 1.0
public typealias AIEnhanceProgress = @Sendable (Double) -> Void

// MARK: - Main Actor

/// AI enhancement actor for super-resolution processing
public actor AIEnhancer {
    private let config: AIEnhanceConfig
    private var processor: RealESRGANProcessor?
    private let logger = Logger(subsystem: "LivePhotoCore", category: "AIEnhancer")

    /// Scale factor (4 for Real-ESRGAN x4plus)
    public static let scaleFactor: Int = 4

    /// Initialize with configuration
    public init(config: AIEnhanceConfig = .standard) {
        self.config = config
    }

    // MARK: - Device Capability

    /// Check if AI enhancement is available on this device
    public static func isAvailable() -> Bool {
        // Require iOS 17+
        guard #available(iOS 17.0, *) else {
            return false
        }

        // Check device memory (require at least 4GB)
        let totalMemory = ProcessInfo.processInfo.physicalMemory
        let memoryGB = Double(totalMemory) / (1024 * 1024 * 1024)
        guard memoryGB >= 4.0 else {
            return false
        }

        // Neural Engine is available on A12+ (iPhone XS and later)
        // iOS 17 requirement ensures A12+ is present
        return true
    }
    
    // MARK: - Model Download (ODR)
    
    /// Check if AI model needs to be downloaded
    public static func needsDownload() async -> Bool {
        let available = await ODRManager.shared.isModelAvailable()
        return !available
    }
    
    /// Get current model download state
    public static func getDownloadState() async -> ModelDownloadState {
        await ODRManager.shared.getDownloadState()
    }
    
    /// Download AI model with progress callback
    /// - Parameter progress: Progress callback (0.0 to 1.0)
    public static func downloadModel(progress: @escaping @Sendable (Double) -> Void) async throws {
        try await ODRManager.shared.downloadModel(progress: progress)
    }
    
    /// Release ODR resources when AI enhancement is no longer needed
    public static func releaseModelResources() async {
        await ODRManager.shared.releaseResources()
    }

    // MARK: - Model Management

    /// Preload the model (call during app launch or settings change)
    public func preloadModel() async throws {
        guard AIEnhancer.isAvailable() else {
            throw AIEnhanceError.deviceNotSupported
        }

        guard processor == nil else {
            logger.debug("Model already loaded")
            return
        }

        logger.info("Preloading Real-ESRGAN model...")

        processor = RealESRGANProcessor()
        try await processor?.loadModel()

        logger.info("Model preloaded successfully")
    }

    /// Release model from memory
    public func unloadModel() async {
        await processor?.unloadModel()
        processor = nil
        logger.info("Model unloaded")
    }

    // MARK: - Enhancement

    /// Enhance a single image with AI super-resolution
    /// - Parameters:
    ///   - image: Input CGImage to enhance
    ///   - progress: Optional progress callback (0.0 to 1.0)
    /// - Returns: Enhanced result with metadata
    public func enhance(
        image: CGImage,
        progress: AIEnhanceProgress? = nil
    ) async throws -> AIEnhanceResult {
        guard config.enabled else {
            throw AIEnhanceError.inputImageInvalid
        }

        guard AIEnhancer.isAvailable() else {
            throw AIEnhanceError.deviceNotSupported
        }

        let startTime = CFAbsoluteTimeGetCurrent()
        let originalSize = CGSize(width: image.width, height: image.height)

        logger.info("Starting AI enhancement: \(image.width)x\(image.height)")

        // Ensure model is loaded
        if processor == nil {
            try await preloadModel()
        }

        guard let processor = processor else {
            throw AIEnhanceError.modelNotFound
        }

        // Choose processor based on image size
        // - Small images (≤ 512x512): use WholeImageProcessor (faster, single inference)
        // - Large images (> 512 in either dimension): use TiledImageProcessor (preserves detail)
        let usesTiling = image.width > RealESRGANProcessor.inputSize || image.height > RealESRGANProcessor.inputSize

        let enhancedImage: CGImage
        if usesTiling {
            logger.info("Using tiled processing for large image")
            let tiledProcessor = TiledImageProcessor()
            enhancedImage = try await tiledProcessor.processImage(
                image,
                processor: processor,
                progress: progress
            )
        } else {
            logger.info("Using whole image processing for small image")
            let wholeImageProcessor = WholeImageProcessor()
            enhancedImage = try await wholeImageProcessor.processImage(
                image,
                processor: processor,
                progress: progress
            )
        }

        let processingTime = (CFAbsoluteTimeGetCurrent() - startTime) * 1000
        let enhancedSize = CGSize(width: enhancedImage.width, height: enhancedImage.height)

        logger.info(
            "AI enhancement complete: \(Int(originalSize.width))x\(Int(originalSize.height)) -> \(Int(enhancedSize.width))x\(Int(enhancedSize.height)) in \(Int(processingTime))ms"
        )

        return AIEnhanceResult(
            enhancedImage: enhancedImage,
            originalSize: originalSize,
            enhancedSize: enhancedSize,
            processingTimeMs: processingTime
        )
    }
}