import XCTest
@testable import MLXCore

/// Tests for the LM-Studio-style `<author>/<repo>` on-disk layout in
/// DownloadManager. New downloads land in the 2-level layout; existing flat
/// dirs continue to resolve via the dual-scan fallback. No auto-migration —
/// users redownload or move dirs manually.
final class DownloadManagerLayoutTests: XCTestCase {
    private var tempRoot: String!

    override func setUpWithError() throws {
        tempRoot = (NSTemporaryDirectory() as NSString)
            .appendingPathComponent("mlx-serve-tests-\(UUID().uuidString)")
        try FileManager.default.createDirectory(atPath: tempRoot, withIntermediateDirectories: true)
    }

    override func tearDownWithError() throws {
        try? FileManager.default.removeItem(atPath: tempRoot)
    }

    // MARK: - Path resolution

    func testNewLayoutDirSplitsAuthorAndName() {
        let p = DownloadManager.newLayoutDir(rootDir: tempRoot, repoId: "mlx-community/Qwen3.6-27B-mtp")
        XCTAssertEqual(p, (tempRoot as NSString)
            .appendingPathComponent("mlx-community")
            .appending("/Qwen3.6-27B-mtp"))
    }

    func testNewLayoutDirBareNameFallsBackToTopLevel() {
        // No author component — caller passed a bare name. Land at top level so
        // we don't fabricate an author dir.
        let p = DownloadManager.newLayoutDir(rootDir: tempRoot, repoId: "Qwen3.6-27B-mtp")
        XCTAssertEqual(p, (tempRoot as NSString).appendingPathComponent("Qwen3.6-27B-mtp"))
    }

    func testExistingModelDirPrefersNewLayout() throws {
        // Set up both: legacy flat AND new <author>/<name>.
        let name = "demo"
        let legacy = (tempRoot as NSString).appendingPathComponent(name)
        let nested = ((tempRoot as NSString).appendingPathComponent("acme") as NSString)
            .appendingPathComponent(name)
        try makeFakeModel(at: legacy)
        try makeFakeModel(at: nested)

        let resolved = DownloadManager.existingModelDir(rootDir: tempRoot, repoId: "acme/\(name)")
        XCTAssertEqual(resolved, nested, "new layout should win over legacy when both exist")
    }

    func testExistingModelDirFallsBackToLegacy() throws {
        // Only legacy exists. With a 2-level repoId we still want it found.
        let legacy = (tempRoot as NSString).appendingPathComponent("legacy-only")
        try makeFakeModel(at: legacy)

        let resolved = DownloadManager.existingModelDir(rootDir: tempRoot, repoId: "mlx-community/legacy-only")
        XCTAssertEqual(resolved, legacy, "legacy flat layout must remain discoverable until migrated")
    }

    func testExistingModelDirReturnsNilWhenAbsent() {
        XCTAssertNil(DownloadManager.existingModelDir(rootDir: tempRoot, repoId: "nobody/missing"))
    }

    // MARK: - File selection (recursive tree, incl. mtp/ sidecar)

    /// Regression for the silent MTP-sidecar drop: a model download must pull
    /// the nested `mtp/weights.safetensors` head (else the model loses its
    /// speculative-decoding speedup), while NOT pulling unrelated nested
    /// subdirectories (e.g. `original/` alternate-precision shadow weights).
    func testSelectNeededFilesIncludesMtpSidecarSkipsOtherNestedDirs() {
        let entries: [[String: Any]] = [
            ["path": "config.json", "type": "file", "size": 100],
            ["path": "model-00001-of-00003.safetensors", "type": "file", "size": 5_000_000_000],
            ["path": "tokenizer.json", "type": "file", "size": 19_000_000],
            ["path": "chat_template.jinja", "type": "file", "size": 7_000],
            ["path": "README.md", "type": "file", "size": 5_000],
            ["path": ".DS_Store", "type": "file", "size": 6_000],
            ["path": "mtp", "type": "directory", "size": 0],
            ["path": "mtp/weights.safetensors", "type": "file", "size": 524_000_000],
            ["path": "original/model.safetensors", "type": "file", "size": 50_000_000_000],
        ]
        let paths = Set(DownloadManager.selectNeededFiles(from: entries).map { $0.0 })

        XCTAssertTrue(paths.contains("mtp/weights.safetensors"), "MTP sidecar must be downloaded")
        XCTAssertTrue(paths.contains("config.json"))
        XCTAssertTrue(paths.contains("model-00001-of-00003.safetensors"))
        XCTAssertTrue(paths.contains("tokenizer.json"))
        XCTAssertTrue(paths.contains("chat_template.jinja"))

        XCTAssertFalse(paths.contains("mtp"), "a directory entry is not a downloadable file")
        XCTAssertFalse(paths.contains("README.md"), "non-weight markdown skipped")
        XCTAssertFalse(paths.contains(".DS_Store"), "files without a needed extension skipped")
        XCTAssertFalse(paths.contains("original/model.safetensors"),
                       "nested non-mtp shadow weights must not be pulled")

        // Sidecar size is threaded through for the progress/space pre-check.
        let sidecar = DownloadManager.selectNeededFiles(from: entries).first { $0.0 == "mtp/weights.safetensors" }
        XCTAssertEqual(sidecar?.1, 524_000_000)
    }

