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

import Foundation

/// Live Hugging Face Hub client: search, model info, file listings.
/// Contract documented in docs/MODELS.md §1. All responses are Codable
/// structs — no dictionary spelunking.
struct HubService: Sendable {
    /// Search scope shown in the Discover UI.
    enum Scope: String, CaseIterable, Sendable {
        case featured
        case mlxCommunity
        case allMLX
    }

    enum Sort: String, CaseIterable, Sendable {
        case downloads
        case likes
        case newest

        var apiValue: String {
            switch self {
            case .downloads: "downloads"
            case .likes: "likes"
            case .newest: "createdAt"
            }
        }
    }

    /// One search result from /api/models.
    struct ModelSummary: Codable, Hashable, Sendable, Identifiable {
        var id: String
        var likes: Int?
        var downloads: Int?
        var gated: GatedValue?
        var tags: [String]?
        var createdAt: Date?
        var config: ModelConfigInfo?

        var isGated: Bool { gated?.isGated ?? false }
        var architecture: String? { config?.modelType }

        /// True when the architecture is known-supported by the MLX engine.
        var isSupportedArchitecture: Bool {
            guard let architecture else { return true }  // unknown → don't block, warn later
            return SupportedArchitectures.contains(architecture)
        }
    }

    struct ModelConfigInfo: Codable, Hashable, Sendable {
        var modelType: String?

        enum CodingKeys: String, CodingKey {
            case modelType = "model_type"
        }
    }

    /// `gated` is `false` or a string ("auto"/"manual").
    enum GatedValue: Codable, Hashable, Sendable {
        case bool(Bool)
        case mode(String)

        var isGated: Bool {
            switch self {
            case .bool(let b): b
            case .mode: true
            }
        }

        init(from decoder: Decoder) throws {
            let container = try decoder.singleValueContainer()
            if let b = try? container.decode(Bool.self) {
                self = .bool(b)
            } else {
                self = .mode(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 .mode(let s): try container.encode(s)
            }
        }
    }

    /// One file from /tree/main?recursive=true.
    struct RepoFile: Codable, Hashable, Sendable {
        var type: String
        var path: String
        var size: Int64?
        var lfs: LFSInfo?

        struct LFSInfo: Codable, Hashable, Sendable {
            var oid: String?
            var size: Int64?
        }

        var isFile: Bool { type == "file" }

        /// Files the MLX stack needs (mirrors the package's own filter).
        var isModelFile: Bool {
            path.hasSuffix(".safetensors") || path.hasSuffix(".json") || path.hasSuffix(".jinja")
        }
    }

    enum HubError: LocalizedError {
        case http(Int, String)
        case rateLimited(retryAfterSeconds: Int?)
        case gatedOrMissing(String)
        case invalidResponse

        var errorDescription: String? {
            switch self {
            case .http(let code, let repo):
                "Hugging Face returned HTTP \(code) for \(repo)."
            case .rateLimited(let retry):
                "Hugging Face rate limit reached. Try again in \(retry.map { "\($0)s" } ?? "a moment")."
            case .gatedOrMissing(let repo):
                "\(repo) requires access approval or does not exist. For gated models (Llama, Gemma), add a Hugging Face token in Settings."
            case .invalidResponse:
                "Unexpected response from Hugging Face."
            }
        }
    }

    var token: String?

    private static let base = URL(string: "https://huggingface.co")!

    private var session: URLSession { URLSession.shared }

    private func request(_ url: URL) -> URLRequest {
        var request = URLRequest(url: url)
        request.setValue("ZyquoLocal/1.0.0", forHTTPHeaderField: "User-Agent")
        if let token, !token.isEmpty {
            request.setValue("Bearer \(token)", forHTTPHeaderField: "Authorization")
        }
        return request
    }

    private static let decoder: JSONDecoder = {
        let decoder = JSONDecoder()
        decoder.dateDecodingStrategy = .custom { d in
            let raw = try d.singleValueContainer().decode(String.self)
            if let date = try? Date(raw, strategy: Date.ISO8601FormatStyle(includingFractionalSeconds: true)) {
                return date
            }
            if let date = try? Date(raw, strategy: .iso8601) {
                return date
            }
            throw DecodingError.dataCorrupted(.init(codingPath: d.codingPath, debugDescription: "Bad date \(raw)"))
        }
        return decoder
    }()

    private func get<T: Decodable>(_ url: URL, as type: T.Type, context: String) async throws -> T {
        let (data, response) = try await session.data(for: request(url))
        guard let http = response as? HTTPURLResponse else { throw HubError.invalidResponse }
        switch http.statusCode {
        case 200:
            return try Self.decoder.decode(T.self, from: data)
        case 401, 403:
            throw HubError.gatedOrMissing(context)
        case 429:
            let retry = http.value(forHTTPHeaderField: "retry-after").flatMap(Int.init)
            throw HubError.rateLimited(retryAfterSeconds: retry)
        default:
            throw HubError.http(http.statusCode, context)
        }
    }

