sleepy_agent_ios/SleepyAgent/Services/ModelDownloadService.swift

import Foundation

// MARK: - Errors

enum ModelDownloadError: LocalizedError {
    case invalidURL
    case networkError(Error)
    case serverError(Int)
    case insufficientStorage
    case fileError(Error)
    case invalidResponse
    case checksumMismatch
    case cancelled
    
    var errorDescription: String? {
        switch self {
        case .invalidURL:
            return "Invalid download URL"
        case .networkError(let error):
            return "Network error: \(error.localizedDescription)"
        case .serverError(let code):
            return "Server error: HTTP \(code)"
        case .insufficientStorage:
            return "Insufficient storage space"
        case .fileError(let error):
            return "File error: \(error.localizedDescription)"
        case .invalidResponse:
            return "Invalid server response"
        case .checksumMismatch:
            return "Downloaded file checksum mismatch"
        case .cancelled:
            return "Download cancelled"
        }
    }
}

// MARK: - Progress & State

enum DownloadState: Equatable {
    case idle
    case checking
    case downloading(progress: Double, bytesDownloaded: Int64, totalBytes: Int64)
    case paused(bytesDownloaded: Int64, totalBytes: Int64)
    case completed(url: URL)
    case error(ModelDownloadError)
    
    static func == (lhs: DownloadState, rhs: DownloadState) -> Bool {
        switch (lhs, rhs) {
        case (.idle, .idle), (.checking, .checking):
            return true
        case let (.downloading(p1, b1, t1), .downloading(p2, b2, t2)):
            return abs(p1 - p2) < 0.001 && b1 == b2 && t1 == t2
        case let (.paused(b1, t1), .paused(b2, t2)):
            return b1 == b2 && t1 == t2
        case let (.completed(u1), .completed(u2)):
            return u1 == u2
        case let (.error(e1), .error(e2)):
            return e1.localizedDescription == e2.localizedDescription
        default:
            return false
        }
    }
}

// MARK: - Model Info

struct ModelInfo {
    let url: URL
    let filename: String
    let expectedSize: Int64?
    let variant: ModelVariant
    
    enum ModelVariant: String, CaseIterable {
        case e2b = "gemma-4-E2B-it"
        case e4b = "gemma-4-E4B-it"
        
        var filename: String {
            "\(rawValue).litertlm"
        }
        
        var huggingFaceURL: URL {
            URL(string: "https://huggingface.co/litert-community/\(rawValue)-litert-lm/resolve/main/\(filename)")!
        }
        
        var expectedSize: Int64 {
            switch self {
            case .e2b:
                return 2_717_263_232 // ~2.53 GB
            case .e4b:
                return 4_831_838_208 // ~4.5 GB
            }
        }
    }
    
    static func e2b() -> ModelInfo {
        ModelInfo(
            url: ModelVariant.e2b.huggingFaceURL,
            filename: ModelVariant.e2b.filename,
            expectedSize: ModelVariant.e2b.expectedSize,
            variant: .e2b
        )
    }
    
    static func e4b() -> ModelInfo {
        ModelInfo(
            url: ModelVariant.e4b.huggingFaceURL,
            filename: ModelVariant.e4b.filename,
            expectedSize: ModelVariant.e4b.expectedSize,
            variant: .e4b
        )
    }
}

// MARK: - Download Task