    // MARK: - Drafter discovery

    func testDiscoverDraftersFindsAllPublishedVariants() throws {
        // Drafters live under different authors today (mlx-community for the
        // older bf16 quants, google for the 12B official upload). The discoverer
        // must surface every variant regardless of its author prefix.
        for variant in GemmaVariant.allCases {
            let parts = variant.drafterRepoId.split(separator: "/")
            let dir = ((tempRoot as NSString).appendingPathComponent(String(parts[0])) as NSString)
                .appendingPathComponent(String(parts[1]))
            try makeDrafterDir(at: dir)
        }
        let found = DownloadManager.discoverDrafters(in: [tempRoot])
        XCTAssertEqual(Set(found.map { $0.variant }), Set(GemmaVariant.allCases))
    }

    func testDiscoverDraftersSkipsDirsWithWrongModelType() throws {
        // Wrong dirname: looks Gemma-shaped but isn't on the list.
        let bogus = ((tempRoot as NSString).appendingPathComponent("mlx-community") as NSString)
            .appendingPathComponent("gemma-4-other-it-assistant-bf16")
        try makeDrafterDir(at: bogus)
        // Right dirname but wrong model_type — NOT a drafter.
        let lookalike = ((tempRoot as NSString).appendingPathComponent("mlx-community") as NSString)
            .appendingPathComponent(GemmaVariant.E2B.drafterDirName)
        try FileManager.default.createDirectory(atPath: lookalike, withIntermediateDirectories: true)
        let cfg = (lookalike as NSString).appendingPathComponent("config.json")
        try "{\"model_type\":\"gemma4\"}".write(toFile: cfg, atomically: true, encoding: .utf8)

        XCTAssertTrue(DownloadManager.discoverDrafters(in: [tempRoot]).isEmpty)
    }

    func testDiscoverDraftersFirstRootWins() throws {
        // Same variant in two roots — earlier root takes precedence so a
        // user copy in ~/.mlx-serve/ wins over a leftover LM Studio copy.
        let alt = (NSTemporaryDirectory() as NSString)
            .appendingPathComponent("mlx-serve-tests-alt-\(UUID().uuidString)")
        defer { try? FileManager.default.removeItem(atPath: alt) }
        let primary = ((tempRoot as NSString).appendingPathComponent("mlx-community") as NSString)
            .appendingPathComponent(GemmaVariant.E4B.drafterDirName)
        let secondary = ((alt as NSString).appendingPathComponent("mlx-community") as NSString)
            .appendingPathComponent(GemmaVariant.E4B.drafterDirName)
        try makeDrafterDir(at: primary)
        try makeDrafterDir(at: secondary)

        let found = DownloadManager.discoverDrafters(in: [tempRoot, alt])
        XCTAssertEqual(found.count, 1)
        XCTAssertEqual(found.first?.url.path, primary)
    }

