import Foundation
import NucleicCore
import Testing

@testable import NucleicApp

@Suite("Purpose ML classifier", .serialized)
struct PurposeMLClassifierTests {
    private struct Golden: Decodable {
        let prompt: String
        let normalizedPrompt: String
        let inputIds: [Int32]
        let attentionMask: [Int32]
        let tokenTypeIds: [Int32]
    }

    private struct LabeledPrompt: Decodable {
        let prompt: String
        let purpose: String
        let slice: String?
    }

    private struct FixedModel: PurposeModelClassifying {
        let verdict: PurposeVerdict?
        func classify(prompt: String) async -> PurposeVerdict? { verdict }
    }

    private actor CountingModel: PurposeModelClassifying {
        private(set) var calls = 0
        func classify(prompt: String) async -> PurposeVerdict? {
            calls += 1
            return PurposeVerdict(
                purpose: .planning, confidence: .high, source: .custom,
                reason: "should have been skipped")
        }
    }

    private struct FixedIntelligence: IntelligenceProviding {
        let verdict: PurposeVerdict
        func summarize(session: Session, events: [AgentEvent]) async -> String { "" }
        func sessionName(fromFirstMessage message: String) async -> String? { nil }
        func classifyPromptPurpose(prompt: String) async -> PurposeVerdict { verdict }
    }

