//
//  VerifyHarness.swift
//  Zyquo Agent
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  Phase 7.1 — provider tool-calling verification, run against LIVE APIs with
//  the exact production provider clients. For every agent-capable model whose
//  provider has an API key available (environment or vault), the harness:
//
//    1. Streams a chat request offering ONE tool (`list_files`) and asserts the
//       model emits a tool call with the right name, syntactically valid JSON
//       arguments, and a normalized StopReason of `.toolUse`.
//    2. Threads a canned tool_result back in the provider's wire format and
//       asserts the model produces a final text answer that mentions at least
//       one of the listed file names, with StopReason `.endTurn`.
//    3. Asserts streaming actually streamed (more than one text/argument delta
//       across the exchange, not a single blob).
//
//  Results are printed live and written as a markdown table to
//  docs/VERIFICATION.md. Providers run sequentially, models within a provider
//  sequentially, with a polite delay between calls. Keys are never logged.
//
//  Usage:  ZyquoAgent --verify [--provider <id>] [--model <id | provider/id>]
//

import Foundation

enum VerifyHarness {

    // MARK: - Constants

    /// The single tool offered to every model. {path?: string} — an empty
    /// arguments object is a perfectly valid call.
    static let listFilesTool = ToolSpec(
        name: "list_files",
        description: "List the files in the current workspace directory.",
        parametersJSONSchema: """
        {"type":"object","properties":{"path":{"type":"string","description":"Optional subdirectory to list, relative to the workspace root."}}}
        """
    )

    static let systemPrompt = """
        You are Zyquo Agent, an autonomous agent operating inside a workspace directory on the user's Mac. \
        You must use the provided tools to inspect the workspace — never guess or fabricate file listings. \
        Call a tool when you need information. Once you have the tool result, answer the user in plain text \
        (naming the files you found) and stop; do not call the same tool again.
        """

    static let userPrompt = "List the files in the workspace using the tool."

    /// Canned tool_result threaded back on step 2.
    static let cannedListing = "MEMORY.md\nnotes.txt\nresult.csv"
    static let expectedNames = ["MEMORY.md", "notes.txt", "result.csv"]

    /// Whole-model budget (both steps + retries within a step).
    static let perModelTimeout: TimeInterval = 90

    // MARK: - Result model

    struct ModelReport {
        let model: AIModel
        var toolCall = false
        var threading = false
        var streaming = false
        var latency: TimeInterval?
        var notes: [String] = []

        var green: Bool { toolCall && threading && streaming }
    }

    private struct VerifyTimeout: Error {}

    /// Everything observed while one streamed response was consumed.
    private struct StreamOutcome {
        var text = ""
        var textDeltas = 0
        var reasoningDeltas = 0
        var toolCallStarts = 0
        var argDeltas = 0
        var toolCalls: [ToolCall] = []
        var rawReason: String?
        var stop: StopReason = .other(nil)
    }

    // MARK: - Entry

    static func run(arguments: [String]) async -> Int32 {
        let onlyProvider = value(after: "--provider", in: arguments).flatMap { ProviderID(rawValue: $0) }
        let onlyModel = value(after: "--model", in: arguments)

        var models = ModelCatalogData.all.filter(\.agentCapable)
        if let onlyProvider {
            models = models.filter { $0.provider == onlyProvider }
        }
        if let onlyModel {
            models = models.filter { $0.id == onlyModel || "\($0.provider.rawValue)/\($0.id)" == onlyModel }
            if models.isEmpty {
                FileHandle.standardError.write(Data("No agent-capable model matches “\(onlyModel)”.\n".utf8))
                return 64
            }
        }

        // Group by provider, in the canonical provider order.
        let providers = ProviderID.builtIn.filter { provider in
            models.contains { $0.provider == provider }
        }
        var keyless: [ProviderID] = []
        var reports: [ModelReport] = []

        let total = providers.reduce(0) { count, provider in
            AgentCLI.resolveAPIKey(for: provider) == nil
                ? count
                : count + models.filter { $0.provider == provider }.count
        }
        print("Zyquo Agent — provider tool-calling verification")
        print("  \(total) agent-capable model(s) across \(providers.count) provider(s)\n")

        for provider in providers {
            let providerModels = models.filter { $0.provider == provider }
            guard let key = AgentCLI.resolveAPIKey(for: provider) else {
                keyless.append(provider)
                print("⚠️  \(provider.displayName): no API key in environment or vault — \(providerModels.count) model(s) skipped")
                continue
            }
            let client = ProviderRegistry.client(for: provider)
            print("── \(provider.displayName) (\(providerModels.count) model(s)) " + String(repeating: "─", count: 30))
            for model in providerModels {
                var report = await checkWithRetry(model: model, client: client, key: key)
                if !report.green {
                    report.notes.insert("failed twice (attempt 1 + retry)", at: 0)
                }
                reports.append(report)
                printLine(report)
                try? await Task.sleep(nanoseconds: interCallDelay(for: provider))
            }
        }

        writeReport(reports: reports, keyless: keyless)

        let green = reports.filter(\.green).count
        print("\n\(reports.count) model(s) tested · \(green) green · \(reports.count - green) failed")
        print("Full table: docs/VERIFICATION.md")
        return green == reports.count ? 0 : 1
    }