    func testGemmaVariantParsing() {
        XCTAssertEqual(DownloadManager.gemmaVariantFor(modelPath: "/m/gemma-4-e4b-it-4bit", isMoE: false), .E4B)
        XCTAssertEqual(DownloadManager.gemmaVariantFor(modelPath: "/m/gemma-4-e2b-it-8bit", isMoE: false), .E2B)
        XCTAssertEqual(DownloadManager.gemmaVariantFor(modelPath: "/m/gemma-4-12b-it-4bit", isMoE: false), .gemma12B)
        XCTAssertEqual(DownloadManager.gemmaVariantFor(modelPath: "/m/gemma-4-31b-it-4bit", isMoE: false), .gemma31B)
        XCTAssertEqual(DownloadManager.gemmaVariantFor(modelPath: "/m/gemma-4-26b-a4b-it-4bit", isMoE: true), .moe26B)
        // isMoE alone should also pick MoE so we route correctly even if the
        // path doesn't include the size designator.
        XCTAssertEqual(DownloadManager.gemmaVariantFor(modelPath: "/m/something-weird", isMoE: true), .moe26B)
        XCTAssertNil(DownloadManager.gemmaVariantFor(modelPath: "/m/qwen3-7b-4bit", isMoE: false))
    }

    // MARK: - Per-variant drafter repo paths

    /// All Gemma 4 drafters use the uniform mlx-community bf16 path —
    /// pinned because mlx-community only publishes 8bit for the new 12B
    /// drafter, and an earlier wholesale switch to 8bit was reverted after
    /// HF 401'd on the four older variants. Keep one suffix for consistency.
    func testDrafterRepoIdMatchesPublishedConvention() {
        XCTAssertEqual(GemmaVariant.E2B.drafterRepoId,      "mlx-community/gemma-4-E2B-it-assistant-bf16")
        XCTAssertEqual(GemmaVariant.E4B.drafterRepoId,      "mlx-community/gemma-4-E4B-it-assistant-bf16")
        XCTAssertEqual(GemmaVariant.gemma12B.drafterRepoId, "mlx-community/gemma-4-12B-it-assistant-bf16")
        XCTAssertEqual(GemmaVariant.moe26B.drafterRepoId,   "mlx-community/gemma-4-26B-A4B-it-assistant-bf16")
        XCTAssertEqual(GemmaVariant.gemma31B.drafterRepoId, "mlx-community/gemma-4-31B-it-assistant-bf16")
    }

    /// The 12B drafter declares `model_type: "gemma4_unified_assistant"` —
    /// a newer "unified" architecture spanning dense + MoE targets, distinct
    /// from the original `gemma4_assistant`. Both must classify as drafters
    /// so the dir doesn't surface as a base model in the tray-menu picker
    /// and doesn't trip the red "Unsupported architecture" label.
    func testDiscoverDraftersAcceptsUnifiedAssistantModelType() throws {
        let dir = ((tempRoot as NSString).appendingPathComponent("mlx-community") as NSString)
            .appendingPathComponent(GemmaVariant.gemma12B.drafterDirName)
        try FileManager.default.createDirectory(atPath: dir, withIntermediateDirectories: true)
        let cfg = (dir as NSString).appendingPathComponent("config.json")
        try "{\"model_type\":\"gemma4_unified_assistant\"}".write(toFile: cfg, atomically: true, encoding: .utf8)

        let found = DownloadManager.discoverDrafters(in: [tempRoot])
        XCTAssertEqual(found.first?.variant, .gemma12B)
    }

    // MARK: - GGUF classification & discovery

    func testGgufModelTypeRoutesDsv4ToDs4AndOthersToLlama() {
        // DeepSeek-V4-Flash → ds4 engine (case-insensitive).
        XCTAssertEqual(DownloadManager.ggufModelType(forBasename: "DeepSeek-V4-Flash-Q4_K_M.gguf"), "deepseek_v4")
        XCTAssertEqual(DownloadManager.ggufModelType(forBasename: "deepseek-v4-flash-bf16.gguf"), "deepseek_v4")
        // Any other GGUF → llama.cpp engine ("gguf").
        XCTAssertEqual(DownloadManager.ggufModelType(forBasename: "qwen2.5-0.5b-instruct-q4_k_m.gguf"), "gguf")
        XCTAssertEqual(DownloadManager.ggufModelType(forBasename: "Meta-Llama-3.1-8B-Instruct.Q4_K_M.gguf"), "gguf")
        // Not a GGUF → nil (won't be surfaced via the GGUF fast-path).
        XCTAssertNil(DownloadManager.ggufModelType(forBasename: "model.safetensors"))
        XCTAssertNil(DownloadManager.ggufModelType(forBasename: "config.json"))
    }

