markbaseengine/Sources/MarkBaseServer/ModelDownloader.swift

import Foundation

// ─────────────────────────────────────────────────────────────
// Model Downloader - Download models from HuggingFace
// ─────────────────────────────────────────────────────────────

/// Download progress
public struct DownloadProgress: Sendable {
    public let current: Int64
    public let total: Int64
    public let percentage: Double
    public let speed: Double // bytes per second
    public let eta: TimeInterval // seconds
    
    public init(current: Int64, total: Int64, speed: Double = 0) {
        self.current = current
        self.total = total
        self.percentage = total > 0 ? Double(current) / Double(total) : 0
        self.speed = speed
        self.eta = speed > 0 ? Double(total - current) / speed : 0
    }
}

/// Model downloader
public final class ModelDownloader: @unchecked Sendable {
    public static let shared = ModelDownloader()
    
    private let session: URLSession
    private var progressHandler: ((DownloadProgress) -> Void)?
    
    private init() {
        let config = URLSessionConfiguration.default
        config.timeoutIntervalForRequest = 300
        config.timeoutIntervalForResource = 3600
        self.session = URLSession(configuration: config)
    }
    
    /// Set progress handler
    public func onProgress(_ handler: @escaping (DownloadProgress) -> Void) {
        self.progressHandler = handler
    }
    
    /// Download model from HuggingFace
    public func downloadModel(
        repoId: String,
        to destinationDir: String,
        revision: String = "main"
    ) async throws {
        let destinationURL = URL(fileURLWithPath: destinationDir)
        
        // Create destination directory
        try FileManager.default.createDirectory(
            at: destinationURL,
            withIntermediateDirectories: true
        )
        
        // List files in repo
        let files = try await listRepoFiles(repoId: repoId, revision: revision)
        
        // Download each file
        for file in files {
            let fileURL = destinationURL.appendingPathComponent(file)
            
            // Skip if already exists
            if FileManager.default.fileExists(atPath: fileURL.path) {
                print("Skipping \(file) (already exists)")
                continue
            }
            
            print("Downloading \(file)...")
            try await downloadFile(
                repoId: repoId,
                file: file,
                revision: revision,
                to: fileURL
            )
        }
        
        print("✓ Model downloaded to \(destinationDir)")
    }
    
    /// List files in repo
    private func listRepoFiles(repoId: String, revision: String) async throws -> [String] {
        let url = URL(string: "https://huggingface.co/api/models/\(repoId)/tree/\(revision)")!
        let (data, _) = try await session.data(from: url)
        
        struct FileInfo: Codable {
            let path: String
            let type: String
        }
        
        let files = try JSONDecoder().decode([FileInfo].self, from: data)
        return files.filter { $0.type == "file" }.map { $0.path }
    }
    
    /// Download single file
    private func downloadFile(
        repoId: String,
        file: String,
        revision: String,
        to destination: URL
    ) async throws {
        let downloadURL = URL(
            string: "https://huggingface.co/\(repoId)/resolve/\(revision)/\(file)"
        )!
        
        let (tempURL, response) = try await session.download(from: downloadURL)
        
        // Get total size
        let total = response.expectedContentLength
        
        // Move to destination
        try FileManager.default.moveItem(at: tempURL, to: destination)
    }
    
    /// Download with progress
    private func downloadWithProgress(
        from url: URL,
        to destination: URL
    ) async throws {
        var request = URLRequest(url: url)
        request.timeoutInterval = 300
        
        let (bytes, response) = try await session.bytes(for: request)
        let total = response.expectedContentLength
        
        var current: Int64 = 0
        var startTime = Date()
        
        // Create file
        let fileHandle = try FileHandle(forWritingTo: destination)
        defer { try? fileHandle.close() }
        
        for try await byte in bytes {
            try fileHandle.write(contentsOf: [byte])
            current += 1
            
            // Update progress every 1MB
            if current % (1024 * 1024) == 0 {
                let elapsed = Date().timeIntervalSince(startTime)
                let speed = elapsed > 0 ? Double(current) / elapsed : 0
                
                let progress = DownloadProgress(
                    current: current,
                    total: total,
                    speed: speed
                )
                
                progressHandler?(progress)
            }
        }
    }
}

/// Model repository
public struct ModelRepository {
    public let id: String
    public let name: String
    public let description: String
    public let downloads: Int
    public let likes: Int
    
    public init(
        id: String,
        name: String,
        description: String,
        downloads: Int,
        likes: Int
    ) {
        self.id = id
        self.name = name
        self.description = description
        self.downloads = downloads
        self.likes = likes
    }
}