grammargem/Sources/GrammaGem/AI/ModelManager.swift

import Foundation

/// Downloads the local LLM weights from Hugging Face on demand. This is a *real*
/// download (the model files are fetched to disk with progress) — one of the
/// only three network touchpoints in the product, and it never sends user text.
@MainActor
final class ModelManager: ObservableObject {
    enum State: Equatable {
        case notDownloaded
        case downloading(progress: Double)
        case ready
        case failed(String)
    }

    @Published private(set) var state: State = .notDownloaded
    @Published private(set) var statusText: String = ""
    @Published var selectedRepo: String = AppConfig.Model.defaultRepo

    private var task: Task<Void, Never>?
    private var generation = 0

    init() {
        if isModelPresent(repo: selectedRepo) { state = .ready }
    }

    var modelsDirectory: URL {
        let base = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0]
        return base.appendingPathComponent("GrammaGem/Models", isDirectory: true)
    }

    func modelDir(_ repo: String) -> URL {
        modelsDirectory.appendingPathComponent(repo.replacingOccurrences(of: "/", with: "_"), isDirectory: true)
    }

    /// Considered present once the completion marker is written.
    func isModelPresent(repo: String) -> Bool {
        FileManager.default.fileExists(atPath: modelDir(repo).appendingPathComponent(".complete").path)
    }

    /// Kick off (or resume to) a real download. Safe to call from the UI.
    func startDownload(repo: String? = nil) {
        let target = repo ?? selectedRepo
        selectedRepo = target
        if isModelPresent(repo: target) { state = .ready; return }
        task?.cancel()
        generation += 1
        let gen = generation
        task = Task { await self.run(repo: target, generation: gen) }
    }

    func cancel() {
        task?.cancel()
        task = nil
        generation += 1 // invalidate any in-flight run's state writes
        statusText = ""
        state = isModelPresent(repo: selectedRepo) ? .ready : .notDownloaded
    }

    // MARK: - Download

    private func run(repo: String, generation gen: Int) async {
        if isModelPresent(repo: repo) { set(gen, .ready, ""); return }

        let dir = modelDir(repo)
        set(gen, .downloading(progress: 0), "Preparing…")
        do {
            try FileManager.default.createDirectory(at: dir, withIntermediateDirectories: true)
            let files = try await fetchFileList(repo: repo)
            guard !files.isEmpty else { throw ModelError.empty }
            let total = max(files.reduce(Int64(0)) { $0 + $1.size }, 1)
            var done: Int64 = 0

            for file in files {
                try Task.checkCancellation()
                set(gen, .downloading(progress: min(1, Double(done) / Double(total))),
                    "Downloading \(file.path) (\(byteString(file.size)))…")
                let url = URL(string: "\(AppConfig.Model.huggingFaceBase)/\(repo)/resolve/main/\(file.path)")!
                let (tmp, response) = try await URLSession.shared.download(from: url)
                if let http = response as? HTTPURLResponse, http.statusCode != 200 {
                    throw ModelError.http(http.statusCode, file.path)
                }
                let dest = dir.appendingPathComponent(file.path)
                try FileManager.default.createDirectory(
                    at: dest.deletingLastPathComponent(), withIntermediateDirectories: true)
                if FileManager.default.fileExists(atPath: dest.path) {
                    try FileManager.default.removeItem(at: dest)
                }
                try FileManager.default.moveItem(at: tmp, to: dest)

                // Reject a wrong/truncated body (e.g. an HTML interstitial served
                // with HTTP 200) by checking the on-disk size against the expected.
                if file.size > 0 {
                    let actual = (try? FileManager.default.attributesOfItem(atPath: dest.path)[.size] as? Int64) ?? nil
                    if let actual, actual != file.size {
                        throw ModelError.corrupt(file.path)
                    }
                }
                done += file.size
                set(gen, .downloading(progress: min(1, Double(done) / Double(total))), nil)
            }

            try Data().write(to: dir.appendingPathComponent(".complete"))
            set(gen, .ready, "")
            Log.ai.info("Model downloaded: \(repo, privacy: .public)")
        } catch is CancellationError {
            cleanupPartial(repo: repo)
            set(gen, isModelPresent(repo: repo) ? .ready : .notDownloaded, "")
        } catch let urlError as URLError where urlError.code == .cancelled {
            cleanupPartial(repo: repo)
            set(gen, isModelPresent(repo: repo) ? .ready : .notDownloaded, "")
        } catch {
            cleanupPartial(repo: repo)
            set(gen, .failed(error.localizedDescription), "")
            Log.ai.error("Model download failed: \(error.localizedDescription, privacy: .public)")
        }
    }

    /// Apply a state/status update only if this run hasn't been superseded.
    private func set(_ gen: Int, _ newState: State, _ status: String?) {
        guard gen == generation else { return }
        state = newState
        if let status { statusText = status }
    }

    /// Remove an incomplete model directory so partial files don't linger.
    private func cleanupPartial(repo: String) {
        guard !isModelPresent(repo: repo) else { return }
        try? FileManager.default.removeItem(at: modelDir(repo))
    }

    /// List the repo's files (+ sizes) from the Hugging Face tree API.
    private func fetchFileList(repo: String) async throws -> [(path: String, size: Int64)] {
        let url = URL(string: "\(AppConfig.Model.huggingFaceBase)/api/models/\(repo)/tree/main?recursive=1")!
        let (data, response) = try await URLSession.shared.data(from: url)
        if let http = response as? HTTPURLResponse, http.statusCode != 200 {
            throw ModelError.http(http.statusCode, "file list")
        }
        struct Entry: Decodable { let type: String; let path: String; let size: Int64? }
        let entries = try JSONDecoder().decode([Entry].self, from: data)
        return entries.filter { $0.type == "file" }.map { ($0.path, $0.size ?? 0) }
    }

    private func byteString(_ bytes: Int64) -> String {
        ByteCountFormatter.string(fromByteCount: bytes, countStyle: .file)
    }

    enum ModelError: LocalizedError {
        case empty
        case http(Int, String)
        case corrupt(String)
        var errorDescription: String? {
            switch self {
            case .empty: return "No model files were found for this repository."
            case .http(let code, let what): return "Download failed (\(code)) while fetching \(what)."
            case .corrupt(let what): return "Downloaded file looked wrong (\(what)); please retry."
            }
        }
    }
}