    func testGgufModelTypesAreSupportedArchitectures() {
        // Both engines' GGUF modelTypes must pass the architecture gate so the
        // model browser doesn't flag them "Unsupported architecture".
        for mt in ["gguf", "deepseek_v4"] {
            let m = LocalModel(
                id: "test:\(mt)", name: mt, path: "/tmp/x.gguf",
                sizeFormatted: "1 GB", modelType: mt, source: .custom, kind: .base
            )
            XCTAssertTrue(m.isSupportedArchitecture, "\"\(mt)\" must be in supportedModelTypes")
        }
    }

    func testQwen3MoeIsSupportedArchitecture() {
        // Qwen3-30B-A3B / Qwen3-Coder-30B-A3B ship model_type "qwen3_moe".
        // A locally-discovered checkpoint must NOT be flagged "Unsupported
        // architecture" in the model manager (issue #19).
        for mt in ["qwen3_moe", "qwen3_moe_text"] {
            let m = LocalModel(
                id: "test:\(mt)", name: mt, path: "/tmp/Qwen3-Coder-30B-A3B-8bit",
                sizeFormatted: "32 GB", modelType: mt, source: .custom, kind: .base
            )
            XCTAssertTrue(m.isSupportedArchitecture, "\"\(mt)\" must be in supportedModelTypes")
        }
    }

    func testGemma3TextIsSupportedArchitecture() {
        // Text-only Gemma 3 (Gemma3ForCausalLM) ships model_type "gemma3_text"
        // — e.g. mlx-community/gemma-3-12b-it-qat-abliterated-lm-4bit. A locally
        // discovered checkpoint must NOT be flagged "Unsupported architecture"
        // in the model manager (mirrors the gemma4_text / qwen3_moe_text tags).
        for mt in ["gemma3", "gemma3_text"] {
            let m = LocalModel(
                id: "test:\(mt)", name: mt, path: "/tmp/gemma-3-12b-it-qat-abliterated-lm-4bit",
                sizeFormatted: "7 GB", modelType: mt, source: .custom, kind: .base
            )
            XCTAssertTrue(m.isSupportedArchitecture, "\"\(mt)\" must be in supportedModelTypes")
        }
    }

    // MARK: - mmproj sidecar filtering

    /// `mmproj-*.gguf` files are CLIP / audio encoders, not language models —
    /// llama.cpp refuses them with "unsupported model architecture: 'clip'".
    /// The model-picker must skip them when scanning a vision-enabled folder
    /// (Gemma 4 VL, Qwen 3.6 VL, etc. ship both files side-by-side).
    func testIsMmprojGgufMatchesRealSidecars() {
        // Real mmproj basenames seen across the model zoo.
        XCTAssertTrue(DownloadManager.isMmprojGguf("mmproj-gemma-4-E4B-it-BF16.gguf"))
        XCTAssertTrue(DownloadManager.isMmprojGguf("mmproj-gemma-4-E2B-it-BF16.gguf"))
        XCTAssertTrue(DownloadManager.isMmprojGguf("mmproj-Qwen3.6-27B-VL-BF16.gguf"))
        XCTAssertTrue(DownloadManager.isMmprojGguf("MMPROJ-foo.gguf"))   // case-insensitive
        XCTAssertTrue(DownloadManager.isMmprojGguf("mmproj.gguf"))       // bare prefix
        // Real LLM .gguf files MUST NOT match.
        XCTAssertFalse(DownloadManager.isMmprojGguf("gemma-4-E4B-it-Q4_K_M.gguf"))
        XCTAssertFalse(DownloadManager.isMmprojGguf("Qwen3.5-4B-IQ4_NL.gguf"))
        XCTAssertFalse(DownloadManager.isMmprojGguf("DeepSeek-V4-Flash-Q4_K_M.gguf"))
        // Suffix-only — "model-mmproj.gguf" is NOT the wild-type convention.
        XCTAssertFalse(DownloadManager.isMmprojGguf("model-mmproj.gguf"))
        // Non-.gguf — not a sidecar.
        XCTAssertFalse(DownloadManager.isMmprojGguf("mmproj-readme.md"))
        XCTAssertFalse(DownloadManager.isMmprojGguf("mmproj"))
    }