    @Test func tokenizerMatchesExportGoldensExactly() throws {
        let vocabularyURL = try #require(PurposeClassifierResources.url(
            named: "vocab", extensions: ["txt"]))
        let goldensURL = try #require(Bundle.module.url(
            forResource: "tokenizer-goldens", withExtension: "json",
            subdirectory: "Fixtures"))
        let tokenizer = try PurposeWordPieceTokenizer(
            vocabularyText: String(contentsOf: vocabularyURL, encoding: .utf8))
        let goldens = try JSONDecoder().decode(
            [Golden].self, from: Data(contentsOf: goldensURL))

        #expect(goldens.count == 4)
        for golden in goldens {
            #expect(PurposeWordPieceTokenizer.normalizePrompt(golden.prompt) == golden.normalizedPrompt)
            let encoded = tokenizer.encode(golden.prompt)
            #expect(encoded.inputIDs == golden.inputIds)
            #expect(encoded.attentionMask == golden.attentionMask)
            #expect(encoded.tokenTypeIDs == golden.tokenTypeIds)
        }
    }

    @Test func calibratedLocalVerdictWinsBeforeAFM() async {
        let local = PurposeVerdict(
            purpose: .review, confidence: .medium, source: .custom,
            reason: "test local model")
        let afm = PurposeVerdict(
            purpose: .debugging, confidence: .high, source: .afm,
            reason: "should not win")
        let result = await PromptPurposeService.classifyLayered(
            "Please consider the requested behavior across the application",
            modelClassifier: FixedModel(verdict: local),
            intelligence: FixedIntelligence(verdict: afm),
            escalation: nil)
        #expect(result == local)
    }

    @Test func highHeuristicSkipsEveryModelLayer() async {
        let model = CountingModel()
        let result = await PromptPurposeService.classifyLayered(
            "Fix the typo and update the version in the README",
            modelClassifier: model,
            intelligence: FixedIntelligence(verdict: .general),
            escalation: nil)
        #expect(result.source == .heuristic)
        #expect(result.confidence == .high)
        #expect(result.purpose == .quickFix)
        #expect(await model.calls == 0)
    }

    @Test func lowLocalDisagreementFallsThroughToAFM() async {
        let local = PurposeVerdict(
            purpose: .quickFix, confidence: .low, source: .custom,
            reason: "uncertain local model")
        let afm = PurposeVerdict(
            purpose: .planning, confidence: .high, source: .afm,
            reason: "AFM resolved ambiguity")
        let result = await PromptPurposeService.classifyLayered(
            "Please consider these details and decide what should happen next",
            modelClassifier: FixedModel(verdict: local),
            intelligence: FixedIntelligence(verdict: afm),
            escalation: nil)
        #expect(result == afm)
    }

    @Test func lowLocalAgreementIsAccepted() async {
        let text = "Review the behavior and tell me whether it is safe"
        let heuristic = HeuristicPurposeClassifier.classify(text)
        #expect(heuristic.purpose == .review)
        let local = PurposeVerdict(
            purpose: .review, confidence: .low, source: .custom,
            reason: "guardrail agreement")
        let result = await PromptPurposeService.classifyLayered(
            text,
            modelClassifier: FixedModel(verdict: local),
            intelligence: FixedIntelligence(verdict: .general),
            escalation: nil)
        #expect(result == local)
    }

    @Test func bundledW8A8ModelLoadsAndPredicts() async {
        let defaults = UserDefaults.standard
        let oldEnabled = defaults.object(forKey: PurposeMLClassifier.enabledKey)
        let oldModel = defaults.object(forKey: PurposeMLClassifier.modelKey)
        defer {
            if let oldEnabled { defaults.set(oldEnabled, forKey: PurposeMLClassifier.enabledKey) }
            else { defaults.removeObject(forKey: PurposeMLClassifier.enabledKey) }
            if let oldModel { defaults.set(oldModel, forKey: PurposeMLClassifier.modelKey) }
            else { defaults.removeObject(forKey: PurposeMLClassifier.modelKey) }
        }
        defaults.set(true, forKey: PurposeMLClassifier.enabledKey)
        defaults.set("lite", forKey: PurposeMLClassifier.modelKey)

        let result = await PurposeMLClassifier().classify(
            prompt: "Diagnose why this request crashes and identify the root cause")
        #expect(result?.source == .custom)
        #expect(result?.reason.contains(PurposeMLClassifier.artifactVersion) == true)
    }

    /// Runs the shipped runtime—not PyTorch or ONNX—over the canonical frozen corpus. This is
    /// intentionally a normal Darwin test: a tokenizer or Core ML conversion drift must block a
    /// release, and the warm model evaluates the ~1.2K prompts in only a few seconds.
    @Test func bundledW8A8MeetsFrozenRuntimeGates() async throws {
        let defaults = UserDefaults.standard
        let oldEnabled = defaults.object(forKey: PurposeMLClassifier.enabledKey)
        let oldModel = defaults.object(forKey: PurposeMLClassifier.modelKey)
        defer {
            if let oldEnabled { defaults.set(oldEnabled, forKey: PurposeMLClassifier.enabledKey) }
            else { defaults.removeObject(forKey: PurposeMLClassifier.enabledKey) }
            if let oldModel { defaults.set(oldModel, forKey: PurposeMLClassifier.modelKey) }
            else { defaults.removeObject(forKey: PurposeMLClassifier.modelKey) }
        }
        defaults.set(true, forKey: PurposeMLClassifier.enabledKey)
        defaults.set("lite", forKey: PurposeMLClassifier.modelKey)

        let repository = URL(fileURLWithPath: #filePath)
            .deletingLastPathComponent() // NucleicAppTests
            .deletingLastPathComponent() // Tests
            .deletingLastPathComponent() // repository root
        let frozenURL = repository.appendingPathComponent(
            "ml/purpose-classifier/data/frozen-test-v1.jsonl")
        let fixturesURL = repository.appendingPathComponent(
            "Tests/NucleicCoreTests/Fixtures/purpose-prompts.json")
        let decoder = JSONDecoder()
        let frozenText = try String(contentsOf: frozenURL, encoding: .utf8)
        var records = try frozenText.split(separator: "\n").map {
            try decoder.decode(LabeledPrompt.self, from: Data($0.utf8))
        }
        let fixtures = try decoder.decode(
            [LabeledPrompt].self, from: Data(contentsOf: fixturesURL))
        records += fixtures.filter { $0.purpose != PromptPurpose.general.rawValue }.map {
            LabeledPrompt(prompt: $0.prompt, purpose: $0.purpose, slice: "shipped-fixture")
        }

        let classifier = PurposeMLClassifier()
        var scoredCorrect = 0
        var scoredCount = 0
        var actualByPurpose: [String: Int] = [:]
        var correctByPurpose: [String: Int] = [:]
        var vagueCount = 0
        var vagueLow = 0
        var hardCount = 0
        var hardCorrect = 0
        var fixtureCount = 0
        var fixtureCorrect = 0
        var latency: [Double] = []
        let hardSlices: Set<String> = ["boundary", "mixed", "pasted-context"]

        for record in records {
            let started = ProcessInfo.processInfo.systemUptime
            let verdict = try #require(await classifier.classify(prompt: record.prompt))
            latency.append((ProcessInfo.processInfo.systemUptime - started) * 1_000)
            if record.slice == "vague-eval" {
                vagueCount += 1
                if verdict.confidence == .low { vagueLow += 1 }
            } else {
                scoredCount += 1
                actualByPurpose[record.purpose, default: 0] += 1
                if verdict.purpose.rawValue == record.purpose {
                    scoredCorrect += 1
                    correctByPurpose[record.purpose, default: 0] += 1
                }
                if let slice = record.slice, hardSlices.contains(slice) {
                    hardCount += 1
                    if verdict.purpose.rawValue == record.purpose { hardCorrect += 1 }
                }
                if record.slice == "shipped-fixture" {
                    fixtureCount += 1
                    if verdict.purpose.rawValue == record.purpose { fixtureCorrect += 1 }
                }
            }
        }

        let accuracy = Double(scoredCorrect) / Double(scoredCount)
        let minimumRecall = actualByPurpose.map { purpose, count in
            Double(correctByPurpose[purpose, default: 0]) / Double(count)
        }.min() ?? 0
        let macroRecall = actualByPurpose.map { purpose, count in
            Double(correctByPurpose[purpose, default: 0]) / Double(count)
        }.reduce(0, +) / Double(actualByPurpose.count)
        let hardAccuracy = Double(hardCorrect) / Double(hardCount)
        let fixtureAccuracy = Double(fixtureCorrect) / Double(fixtureCount)
        let vagueLowRate = Double(vagueLow) / Double(vagueCount)
        let orderedLatency = latency.sorted()
        let p95 = orderedLatency[Int(Double(orderedLatency.count - 1) * 0.95)]
        print(String(
            format: "W8A8 frozen: accuracy=%.4f macro_recall=%.4f min_recall=%.4f "
                + "hard=%.4f fixtures=%.4f vague_low=%.4f p95=%.2fms",
            accuracy, macroRecall, minimumRecall, hardAccuracy, fixtureAccuracy,
            vagueLowRate, p95))

        #expect(scoredCount == 937)
        #expect(accuracy >= 0.95)
        #expect(minimumRecall >= 0.85)
        #expect(vagueLowRate >= 0.90)
        #expect(p95 <= 20)
    }
}