    // MARK: - Per-model check

    /// One attempt; on any failure, a single automatic retry (transient
    /// hiccups, rate-limit blips). Notes reflect the final attempt.
    private static func checkWithRetry(model: AIModel, client: any ProviderClient, key: String) async -> ModelReport {
        var start = Date()
        var first = await check(model: model, client: client, key: key)
        first.latency = Date().timeIntervalSince(start)
        if first.green { return first }
        try? await Task.sleep(nanoseconds: 2_000_000_000)
        start = Date()
        var second = await check(model: model, client: client, key: key)
        second.latency = Date().timeIntervalSince(start)
        if second.green {
            second.notes.append("passed on retry")
        }
        return second
    }

    private static func check(model: AIModel, client: any ProviderClient, key: String) async -> ModelReport {
        var report = ModelReport(model: model)
        let start = Date()

        var parameters = ChatParameters(maxTokens: 300)
        if model.parameterSupport.reasoningEffort { parameters.reasoningEffort = "low" }
        if model.parameterSupport.thinkingToggle { parameters.thinkingEnabled = false }

        // ---- Step 1: the model must call list_files -------------------------
        let step1 = ChatRequest(
            model: model,
            systemPrompt: systemPrompt,
            messages: [Message(role: .user, text: userPrompt)],
            parameters: parameters,
            tools: [listFilesTool],
            toolChoice: .auto
        )
        let outcome1: StreamOutcome
        do {
            outcome1 = try await collectPatiently(client, step1, apiKey: key, deadline: perModelTimeout)
        } catch {
            report.notes.append("step 1: \(describe(error))")
            return report
        }

        let calls = outcome1.toolCalls
        if calls.isEmpty {
            report.notes.append("step 1: no tool call emitted (stop=\(outcome1.rawReason ?? "nil"))"
                + (outcome1.text.isEmpty ? "" : "; text: “\(truncate(outcome1.text, 80))”"))
            if outcome1.stop == .maxTokens {
                report.notes.append("token-starved at 300 max_tokens")
            }
            return report
        }
        var callProblems: [String] = []
        if !calls.contains(where: { $0.name == listFilesTool.name }) {
            callProblems.append("wrong tool name “\(calls.map(\.name).joined(separator: ","))”")
        }
        for call in calls where call.argumentsDictionary == nil {
            callProblems.append("arguments not valid JSON object: \(truncate(call.argumentsJSON, 80))")
        }
        if outcome1.stop != .toolUse {
            callProblems.append("stop reason not toolUse (raw=\(outcome1.rawReason ?? "nil"))")
        }
        report.toolCall = callProblems.isEmpty
        report.notes.append(contentsOf: callProblems.map { "step 1: \($0)" })
        if calls.count > 1 {
            report.notes.append("emitted \(calls.count) parallel calls")
        }
        guard report.toolCall else { return report }

        // ---- Step 2: thread the canned tool_result back ----------------------
        let assistantTurn = Message.assistantToolCalls(calls, text: outcome1.text)
        let results = calls.map { ToolResult(toolCallID: $0.id, content: cannedListing) }
        let step2 = ChatRequest(
            model: model,
            systemPrompt: systemPrompt,
            messages: [
                Message(role: .user, text: userPrompt),
                assistantTurn,
                Message.toolResultsMessage(results),
            ],
            parameters: parameters,
            tools: [listFilesTool],
            toolChoice: .auto
        )
        let remaining = max(15, perModelTimeout - Date().timeIntervalSince(start))
        let outcome2: StreamOutcome
        do {
            outcome2 = try await collectPatiently(client, step2, apiKey: key, deadline: remaining)
        } catch {
            report.notes.append("step 2: \(describe(error))")
            return report
        }

        var threadProblems: [String] = []
        let answer = outcome2.text.trimmingCharacters(in: .whitespacesAndNewlines)
        if !outcome2.toolCalls.isEmpty {
            threadProblems.append("called the tool again after receiving the result")
        }
        if answer.isEmpty {
            threadProblems.append("empty final answer (stop=\(outcome2.rawReason ?? "nil"))")
        } else if !expectedNames.contains(where: { answer.localizedCaseInsensitiveContains($0) }) {
            threadProblems.append("final answer names none of the files: “\(truncate(answer, 80))”")
        }
        if outcome2.stop != .endTurn {
            threadProblems.append("stop reason not endTurn (raw=\(outcome2.rawReason ?? "nil"))")
        }
        report.threading = threadProblems.isEmpty
        report.notes.append(contentsOf: threadProblems.map { "step 2: \($0)" })

        // ---- Streaming: deltas must have arrived incrementally ---------------
        let deltas = outcome1.textDeltas + outcome1.argDeltas + outcome2.textDeltas
        report.streaming = deltas > 1
        if !report.streaming {
            report.notes.append("streaming: only \(deltas) delta event(s) — arrived as a single blob")
        }
        return report
    }