    func testIsSupportedGgufExcludesMmprojSidecars() {
        // Real LLM .gguf is supported.
        XCTAssertTrue(DownloadManager.isSupportedGguf("gemma-4-E4B-it-Q4_K_M.gguf"))
        XCTAssertTrue(DownloadManager.isSupportedGguf("DeepSeek-V4-Flash-Q4_K_M.gguf"))
        // mmproj sidecars are NOT — this is the regression that made the model
        // picker hand the wrong .gguf to the server.
        XCTAssertFalse(DownloadManager.isSupportedGguf("mmproj-gemma-4-E4B-it-BF16.gguf"))
        XCTAssertFalse(DownloadManager.isSupportedGguf("mmproj-Qwen3.6-27B-VL-BF16.gguf"))
        // Non-.gguf: not a GGUF at all.
        XCTAssertFalse(DownloadManager.isSupportedGguf("config.json"))
    }

    // MARK: - Cancellation cleanup (full wipe)
    //
    // User-cancel removes the ENTIRE download dir — completed shards, config,
    // and `.partial` files alike — so a cancel leaves zero footprint: no remnant
    // that masquerades as a complete model, no undeletable config-only orphan.
    // (Network-error resume is a separate path that keeps `.partial`s; it does
    // NOT go through this wipe.)

    func testCancelWipeRemovesCompletedShardsNotJustPartials() throws {
        // The late-cancel case: some shards finished before the user hit cancel.
        let repoId = "acme/demo"
        let author = (tempRoot as NSString).appendingPathComponent("acme")
        let dir = DownloadManager.newLayoutDir(rootDir: tempRoot, repoId: repoId)
        try FileManager.default.createDirectory(atPath: dir, withIntermediateDirectories: true)
        let finalCfg = (dir as NSString).appendingPathComponent("config.json")
        let doneShard = (dir as NSString).appendingPathComponent("model-00001.safetensors")
        let inFlight = (dir as NSString).appendingPathComponent("model-00002.safetensors.partial")
        try "{}".write(toFile: finalCfg, atomically: true, encoding: .utf8)
        FileManager.default.createFile(atPath: doneShard, contents: Data("w".utf8))
        FileManager.default.createFile(atPath: inFlight, contents: Data())

        DownloadManager.removeModelFiles(at: dir, roots: [tempRoot])

        let fm = FileManager.default
        XCTAssertFalse(fm.fileExists(atPath: dir), "whole download dir must be gone, completed shards included")
        XCTAssertFalse(fm.fileExists(atPath: author), "now-empty author dir must be pruned")
        XCTAssertTrue(fm.fileExists(atPath: tempRoot), "scan root must survive")
    }

    func testCancelWipeRemovesConfigOnlyOrphan() throws {
        // The previously-invisible case: config downloaded, no shard yet. Must
        // not be left behind (it wouldn't surface as a deletable LocalModel).
        let repoId = "acme/justconfig"
        let dir = DownloadManager.newLayoutDir(rootDir: tempRoot, repoId: repoId)
        try makeFakeModel(at: dir)   // config.json only

        DownloadManager.removeModelFiles(at: dir, roots: [tempRoot])

        XCTAssertFalse(FileManager.default.fileExists(atPath: dir))
    }

    func testCancelWipeNoOpWhenDirMissing() {
        // Cancelling a fresh download that bailed before mkdir must not crash.
        let dir = DownloadManager.newLayoutDir(rootDir: tempRoot, repoId: "ghost/never-created")
        XCTAssertFalse(DownloadManager.removeModelFiles(at: dir, roots: [tempRoot]))
    }

    // MARK: - Delete by on-disk path
    //
    // `LocalModelRow` used to delete via `deleteModel(repoId: model.id)`, but a
    // LocalModel's `id` is source-prefixed (`"mlxServe:author/name"`), so the
    // repoId-based path resolver looked under `<root>/mlxServe:author/name` and
    // deleted nothing — the trash button silently no-op'd and the user had to
    // `rm -rf` from a terminal. Deletion now keys off the model's real resolved
    // `path` instead. These pin that behavior.

    func testRemoveModelFilesDeletesNestedDirAndPrunesAuthor() throws {
        let author = (tempRoot as NSString).appendingPathComponent("acme")
        let modelDir = (author as NSString).appendingPathComponent("demo")
        try makeFakeModel(at: modelDir)
        let partial = (modelDir as NSString).appendingPathComponent("model.safetensors.partial")
        FileManager.default.createFile(atPath: partial, contents: Data())

        let removed = DownloadManager.removeModelFiles(at: modelDir, roots: [tempRoot])

        let fm = FileManager.default
        XCTAssertTrue(removed)
        XCTAssertFalse(fm.fileExists(atPath: modelDir), "model dir (incl. .partial) must be gone")
        XCTAssertFalse(fm.fileExists(atPath: author), "now-empty author dir must be pruned")
        XCTAssertTrue(fm.fileExists(atPath: tempRoot), "the scan root itself must never be removed")
    }

