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

import Foundation

/// Phase 7 verification harness: `ZyquoLocal --verify [--keep]`.
/// Proves downloading and running real models works end-to-end on this Mac,
/// and dry-verifies the entire Featured catalog against the live Hub.
enum VerifyRunner {
    /// Models under test: span architectures and sizes that fit this Mac.
    private static let testRepos = [
        "mlx-community/Qwen3-0.6B-4bit",                 // tiny, qwen3
        "mlx-community/Llama-3.2-1B-Instruct-4bit",      // tiny, llama
        "mlx-community/SmolLM3-3B-4bit",                 // 3B, smollm3
        "mlx-community/gemma-3-4b-it-qat-4bit",          // 4B, gemma3
        "mlx-community/Qwen2.5-Coder-7B-Instruct-4bit",  // coding, qwen2
        "mlx-community/DeepSeek-R1-0528-Qwen3-8B-4bit",  // reasoning distill
    ]

    struct RunResult {
        var repoID: String
        var download = "–"
        var load = "–"
        var generate = "–"
        var multiTurn = "–"
        var cancel = "–"
        var unload = "–"
        var tokensPerSecond: Double?
        var ttft: TimeInterval?
        var note = ""
    }

    @MainActor
    static func run(arguments: [String]) async {
        let keepAll = arguments.contains("--keep")
        var results: [RunResult] = []

        let store = ModelStore()
        let hub = HubService(token: nil)
        let downloads = DownloadManager(hub: hub, store: store)

        for repoID in testRepos {
            var result = RunResult(repoID: repoID)
            log("\n═══ \(repoID) ═══")

            // 1 — download (or reuse)
            if store.model(for: repoID) == nil {
                log("downloading…")
                await downloads.download(repoID: repoID)
                var done = false
                var lastPercent = -1
                while !done {
                    try? await Task.sleep(for: .milliseconds(500))
                    guard let t = downloads.task(for: repoID) else {
                        result.note = "download task vanished"
                        break
                    }
                    switch t.state {
                    case .completed:
                        done = true
                    case .failed:
                        result.note = t.errorDescription ?? "download failed"
                        done = true
                    default:
                        let percent = Int(t.fractionCompleted * 100)
                        if percent / 10 != lastPercent / 10 {
                            lastPercent = percent
                            log("  \(percent)%  \(downloads.speeds[repoID].map(formatSpeed) ?? "")")
                        }
                    }
                }
            }
            store.rescan()
            guard let model = store.model(for: repoID) else {
                result.download = "❌"
                results.append(result)
                log("  download ❌ \(result.note)")
                continue
            }
            guard ModelStore.isValidModelDirectory(model.directory) else {
                result.download = "❌"
                result.note = "invalid directory after download"
                results.append(result)
                continue
            }
            result.download = "✅"
            log("  download ✅ (\(formatBytes(model.sizeBytes)), arch=\(model.architecture ?? "?"))")

            // 2 — load
            let engine = InferenceEngine()
            do {
                try await engine.load(model: model)
                result.load = "✅"
                log("  load ✅")
            } catch {
                result.load = "❌"
                result.note = error.localizedDescription
                results.append(result)
                log("  load ❌ \(error.localizedDescription)")
                continue
            }

            // 3 — deterministic generation
            do {
                let conversation = Conversation(
                    params: GenerationParams(temperature: 0, maxTokens: 600))
                try await engine.startSession(conversation: conversation)
                let (text, stats) = try await collect(
                    engine: engine, prompt: "Reply with exactly: OK",
                    params: conversation.params)
                var parser = ThinkTagParser()
                let (visible, _, _) = parser.consume(text)
                let cleaned = visible.trimmingCharacters(in: .whitespacesAndNewlines)
                if cleaned.isEmpty {
                    result.generate = "❌"
                    result.note = "empty output"
                } else {
                    result.generate = "✅"
                    result.ttft = stats?.timeToFirstToken
                    log("  generate ✅ “\(String(cleaned.prefix(60)))”")
                }
            } catch {
                result.generate = "❌"
                result.note = error.localizedDescription
                log("  generate ❌ \(error.localizedDescription)")
            }

            // 4 — multi-turn context carry-over
            do {
                let conversation = Conversation(
                    params: GenerationParams(temperature: 0, maxTokens: 800))
                try await engine.startSession(conversation: conversation)
                _ = try await collect(
                    engine: engine,
                    prompt: "My favorite color is vermilion. Just say: noted.",
                    params: conversation.params)
                let (answer, stats) = try await collect(
                    engine: engine, prompt: "What is my favorite color? Answer in one word.",
                    params: conversation.params)
                if answer.localizedCaseInsensitiveContains("vermilion") {
                    result.multiTurn = "✅"
                    result.tokensPerSecond = stats?.tokensPerSecond
                    log("  multi-turn ✅")
                } else {
                    result.multiTurn = "❌"
                    result.note = "no context carry-over: “\(String(answer.suffix(80)))”"
                    log("  multi-turn ❌ \(result.note)")
                }
                if result.tokensPerSecond == nil { result.tokensPerSecond = stats?.tokensPerSecond }
            } catch {
                result.multiTurn = "❌"
                result.note = error.localizedDescription
            }

            // 5 — streaming cancellation
            do {
                let conversation = Conversation(params: GenerationParams(temperature: 0.7))
                try await engine.startSession(conversation: conversation)
                let stream = try await engine.generate(
                    prompt: "Write a very long story about the ocean.",
                    params: conversation.params)
                var tokens = 0
                let start = Date()
                var sawEnd = false
                let consumer = Task {
                    for try await event in stream {
                        if case .token = event {
                            tokens += 1
                            if tokens == 12 { break }  // cancels via onTermination
                        }
                    }
                }
                _ = try? await consumer.value
                await engine.stopGeneration()
                sawEnd = true
                let elapsed = Date().timeIntervalSince(start)
                if sawEnd && elapsed < 30 {
                    result.cancel = "✅"
                    log("  cancel ✅ (stopped after \(tokens) tokens, \(String(format: "%.1f", elapsed))s)")
                } else {
                    result.cancel = "❌"
                }
            } catch {
                result.cancel = "❌"
                result.note = error.localizedDescription
            }

            // 6 — unload + memory release
            let before = MemoryAdvisor.activeMemoryBytes
            await engine.unload()
            try? await Task.sleep(for: .milliseconds(500))
            let after = MemoryAdvisor.activeMemoryBytes
            if after < max(200_000_000, before / 4) {
                result.unload = "✅"
                log("  unload ✅ (\(formatBytes(Int64(before))) → \(formatBytes(Int64(after))))")
            } else {
                result.unload = "❌"
                result.note += " memory not released (\(formatBytes(Int64(after))))"
                log("  unload ❌ (\(formatBytes(Int64(before))) → \(formatBytes(Int64(after))))")
            }

            results.append(result)
        }

        // Reclaim disk: keep the smallest model for ongoing dev.
        if !keepAll {
            for repoID in testRepos.dropFirst() where store.model(for: repoID) != nil {
                store.delete(repoID: repoID)
                log("deleted \(repoID) to reclaim disk")
            }
        }

        // ── Catalog dry-verification ────────────────────────────────────────
        log("\n═══ Featured catalog dry-verification (30 repos) ═══")
        var catalogRows: [(String, String, String)] = []
        for entry in ModelCatalog.featured {
            do {
                let (files, total) = try await hub.requiredFiles(of: entry.repoID)
                let gb = Double(total) / 1_000_000_000
                let deviation = abs(gb - entry.sizeGB) / entry.sizeGB
                let sizeOK = deviation < 0.10
                catalogRows.append((
                    entry.repoID,
                    "✅ \(files.count) files",
                    sizeOK
                        ? String(format: "✅ %.2f GB", gb)
                        : String(format: "⚠️ %.2f GB (catalog says %.2f)", gb, entry.sizeGB)
                ))
            } catch {
                catalogRows.append((entry.repoID, "❌ \(error.localizedDescription)", "–"))
            }
        }
        for row in catalogRows {
            log("  \(row.0): \(row.1) · \(row.2)")
        }

        // ── Results table ───────────────────────────────────────────────────
        var table = """
            | Model | Download | Load | Generate | Multi-turn | Cancel | Unload | tok/s | TTFT |
            |---|---|---|---|---|---|---|---|---|

            """
        for r in results {
            table += "| \(shortModelName(r.repoID)) | \(r.download) | \(r.load) | \(r.generate) | \(r.multiTurn) | \(r.cancel) | \(r.unload) | \(r.tokensPerSecond.map { String(format: "%.1f", $0) } ?? "–") | \(r.ttft.map { String(format: "%.2fs", $0) } ?? "–") |\n"
        }
        log("\n" + table)

        let allGreen = results.allSatisfy {
            $0.download == "✅" && $0.load == "✅" && $0.generate == "✅"
                && $0.multiTurn == "✅" && $0.cancel == "✅" && $0.unload == "✅"
        }
        let catalogGreen = catalogRows.allSatisfy { $0.1.hasPrefix("✅") }
        log(allGreen && catalogGreen ? "\nVERIFY: ALL GREEN" : "\nVERIFY: FAILURES PRESENT")
        for r in results where !r.note.isEmpty {
            log("  note[\(shortModelName(r.repoID))]: \(r.note)")
        }
        exit(allGreen && catalogGreen ? 0 : 1)
    }

    /// Collects one full generation, returning the raw text + final stats.
    @MainActor
    private static func collect(
        engine: InferenceEngine, prompt: String, params: GenerationParams
    ) async throws -> (String, GenerationStats?) {
        var text = ""
        var stats: GenerationStats?
        let stream = try await engine.generate(prompt: prompt, params: params)
        for try await event in stream {
            switch event {
            case .token(let t): text += t
            case .stats(let s): stats = s
            case .finished: break
            }
        }
        return (text, stats)
    }

    private static func log(_ text: String) {
        FileHandle.standardError.write(Data((text + "\n").utf8))
    }
}