    // MARK: - Stream consumption

    /// `collect` with rate-limit patience: strict per-model RPM tiers and
    /// intermittent capacity waves (Moonshot's kimi-k3 429s with
    /// `engine_overloaded_error` on ~half of requests) clear after a short
    /// wait — retry up to twice with growing pauses. Other errors propagate.
    private static func collectPatiently(
        _ client: any ProviderClient,
        _ request: ChatRequest,
        apiKey: String,
        deadline: TimeInterval
    ) async throws -> StreamOutcome {
        var wait: TimeInterval = 20
        for _ in 0..<2 {
            do {
                return try await collect(client, request, apiKey: apiKey, deadline: deadline)
            } catch ProviderError.rateLimited(_, let retryAfter) {
                try await Task.sleep(nanoseconds: UInt64(max(retryAfter ?? wait, wait) * 1_000_000_000))
                wait *= 2
            }
        }
        return try await collect(client, request, apiKey: apiKey, deadline: deadline)
    }

    /// Consumes one streamed response fully, racing a deadline. Any thrown
    /// provider/network error propagates; a deadline hit throws VerifyTimeout.
    private static func collect(
        _ client: any ProviderClient,
        _ request: ChatRequest,
        apiKey: String,
        deadline: TimeInterval
    ) async throws -> StreamOutcome {
        try await withThrowingTaskGroup(of: StreamOutcome?.self) { group in
            group.addTask {
                var out = StreamOutcome()
                for try await event in client.streamChat(request, apiKey: apiKey) {
                    switch event {
                    case .textDelta(let delta):
                        out.text += delta
                        out.textDeltas += 1
                    case .reasoningDelta:
                        out.reasoningDeltas += 1
                    case .toolCallStarted:
                        out.toolCallStarts += 1
                    case .toolCallArgumentsDelta:
                        out.argDeltas += 1
                    case .toolCalls(let calls):
                        out.toolCalls = calls
                    case .finished(let reason, let stop):
                        out.rawReason = reason
                        out.stop = stop
                    case .citations, .usage:
                        break
                    }
                }
                return out
            }
            group.addTask {
                try await Task.sleep(nanoseconds: UInt64(deadline * 1_000_000_000))
                return nil
            }
            guard let first = try await group.next(), let outcome = first else {
                group.cancelAll()
                throw VerifyTimeout()
            }
            group.cancelAll()
            return outcome
        }
    }

    // MARK: - Reporting