    func testRemoveModelFilesFromGgufFilePath() throws {
        // A GGUF model's `path` is the .gguf file, not its containing dir.
        let modelDir = ((tempRoot as NSString).appendingPathComponent("team") as NSString)
            .appendingPathComponent("mygguf")
        try FileManager.default.createDirectory(atPath: modelDir, withIntermediateDirectories: true)
        let gguf = (modelDir as NSString).appendingPathComponent("model-Q4_K_M.gguf")
        FileManager.default.createFile(atPath: gguf, contents: Data("x".utf8))

        let removed = DownloadManager.removeModelFiles(at: gguf, roots: [tempRoot])

        XCTAssertTrue(removed)
        XCTAssertFalse(FileManager.default.fileExists(atPath: modelDir),
                       "deleting by a file path must remove its containing model dir")
    }

    func testRemoveModelFilesKeepsAuthorWithSurvivingSiblings() throws {
        let author = (tempRoot as NSString).appendingPathComponent("acme")
        let a = (author as NSString).appendingPathComponent("model-a")
        let b = (author as NSString).appendingPathComponent("model-b")
        try makeFakeModel(at: a)
        try makeFakeModel(at: b)

        DownloadManager.removeModelFiles(at: a, roots: [tempRoot])

        let fm = FileManager.default
        XCTAssertFalse(fm.fileExists(atPath: a))
        XCTAssertTrue(fm.fileExists(atPath: author), "author dir must survive while a sibling model remains")
        XCTAssertTrue(fm.fileExists(atPath: b))
    }

    func testRemoveModelFilesLegacyFlatDoesNotPruneRoot() throws {
        // Legacy flat layout: the model dir sits directly under a root, so its
        // "author" IS the root — pruning must stop there.
        let modelDir = (tempRoot as NSString).appendingPathComponent("flatmodel")
        try makeFakeModel(at: modelDir)

        DownloadManager.removeModelFiles(at: modelDir, roots: [tempRoot])

        let fm = FileManager.default
        XCTAssertFalse(fm.fileExists(atPath: modelDir))
        XCTAssertTrue(fm.fileExists(atPath: tempRoot), "a root must never be pruned even when emptied")
    }

    func testRemoveModelFilesRefusesToDeleteARoot() throws {
        // Defensive: never remove a root directory itself.
        let removed = DownloadManager.removeModelFiles(at: tempRoot, roots: [tempRoot])
        XCTAssertFalse(removed)
        XCTAssertTrue(FileManager.default.fileExists(atPath: tempRoot))
    }

    func testRemoveModelFilesMissingPathIsNoOp() {
        let ghost = (tempRoot as NSString).appendingPathComponent("nope/missing")
        XCTAssertFalse(DownloadManager.removeModelFiles(at: ghost, roots: [tempRoot]))
    }

    // MARK: - Cancellation detection
    //
    // Cancelling an in-flight download surfaces as URLSession's
    // NSURLErrorCancelled, NOT Swift's CancellationError — so the retry loop
    // must recognize both, or a cancelled transfer flashes "retrying…" before
    // it finally unwinds.

    func testIsCancellationMatchesCancellationErrorAndURLCancel() {
        XCTAssertTrue(DownloadManager.isCancellation(CancellationError()))
        XCTAssertTrue(DownloadManager.isCancellation(URLError(.cancelled)))
        XCTAssertTrue(DownloadManager.isCancellation(
            NSError(domain: NSURLErrorDomain, code: NSURLErrorCancelled)))
    }

    func testIsCancellationRejectsTransientFailures() {
        // These are genuine transient errors the retry loop must keep retrying.
        XCTAssertFalse(DownloadManager.isCancellation(URLError(.timedOut)))
        XCTAssertFalse(DownloadManager.isCancellation(URLError(.networkConnectionLost)))
        XCTAssertFalse(DownloadManager.isCancellation(URLError(.notConnectedToInternet)))
    }

    // MARK: - Downloaded tab live-refresh trigger
    //
    // The Downloaded tab's "Size on Disk" comes from a disk re-scan
    // (`refreshModels`), which previously only fired on tab entry — so sizes
    // froze mid-download until the user toggled the button. The tab now
    // live-polls, but only while it's showing AND a download is in flight.