    // MARK: - Search

    /// Live Hub search. `scope .allMLX` searches all of HF filtered to the
    /// `mlx` library tag; architecture compatibility is checked via
    /// `config.model_type` (requested with config=true).
    func search(
        query: String,
        scope: Scope,
        sort: Sort = .downloads,
        limit: Int = 40
    ) async throws -> [ModelSummary] {
        var components = URLComponents(
            url: Self.base.appendingPathComponent("api/models"), resolvingAgainstBaseURL: false)!
        var items = [
            URLQueryItem(name: "pipeline_tag", value: "text-generation"),
            URLQueryItem(name: "sort", value: sort.apiValue),
            URLQueryItem(name: "direction", value: "-1"),
            URLQueryItem(name: "limit", value: String(limit)),
            URLQueryItem(name: "config", value: "true"),
        ]
        if !query.isEmpty {
            items.append(URLQueryItem(name: "search", value: query))
        }
        switch scope {
        case .featured, .mlxCommunity:
            items.append(URLQueryItem(name: "author", value: "mlx-community"))
        case .allMLX:
            items.append(URLQueryItem(name: "filter", value: "mlx"))
        }
        components.queryItems = items
        return try await get([ModelSummary].self, from: components.url!, context: "search")
    }

    private func get<T: Decodable>(_ type: T.Type, from url: URL, context: String) async throws -> T {
        try await get(url, as: type, context: context)
    }

    // MARK: - Files

    /// Full file listing with sizes for a repo (paginates if needed).
    func files(of repoID: String) async throws -> [RepoFile] {
        var all: [RepoFile] = []
        var url: URL? = Self.base.appendingPathComponent("api/models/\(repoID)/tree/main")
            .appending(queryItems: [URLQueryItem(name: "recursive", value: "true")])
        while let current = url {
            let (data, response) = try await session.data(for: request(current))
            guard let http = response as? HTTPURLResponse else { throw HubError.invalidResponse }
            guard http.statusCode == 200 else {
                if http.statusCode == 401 || http.statusCode == 403 {
                    throw HubError.gatedOrMissing(repoID)
                }
                throw HubError.http(http.statusCode, repoID)
            }
            all += try Self.decoder.decode([RepoFile].self, from: data)
            url = Self.nextPage(from: http)
        }
        return all
    }

    /// Files required to run the model, with total download size.
    func requiredFiles(of repoID: String) async throws -> (files: [RepoFile], totalBytes: Int64) {
        let required = try await files(of: repoID).filter { $0.isFile && $0.isModelFile }
        let total = required.reduce(Int64(0)) { $0 + ($1.size ?? 0) }
        return (required, total)
    }

    /// Download URL for one file of a repo.
    static func resolveURL(repoID: String, path: String) -> URL {
        base.appendingPathComponent("\(repoID)/resolve/main/\(path)")
    }

    /// Parses RFC-5988 Link header for cursor pagination.
    private static func nextPage(from response: HTTPURLResponse) -> URL? {
        guard let link = response.value(forHTTPHeaderField: "Link") else { return nil }
        for part in link.split(separator: ",") {
            let segments = part.split(separator: ";").map { $0.trimmingCharacters(in: .whitespaces) }
            guard segments.count >= 2, segments.contains(where: { $0 == "rel=\"next\"" }),
                let urlPart = segments.first, urlPart.hasPrefix("<"), urlPart.hasSuffix(">")
            else { continue }
            return URL(string: String(urlPart.dropFirst().dropLast()))
        }
        return nil
    }
}

/// Architectures supported by the MLX Swift LLM layer (verified list from
/// docs/MLX-RESEARCH.md §4 — LLMTypeRegistry @ mlx-swift-lm 3.31.4).
enum SupportedArchitectures {
    static let all: Set<String> = [
        "mistral", "mixtral", "llama", "phi", "phi3", "phimoe", "gemma", "gemma2",
        "gemma3", "gemma3_text", "gemma3n", "gemma4", "gemma4_unified", "gemma4_text",
        "qwen2", "qwen3", "qwen3_moe", "qwen3_next", "qwen3_5", "qwen3_5_moe",
        "qwen3_5_text", "minicpm", "starcoder2", "cohere", "openelm", "internlm2",
        "deepseek_v3", "granite", "granitemoehybrid", "mimo", "mimo_v2_flash",
        "minimax", "glm4", "glm4_moe", "glm4_moe_lite", "acereason", "falcon_h1",
        "bitnet", "smollm3", "ernie4_5", "lfm2", "baichuan_m1", "exaone4", "gpt_oss",
        "lille-130m", "olmoe", "olmo2", "olmo3", "bailing_moe", "lfm2_moe",
        "nanochat", "nemotron_h", "afmoe", "jamba", "mamba2", "mistral3", "apertus",
        "nemotron_labs_diffusion",
    ]

    static func contains(_ architecture: String) -> Bool {
        all.contains(architecture)
    }
}