    private static func printLine(_ report: ModelReport) {
        let marks = [
            report.toolCall ? "tool✅" : "tool❌",
            report.threading ? "thread✅" : "thread❌",
            report.streaming ? "stream✅" : "stream❌",
        ].joined(separator: " ")
        let latency = report.latency.map { String(format: "%.1fs", $0) } ?? "—"
        var line = "  \(report.green ? "✅" : "❌") \(report.model.id)  \(marks)  \(latency)"
        if !report.notes.isEmpty {
            line += "\n      · " + report.notes.joined(separator: "\n      · ")
        }
        print(line)
    }

    private static func writeReport(reports: [ModelReport], keyless: [ProviderID]) {
        var lines: [String] = []
        lines.append("<!--")
        lines.append("  VERIFICATION.md")
        lines.append("  Zyquo Agent")
        lines.append("  Author: Simon-Pierre Boucher")
        lines.append("  Mail: contact@spboucher.ai")
        lines.append("-->")
        lines.append("")
        lines.append("# Provider Tool-Calling Verification — \(ISO8601DateFormatter().string(from: Date()))")
        lines.append("")
        lines.append("Phase 7.1 live sweep (`ZyquoAgent --verify`): every agent-capable model receives")
        lines.append("the `list_files` tool schema, must emit a valid tool call (streamed), consume the")
        lines.append("threaded `tool_result`, and produce a final answer naming the listed files.")
        lines.append("")
        let green = reports.filter(\.green).count
        lines.append("**\(reports.count) models tested · \(green) green · \(reports.count - green) failed**")
        if !keyless.isEmpty {
            lines.append("")
            lines.append("Providers skipped (no key available): \(keyless.map(\.displayName).joined(separator: ", "))")
        }
        lines.append("")
        lines.append("| Provider | Model | Tool call | Threading | Streaming | Latency | Notes |")
        lines.append("|---|---|---|---|---|---|---|")
        for report in reports {
            let latency = report.latency.map { String(format: "%.1fs", $0) } ?? "—"
            let notes = report.notes.joined(separator: "; ")
                .replacingOccurrences(of: "|", with: "\\|")
                .replacingOccurrences(of: "\n", with: " ")
            lines.append("| \(report.model.provider.displayName) | `\(report.model.id)` "
                + "| \(report.toolCall ? "✅" : "❌") | \(report.threading ? "✅" : "❌") "
                + "| \(report.streaming ? "✅" : "❌") | \(latency) | \(notes) |")
        }
        lines.append("")
        let url = URL(fileURLWithPath: FileManager.default.currentDirectoryPath)
            .appendingPathComponent("docs/VERIFICATION.md")
        try? lines.joined(separator: "\n").data(using: .utf8)?.write(to: url)
    }

    // MARK: - Helpers

    /// Gentle pacing between models of one provider (Cerebras free-tier rate
    /// limits are the strictest; everyone else gets a polite second).
    private static func interCallDelay(for provider: ProviderID) -> UInt64 {
        switch provider {
        case .cerebras: return 10_000_000_000
        case .mistral: return 2_000_000_000
        default: return 1_200_000_000
        }
    }

    private static func describe(_ error: Error) -> String {
        if error is VerifyTimeout { return "timed out after \(Int(perModelTimeout))s" }
        if let provider = error as? ProviderError {
            var text = provider.errorDescription ?? "\(provider)"
            if case .badRequest(_, let message) = provider {
                let lowered = (message ?? "").lowercased()
                if lowered.contains("model") &&
                    (lowered.contains("not") || lowered.contains("exist") || lowered.contains("found")
                     || lowered.contains("invalid") || lowered.contains("decommission")
                     || lowered.contains("deprecat") || lowered.contains("access")) {
                    text = "model id rejected by API — " + text
                }
            }
            return truncate(text, 220)
        }
        return truncate(error.localizedDescription, 220)
    }

    private static func truncate(_ text: String, _ limit: Int) -> String {
        let flat = text.replacingOccurrences(of: "\n", with: " ")
        return flat.count > limit ? String(flat.prefix(limit)) + "…" : flat
    }

    private static func value(after flag: String, in arguments: [String]) -> String? {
        guard let index = arguments.firstIndex(of: flag), index + 1 < arguments.count else { return nil }
        return arguments[index + 1]
    }
}