//
//  HubService.swift
//  Zyquo MLX
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//

import Foundation

/// One search result / catalog entry from the Hugging Face Hub.
struct HubModel: Identifiable, Codable, Hashable, Sendable {
    var id: String // repo id, e.g. "mlx-community/Qwen3-4B-4bit"
    var pipelineTag: String?
    var downloads: Int
    var likes: Int
    var tags: [String]
    var gated: Bool

    var name: String { id.split(separator: "/").last.map(String.init) ?? id }

    /// Best-effort type from the pipeline tag (MODELS.md §1.1).
    var modelType: ModelType {
        switch pipelineTag {
        case "image-text-to-text": .vlm
        case "feature-extraction", "sentence-similarity": .embedding
        case "automatic-speech-recognition": .speech
        case "image-to-image", "text-to-image": .imageGeneration
        default: .llm
        }
    }
}

/// A file inside a Hub repo (from the tree endpoint).
struct HubFile: Codable, Sendable {
    var path: String
    var size: Int64
}

enum HubServiceError: LocalizedError {
    case badResponse(Int)
    case gatedRepo(String)

    var errorDescription: String? {
        switch self {
        case .badResponse(let code):
            "Hugging Face returned HTTP \(code). Check your connection (or your token for gated models)."
        case .gatedRepo(let id):
            "\(id) is gated — add a Hugging Face token in Settings to access it."
        }
    }
}

/// Live Hugging Face Hub API client (verified behavior in docs/MODELS.md §1):
/// search, file listing (sizes from the recursive tree), and resolve URLs.
actor HubService {

    static let shared = HubService()

    private let session = URLSession.shared

    /// Search MLX models (library tag `mlx`), sorted by downloads.
    func search(query: String, author: String? = "mlx-community", limit: Int = 30) async throws -> [HubModel] {
        var components = URLComponents(string: "https://huggingface.co/api/models")!
        var items = [
            URLQueryItem(name: "filter", value: "mlx"),
            URLQueryItem(name: "sort", value: "downloads"),
            URLQueryItem(name: "direction", value: "-1"),
            URLQueryItem(name: "limit", value: String(limit)),
        ]
        if !query.isEmpty { items.append(URLQueryItem(name: "search", value: query)) }
        if let author { items.append(URLQueryItem(name: "author", value: author)) }
        components.queryItems = items

        let data = try await get(components.url!)
        struct Item: Codable {
            var id: String
            var pipeline_tag: String?
            var downloads: Int?
            var likes: Int?
            var tags: [String]?
            var gated: GatedValue?
        }
        // `gated` is false | "manual" | "auto" in the live API.
        enum GatedValue: Codable {
            case bool(Bool), string(String)
            init(from decoder: Decoder) throws {
                let container = try decoder.singleValueContainer()
                if let b = try? container.decode(Bool.self) { self = .bool(b) } else {
                    self = .string(try container.decode(String.self))
                }
            }
            func encode(to encoder: Encoder) throws {
                var container = encoder.singleValueContainer()
                switch self {
                case .bool(let b): try container.encode(b)
                case .string(let s): try container.encode(s)
                }
            }
            var isGated: Bool {
                if case .bool(let b) = self { return b }
                return true
            }
        }
        return try JSONDecoder().decode([Item].self, from: data).map {
            HubModel(
                id: $0.id,
                pipelineTag: $0.pipeline_tag,
                downloads: $0.downloads ?? 0,
                likes: $0.likes ?? 0,
                tags: $0.tags ?? [],
                gated: $0.gated?.isGated ?? false)
        }
    }

    /// Recursive file listing with true byte sizes (MODELS.md §1.2 —
    /// `recursive=true` is mandatory for subdirectory layouts).
    func files(repo: String, revision: String = "main") async throws -> [HubFile] {
        let url = URL(string: "https://huggingface.co/api/models/\(repo)/tree/\(revision)?recursive=true")!
        let data = try await get(url)
        struct Entry: Codable {
            var type: String
            var path: String
            var size: Int64?
        }
        return try JSONDecoder().decode([Entry].self, from: data)
            .filter { $0.type == "file" }
            .map { HubFile(path: $0.path, size: $0.size ?? 0) }
    }

    /// Total weight bytes of a repo (for RAM badges before download).
    func weightSize(repo: String) async throws -> Int64 {
        try await files(repo: repo)
            .filter { $0.path.hasSuffix(".safetensors") || $0.path.hasSuffix(".npz") }
            .reduce(0) { $0 + $1.size }
    }

    /// Resolve URL for one file (302s to the CDN; supports Range).
    nonisolated static func resolveURL(repo: String, path: String, revision: String = "main") -> URL {
        URL(string: "https://huggingface.co/\(repo)/resolve/\(revision)/\(path)")!
    }

    private func get(_ url: URL) async throws -> Data {
        var request = URLRequest(url: url)
        if let token = HFTokenStore.token, !token.isEmpty {
            request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
        }
        let (data, response) = try await session.data(for: request)
        guard let http = response as? HTTPURLResponse else { throw HubServiceError.badResponse(0) }
        guard (200..<300).contains(http.statusCode) else {
            throw HubServiceError.badResponse(http.statusCode)
        }
        return data
    }
}

/// Hugging Face token storage in the login Keychain (charter: encrypted
/// vault pattern; never plaintext on disk).
enum HFTokenStore {
    private static let service = "com.zyquo.mlx.hf-token"

    static var token: String? {
        let query: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrService as String: service,
            kSecReturnData as String: true,
        ]
        var result: AnyObject?
        guard SecItemCopyMatching(query as CFDictionary, &result) == errSecSuccess,
            let data = result as? Data
        else { return nil }
        return String(data: data, encoding: .utf8)
    }

    static func save(_ token: String) {
        let base: [String: Any] = [
            kSecClass as String: kSecClassGenericPassword,
            kSecAttrService as String: service,
        ]
        SecItemDelete(base as CFDictionary)
        guard !token.isEmpty else { return }
        var add = base
        add[kSecValueData as String] = Data(token.utf8)
        SecItemAdd(add as CFDictionary, nil)
    }
}