markbaseengine/Sources/MarkBaseServer/ModelFinder.swift

import Foundation

// ─────────────────────────────────────────────────────────────
// Model Finder - Search and Select Models
// ─────────────────────────────────────────────────────────────

/// Model search result
public struct ModelSearchResult: Codable, Sendable {
    public let id: String
    public let modelId: String
    public let author: String
    public let sha: String?
    public let lastModified: String?
    public let isPrivate: Bool
    public let disabled: Bool
    public let gated: String?
    public let pipelineTag: String?
    public let tags: [String]
    public let downloads: Int
    public let likes: Int
    public let libraryName: String?
    public let createdAt: String?
    
    public var displayName: String {
        modelId.split(separator: "/").last.map(String.init) ?? modelId
    }
    
    public var isGGUF: Bool {
        tags.contains { $0.lowercased().contains("gguf") }
    }
    
    public var isSafetensors: Bool {
        tags.contains { $0.lowercased().contains("safetensors") }
    }
}

/// Model search filters
public struct ModelFilters {
    public var search: String?
    public var author: String?
    public var library: String?
    public var task: String?
    public var tags: [String]?
    public var sort: String?
    public var direction: Int? // -1 for descending, 1 for ascending
    public var limit: Int?
    
    public init(
        search: String? = nil,
        author: String? = nil,
        library: String? = nil,
        task: String? = nil,
        tags: [String]? = nil,
        sort: String? = nil,
        direction: Int? = nil,
        limit: Int? = nil
    ) {
        self.search = search
        self.author = author
        self.library = library
        self.task = task
        self.tags = tags
        self.sort = sort
        self.direction = direction
        self.limit = limit
    }
    
    /// Convert to query parameters
    public func toQueryItems() -> [URLQueryItem] {
        var items: [URLQueryItem] = []
        
        if let search = search {
            items.append(URLQueryItem(name: "search", value: search))
        }
        if let author = author {
            items.append(URLQueryItem(name: "author", value: author))
        }
        if let library = library {
            items.append(URLQueryItem(name: "filter", value: library))
        }
        if let task = task {
            items.append(URLQueryItem(name: "pipeline_tag", value: task))
        }
        if let tags = tags {
            for tag in tags {
                items.append(URLQueryItem(name: "filter", value: tag))
            }
        }
        if let sort = sort {
            items.append(URLQueryItem(name: "sort", value: sort))
        }
        if let direction = direction {
            items.append(URLQueryItem(name: "direction", value: "\(direction)"))
        }
        if let limit = limit {
            items.append(URLQueryItem(name: "limit", value: "\(limit)"))
        }
        
        return items
    }
}

/// Model finder
public final class ModelFinder {
    public nonisolated(unsafe) static let shared = ModelFinder()
    
    private let session: URLSession
    
    private init() {
        let config = URLSessionConfiguration.default
        config.timeoutIntervalForRequest = 30
        self.session = URLSession(configuration: config)
    }
    
    /// Search models on HuggingFace
    public func searchModels(filters: ModelFilters) async throws -> [ModelSearchResult] {
        var components = URLComponents(string: "https://huggingface.co/api/models")!
        components.queryItems = filters.toQueryItems()
        
        guard let url = components.url else {
            throw ModelFinderError.invalidURL
        }
        
        let (data, _) = try await session.data(from: url)
        let models = try JSONDecoder().decode([ModelSearchResult].self, from: data)
        
        return models
    }
    
    /// Search by name
    public func searchByName(_ name: String, limit: Int = 20) async throws -> [ModelSearchResult] {
        let filters = ModelFilters(
            search: name,
            sort: "downloads",
            direction: -1,
            limit: limit
        )
        return try await searchModels(filters: filters)
    }
    
    /// Search GGUF models
    public func searchGGUF(
        name: String? = nil,
        limit: Int = 20,
        sortBy: String = "downloads"
    ) async throws -> [ModelSearchResult] {
        let filters = ModelFilters(
            search: name,
            tags: ["gguf"],
            sort: sortBy,
            direction: -1,
            limit: limit
        )
        return try await searchModels(filters: filters)
    }
    
    /// Search Safetensors models
    public func searchSafetensors(
        name: String? = nil,
        limit: Int = 20,
        sortBy: String = "downloads"
    ) async throws -> [ModelSearchResult] {
        let filters = ModelFilters(
            search: name,
            tags: ["safetensors"],
            sort: sortBy,
            direction: -1,
            limit: limit
        )
        return try await searchModels(filters: filters)
    }
    
    /// Get model details
    public func getModelDetails(repoId: String) async throws -> ModelSearchResult {
        let url = URL(string: "https://huggingface.co/api/models/\(repoId)")!
        let (data, _) = try await session.data(from: url)
        return try JSONDecoder().decode(ModelSearchResult.self, from: data)
    }
    
    /// List model files
    public func listModelFiles(repoId: String, revision: String = "main") 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 }
    }
}

/// Model finder errors
public enum ModelFinderError: Error, LocalizedError {
    case invalidURL
    case requestFailed(String)
    case decodeFailed(String)
    
    public var errorDescription: String? {
        switch self {
        case .invalidURL:
            return "Invalid URL"
        case .requestFailed(let detail):
            return "Request failed: \(detail)"
        case .decodeFailed(let detail):
            return "Failed to decode response: \(detail)"
        }
    }
}