    func testShouldLivePollOnlyWhenTabOpenAndDownloading() {
        XCTAssertTrue(ModelBrowserView.shouldLivePoll(downloadedTab: true, hasActiveDownloads: true))
        // No active downloads → nothing to refresh, don't spin.
        XCTAssertFalse(ModelBrowserView.shouldLivePoll(downloadedTab: true, hasActiveDownloads: false))
        // Not in the Downloaded tab → the sizes aren't even visible.
        XCTAssertFalse(ModelBrowserView.shouldLivePoll(downloadedTab: false, hasActiveDownloads: true))
        XCTAssertFalse(ModelBrowserView.shouldLivePoll(downloadedTab: false, hasActiveDownloads: false))
    }

    // MARK: - LocalModel metadata caption
    //
    // The Downloaded tab used to render only a name + delete button. These pin
    // the parsed metadata (params / quant / architecture / engine) that now
    // makes each row actually informative.

    private func localModel(
        name: String, type: String, path: String,
        vision: Bool = false, quantBits: Int? = nil, ctx: Int? = nil,
        numExperts: Int? = nil, activeExperts: Int? = nil
    ) -> LocalModel {
        LocalModel(id: "mlxServe:\(name)", name: name, path: path,
                   sizeFormatted: "10 GB", modelType: type, source: .mlxServe, kind: .base,
                   hasVision: vision, quantBits: quantBits, contextLength: ctx,
                   numExperts: numExperts, activeExperts: activeExperts)
    }

    // Captions below use the real config values dumped from the user's models:
    // Qwen2.5-Coder-32B (qwen2, bits 8, ctx 32768, dense),
    // Qwen3-Coder-30B-A3B (qwen3_moe, bits 8, ctx 262144, 128/8 experts),
    // Qwen3-Coder-Next (qwen3_next, bits 4, ctx 262144, 512/10 experts).

    func testMetadataSummaryDenseFromConfig() {
        let m = localModel(name: "Qwen2.5-Coder-32B-Instruct-8bit", type: "qwen2",
                           path: "/m/Qwen2.5-Coder-32B-Instruct-8bit",
                           quantBits: 8, ctx: 32768)
        XCTAssertEqual(m.paramSize, "32B")
        XCTAssertNil(m.expertSummary, "dense model → no expert token")
        // Format reads "MLX" (the weight format), not the "MLX-Serve" app name.
        XCTAssertEqual(m.metadataSummary, "32B · 8-bit · 32K ctx · qwen2 · MLX")
        XCTAssertTrue(m.hasToolCalling)
        XCTAssertFalse(m.hasVision)
    }

    func testMetadataSummaryMoEShowsConfigExperts() {
        let m = localModel(name: "Qwen3-Coder-30B-A3B-Instruct-8bit", type: "qwen3_moe",
                           path: "/m/Qwen3-Coder-30B-A3B-Instruct-8bit",
                           quantBits: 8, ctx: 262144, numExperts: 128, activeExperts: 8)
        XCTAssertEqual(m.expertSummary, "8/128 experts")
        XCTAssertEqual(m.metadataSummary, "30B · 8-bit · 8/128 experts · 256K ctx · qwen3_moe · MLX")
    }

    func testMetadataSummaryConfigQuantOverridesNameAndNoParamToken() {
        // Name says nothing about params; config says bits 4. Caption must use
        // the config bits and skip the missing param token.
        let n = localModel(name: "Qwen3-Coder-Next-4bit", type: "qwen3_next",
                           path: "/m/Qwen3-Coder-Next-4bit",
                           quantBits: 4, ctx: 262144, numExperts: 512, activeExperts: 10)
        XCTAssertNil(n.paramSize)
        XCTAssertEqual(n.quantization, "4-bit")
        XCTAssertEqual(n.metadataSummary, "4-bit · 10/512 experts · 256K ctx · qwen3_next · MLX")
    }

    func testMetadataSummaryGgufFallsBackToNameQuant() {
        // GGUF has no config.json here → quant comes from the name, format = GGUF.
        let g = LocalModel(id: "mlxServe:team/m", name: "Meta-Llama-3.1-8B-Instruct-Q4_K_M",
                           path: "/m/team/m/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf",
                           sizeFormatted: "5 GB", modelType: "gguf", source: .mlxServe, kind: .base)
        XCTAssertEqual(g.paramSize, "8B")
        XCTAssertEqual(g.quantization, "4-bit")   // Q4_K_M → 4-bit, via name fallback
        XCTAssertEqual(g.metadataSummary, "8B · 4-bit · gguf · GGUF")
        XCTAssertTrue(g.hasToolCalling)
    }