actor ModelDownloadService {
    
    // MARK: - Properties
    
    private let fileManager: FileManager
    private let urlSession: URLSession
    private var currentTask: Task<Void, Never>?
    private var downloadContinuation: AsyncStream<DownloadState>.Continuation?
    
    private var currentState: DownloadState = .idle
    private var tempFileURL: URL?
    private var destinationURL: URL?
    
    /// Current download state
    var state: DownloadState { currentState }
    
    /// Path to the models directory
    let modelsDirectory: URL
    
    // MARK: - Constants
    
    private let chunkSize = 8192 // 8KB chunks
    private let progressUpdateInterval: TimeInterval = 0.5
    
    // MARK: - Initialization
    
    init(
        fileManager: FileManager = .default,
        urlSession: URLSession = .shared
    ) {
        self.fileManager = fileManager
        self.urlSession = urlSession
        
        // Setup models directory in documents
        let documentsDirectory = fileManager.urls(for: .documentDirectory, in: .userDomainMask).first!
        self.modelsDirectory = documentsDirectory.appendingPathComponent("models", isDirectory: true)
        
        // Create models directory if needed
        try? fileManager.createDirectory(at: modelsDirectory, withIntermediateDirectories: true)
    }
    
    // MARK: - Public Methods
    
    /// Checks if a model variant is downloaded.
    func isDownloaded(variant: ModelInfo.ModelVariant) -> Bool {
        let fileURL = modelsDirectory.appendingPathComponent(variant.filename)
        guard fileManager.fileExists(atPath: fileURL.path) else { return false }
        
        let minSize: Int64 = 100_000_000 // At least 100MB
        do {
            let attributes = try fileManager.attributesOfItem(atPath: fileURL.path)
            let fileSize = attributes[.size] as? Int64 ?? 0
            return fileSize > minSize
        } catch {
            return false
        }
    }
    
    /// Gets the URL for a model file.
    func modelFileURL(variant: ModelInfo.ModelVariant) -> URL {
        modelsDirectory.appendingPathComponent(variant.filename)
    }
    
    /// Gets download progress for a variant.
    func downloadProgress(variant: ModelInfo.ModelVariant) -> Double {
        let fileURL = modelsDirectory.appendingPathComponent(variant.filename)
        guard fileManager.fileExists(atPath: fileURL.path) else { return 0 }
        
        do {
            let attributes = try fileManager.attributesOfItem(atPath: fileURL.path)
            let fileSize = attributes[.size] as? Int64 ?? 0
            return min(Double(fileSize) / Double(variant.expectedSize), 1.0)
        } catch {
            return 0
        }
    }
    
    /// Downloads a model with progress reporting.
    /// - Parameter modelInfo: The model to download
    /// - Returns: AsyncStream of download state updates
    func download(modelInfo: ModelInfo) -> AsyncStream<DownloadState> {
        // Cancel any existing download
        cancelDownload()
        
        return AsyncStream { continuation in
            self.downloadContinuation = continuation
            
            currentTask = Task {
                await performDownload(modelInfo: modelInfo, continuation: continuation)
            }
            
            continuation.onTermination = { [weak self] _ in
                Task {
                    await self?.cancelDownload()
                }
            }
        }
    }
    
    /// Cancels the current download.
    func cancelDownload() {
        currentTask?.cancel()
        currentTask = nil
        
        // Clean up temp file
        if let tempURL = tempFileURL {
            try? fileManager.removeItem(at: tempURL)
            tempFileURL = nil
        }
        
        updateState(.idle)
    }
    
    /// Deletes a downloaded model.
    func deleteModel(variant: ModelInfo.ModelVariant) throws {
        let fileURL = modelsDirectory.appendingPathComponent(variant.filename)
        if fileManager.fileExists(atPath: fileURL.path) {
            try fileManager.removeItem(at: fileURL)
        }
        
        // Also delete temp file if exists
        let tempURL = modelsDirectory.appendingPathComponent("\(variant.filename).tmp")
        if fileManager.fileExists(atPath: tempURL.path) {
            try fileManager.removeItem(at: tempURL)
        }
    }
    
    /// Gets available storage space in bytes.
    func availableStorage() -> Int64? {
        do {
            let attributes = try fileManager.attributesOfFileSystem(forPath: modelsDirectory.path)
            return attributes[.systemFreeSize] as? Int64
        } catch {
            return nil
        }
    }
    
    /// Formats bytes to human-readable string.
    static func formatBytes(_ bytes: Int64) -> String {
        let formatter = ByteCountFormatter()
        formatter.countStyle = .file
        return formatter.string(fromByteCount: bytes)
    }
    
    // MARK: - Private Methods
    
    private func performDownload(
        modelInfo: ModelInfo,
        continuation: AsyncStream<DownloadState>.Continuation
    ) async {
        updateState(.checking)
        
        destinationURL = modelsDirectory.appendingPathComponent(modelInfo.filename)
        let tempURL = modelsDirectory.appendingPathComponent("\(modelInfo.filename).tmp")
        tempFileURL = tempURL
        
        // Check available storage
        if let available = availableStorage(), let expected = modelInfo.expectedSize {
            if available < expected {
                updateState(.error(.insufficientStorage))
                continuation.finish()
                return
            }
        }
        
        // Check for partial download to resume
        let resumeFrom: Int64
        if fileManager.fileExists(atPath: tempURL.path) {
            do {
                let attributes = try fileManager.attributesOfItem(atPath: tempURL.path)
                resumeFrom = attributes[.size] as? Int64 ?? 0
            } catch {
                resumeFrom = 0
            }
        } else {
            resumeFrom = 0
        }
        
        do {
            try await downloadFile(
                from: modelInfo.url,
                to: tempURL,
                resumeFrom: resumeFrom,
                expectedSize: modelInfo.expectedSize,
                continuation: continuation
            )
            
            // Move temp file to final location
            if fileManager.fileExists(atPath: destinationURL!.path) {
                try fileManager.removeItem(at: destinationURL!)
            }
            try fileManager.moveItem(at: tempURL, to: destinationURL!)
            
            tempFileURL = nil
            updateState(.completed(url: destinationURL!))
            
        } catch is CancellationError {
            updateState(.idle)
        } catch let error as ModelDownloadError {
            updateState(.error(error))
        } catch {
            updateState(.error(.networkError(error)))
        }
        
        continuation.finish()
    }
    
    private func downloadFile(
        from url: URL,
        to destination: URL,
        resumeFrom: Int64,
        expectedSize: Int64?,
        continuation: AsyncStream<DownloadState>.Continuation
    ) async throws {
        var request = URLRequest(url: url)
        request.timeoutInterval = 30
        
        // Set up resume with Range header
        if resumeFrom > 0 {
            request.setValue("bytes=\(resumeFrom)-", forHTTPHeaderField: "Range")
        }
        
        request.setValue("SleepyAgent-iOS/1.0", forHTTPHeaderField: "User-Agent")
        
        let (asyncBytes, response) = try await urlSession.bytes(for: request)
        
        guard let httpResponse = response as? HTTPURLResponse else {
            throw ModelDownloadError.invalidResponse
        }
        
        let validCodes = resumeFrom > 0 ? [200, 206] : [200]
        guard validCodes.contains(httpResponse.statusCode) else {
            throw ModelDownloadError.serverError(httpResponse.statusCode)
        }
        
        // Determine total size
        let totalBytes: Int64
        if let contentLength = httpResponse.value(forHTTPHeaderField: "Content-Length"),
           let length = Int64(contentLength) {
            if httpResponse.statusCode == 206 && resumeFrom > 0 {
                totalBytes = resumeFrom + length
            } else {
                totalBytes = length
            }
        } else if let expected = expectedSize {
            totalBytes = expected
        } else {
            totalBytes = 0
        }
        
        // Open file for writing/appending
        let fileHandle: FileHandle
        if resumeFrom > 0 && fileManager.fileExists(atPath: destination.path) {
            fileHandle = try FileHandle(forWritingTo: destination)
            try fileHandle.seekToEnd()
        } else {
            if fileManager.fileExists(atPath: destination.path) {
                try fileManager.removeItem(at: destination)
            }
            fileManager.createFile(atPath: destination.path, contents: nil)
            fileHandle = try FileHandle(forWritingTo: destination)
        }
        defer {
            try? fileHandle.close()
        }
        
        var bytesDownloaded = resumeFrom
        var lastProgressUpdate = Date()
        
        for try await byte in asyncBytes {
            try Task.checkCancellation()
            
            fileHandle.write(Data([byte]))
            bytesDownloaded += 1
            
            // Update progress periodically
            let now = Date()
            if now.timeIntervalSince(lastProgressUpdate) > progressUpdateInterval {
                let progress = totalBytes > 0 ? Double(bytesDownloaded) / Double(totalBytes) : 0
                updateState(.downloading(
                    progress: progress,
                    bytesDownloaded: bytesDownloaded,
                    totalBytes: totalBytes
                ))
                lastProgressUpdate = now
            }
        }
        
        // Final progress update
        let finalProgress = totalBytes > 0 ? Double(bytesDownloaded) / Double(totalBytes) : 1.0
        updateState(.downloading(
            progress: finalProgress,
            bytesDownloaded: bytesDownloaded,
            totalBytes: totalBytes
        ))
        
        // Verify download completed
        if let expected = expectedSize, bytesDownloaded < expected - 1000 {
            throw ModelDownloadError.networkError(NSError(domain: "Download incomplete", code: -1))
        }
    }
    
    private func updateState(_ state: DownloadState) {
        currentState = state
        downloadContinuation?.yield(state)
    }
}

// MARK: - Background Download Support

extension ModelDownloadService {
    /// Creates a background URLSession configuration for downloads.
    static func backgroundConfiguration(identifier: String) -> URLSessionConfiguration {
        let config = URLSessionConfiguration.background(withIdentifier: identifier)
        config.isDiscretionary = false
        config.sessionSendsLaunchEvents = true
        config.allowsCellularAccess = true
        return config
    }
}