    func testContextFormatting() {
        XCTAssertEqual(LocalModel.formatContext(262144), "256K ctx")
        XCTAssertEqual(LocalModel.formatContext(32768), "32K ctx")
        XCTAssertEqual(LocalModel.formatContext(1048576), "1M ctx")
        XCTAssertEqual(LocalModel.formatContext(512), "512 ctx")
    }

    // MARK: - config.json parsing (the authoritative metadata source)

    func testParseConfigMetadataReadsQuantCtxExpertsVision() throws {
        let dir = (tempRoot as NSString).appendingPathComponent("cfg-moe")
        try FileManager.default.createDirectory(atPath: dir, withIntermediateDirectories: true)
        let cfg = (dir as NSString).appendingPathComponent("config.json")
        try """
        {"model_type":"qwen3_moe","quantization":{"group_size":64,"bits":8},
         "max_position_embeddings":262144,"num_experts":128,"num_experts_per_tok":8}
        """.write(toFile: cfg, atomically: true, encoding: .utf8)

        let meta = DownloadManager.parseConfigMetadata(atPath: cfg)
        XCTAssertEqual(meta.modelType, "qwen3_moe")
        XCTAssertEqual(meta.quantBits, 8)
        XCTAssertEqual(meta.contextLength, 262144)
        XCTAssertEqual(meta.numExperts, 128)
        XCTAssertEqual(meta.activeExperts, 8)
        XCTAssertFalse(meta.hasVision)
    }

    func testParseConfigMetadataQuantizationConfigAndVision() throws {
        // Some checkpoints use `quantization_config`; vision via `vision_config`.
        let dir = (tempRoot as NSString).appendingPathComponent("cfg-vision")
        try FileManager.default.createDirectory(atPath: dir, withIntermediateDirectories: true)
        let cfg = (dir as NSString).appendingPathComponent("config.json")
        try """
        {"model_type":"gemma4","quantization_config":{"bits":4},"vision_config":{"hidden_size":1152}}
        """.write(toFile: cfg, atomically: true, encoding: .utf8)

        let meta = DownloadManager.parseConfigMetadata(atPath: cfg)
        XCTAssertEqual(meta.quantBits, 4)
        XCTAssertTrue(meta.hasVision)
    }

    func testParseConfigMetadataTextOnlyArchSuppressesVision() throws {
        // A `_text` arch with a vestigial vision_config must NOT report vision.
        let dir = (tempRoot as NSString).appendingPathComponent("cfg-text")
        try FileManager.default.createDirectory(atPath: dir, withIntermediateDirectories: true)
        let cfg = (dir as NSString).appendingPathComponent("config.json")
        try #"{"model_type":"qwen3_5_moe_text","vision_config":{"x":1}}"#
            .write(toFile: cfg, atomically: true, encoding: .utf8)

        XCTAssertFalse(DownloadManager.parseConfigMetadata(atPath: cfg).hasVision)
    }

    func testParseConfigMetadataMissingFileDefaults() {
        let meta = DownloadManager.parseConfigMetadata(atPath: "/nope/config.json")
        XCTAssertEqual(meta, DownloadManager.ConfigMetadata())
    }

    // MARK: - Helpers

    /// Minimal model dir layout: just `config.json`. The path-resolution and
    /// migration logic only checks for that file's presence.
    private func makeFakeModel(at path: String) throws {
        try FileManager.default.createDirectory(atPath: path, withIntermediateDirectories: true)
        let cfg = (path as NSString).appendingPathComponent("config.json")
        try "{}".write(toFile: cfg, atomically: true, encoding: .utf8)
    }

    /// Drafter dir: config.json with `model_type: "gemma4_assistant"`.
    private func makeDrafterDir(at path: String) throws {
        try FileManager.default.createDirectory(atPath: path, withIntermediateDirectories: true)
        let cfg = (path as NSString).appendingPathComponent("config.json")
        try "{\"model_type\":\"gemma4_assistant\"}".write(toFile: cfg, atomically: true, encoding: .utf8)
    }
}
