import Foundation

/// The Intelligence slider's discrete stops. Level is *ambition/budget* — how much
/// capability (and therefore cost) a prompt deserves — not an API effort: the router
/// (`IntelligenceRouter`) combines a level with the prompt's classified purpose to pick a
/// concrete (model, effort) pair, so the same level can mean "Sonnet at low" for a typo fix
/// and "Sol at high" for a backend feature. Five stops so the ladder reads like the familiar
/// effort range without being it.
public enum IntelligenceLevel: Int, CaseIterable, Codable, Sendable, Comparable, Identifiable {
    case quick = 1
    case light = 2
    case balanced = 3
    case deep = 4
    case max = 5

    public var id: Int { rawValue }

    /// The stop's short display name (slider caption, Settings picker, session info).
    public var displayName: String {
        switch self {
        case .quick: return "Quick"
        case .light: return "Light"
        case .balanced: return "Balanced"
        case .deep: return "Deep"
        case .max: return "Max"
        }
    }

    /// Position along the slider, 0 (Quick) … 1 (Max) — drives the intensity visuals
    /// (fill length, shimmer speed, glow) so they scale continuously with the level.
    public var fraction: Double { Double(rawValue - 1) / Double(IntelligenceLevel.allCases.count - 1) }

    /// The stop new chats start on: mid-tier work on a mid-tier model. Matches the
    /// `general` routing row's center cell, so an untouched slider is never a surprise bill.
    public static let fallback = IntelligenceLevel.balanced

    public static func < (lhs: IntelligenceLevel, rhs: IntelligenceLevel) -> Bool {
        lhs.rawValue < rhs.rawValue
    }
}

/// What a prompt is *for* — the routing key the purpose classifier produces. The first
/// eight cases are the classifier's fixed label set (mutually distinguishable by keyword
/// and structural signals); `general` is never offered to a model — it's the conservative
/// fallback bucket whose routing row is deliberately mid-tier, so a failed classification
/// can never mean "Fable at max" or "Haiku on a hard task".
public enum PromptPurpose: String, CaseIterable, Codable, Sendable {
    /// Architecture, design docs, migration strategy, multi-step project planning.
    case planning
    /// Server/API/data/algorithm/system implementation work.
    case backendImpl
    /// UI, views, styling, components, layout, animation work.
    case frontendImpl
    /// Typo, one-liner, config tweak, small contained bugfix.
    case quickFix
    /// Restructuring without behavior change (rename, extract, reorganize).
    case refactor
    /// Diagnosing a failure: crash, stack trace, regression, "why does this break".
    case debugging
    /// Reviewing/explaining/answering questions about existing code.
    case review
    /// Docs, commit messages, extraction, formatting — prose over code.
    case writing
    /// Fallback only — not a classifier label (see type comment).
    case general

    /// The labels the model-backed classifier may answer with (`general` excluded — it
    /// exists so "no idea" routes conservatively, not so a model can shrug).
    public static let classifiable: [PromptPurpose] = [
        .planning, .backendImpl, .frontendImpl, .quickFix,
        .refactor, .debugging, .review, .writing,
    ]

    /// Human-readable noun phrase for the resolved-chip tooltip and session info.
    public var displayName: String {
        switch self {
        case .planning: return "Planning"
        case .backendImpl: return "Backend implementation"
        case .frontendImpl: return "Frontend implementation"
        case .quickFix: return "Quick fix"
        case .refactor: return "Refactor"
        case .debugging: return "Debugging"
        case .review: return "Review / explain"
        case .writing: return "Writing / docs"
        case .general: return "General"
        }
    }

    /// The single-word uppercase token the model layers emit (`PLANNING`, `BACKEND`, …) —
    /// shared by the AFM guided-generation guide and the delegated `claude -p` template so
    /// the two prompt twins can't drift (see `IntelligenceDelegate`).
    public var classifierToken: String {
        switch self {
        case .planning: return "PLANNING"
        case .backendImpl: return "BACKEND"
        case .frontendImpl: return "FRONTEND"
        case .quickFix: return "QUICKFIX"
        case .refactor: return "REFACTOR"
        case .debugging: return "DEBUGGING"
        case .review: return "REVIEW"
        case .writing: return "WRITING"
        case .general: return "GENERAL"
        }
    }

    /// Parse a model layer's token back into a purpose (forgiving about case/whitespace).
    /// `GENERAL`/unknown → nil, so callers fall back explicitly rather than silently.
    public static func fromClassifierToken(_ raw: String) -> PromptPurpose? {
        let token = raw.trimmingCharacters(in: .whitespacesAndNewlines).uppercased()
        return classifiable.first { $0.classifierToken == token }
    }
}

/// One classification outcome: the label, how sure the classifying layer was, which layer
/// produced it, and a short human-readable reason (surfaced in the resolved chip's tooltip
/// and the session's routing note).
public struct PurposeVerdict: Sendable, Equatable {
    public enum Confidence: String, Sendable, Comparable {
        case low, medium, high
        private var rank: Int { self == .low ? 0 : self == .medium ? 1 : 2 }
        public static func < (lhs: Confidence, rhs: Confidence) -> Bool { lhs.rank < rhs.rank }
    }

    /// Which layer answered — recorded so the routing note can say "keyword match" vs
    /// "on-device model" vs "haiku", and so tests can assert the layering.
    public enum Source: String, Sendable {
        /// Deterministic keyword/shape analysis (`HeuristicPurposeClassifier`).
        case heuristic
        /// The bundled, task-specific Core ML purpose classifier.
        case custom
        /// Apple Foundation Models guided generation.
        case afm
        /// One-shot agent CLI escalation (`claude -p` on a small SKU).
        case agent
        /// Every layer failed or declined — the conservative `general` bucket.
        case fallback
    }

    public var purpose: PromptPurpose
    public var confidence: Confidence
    public var source: Source
    public var reason: String

    public init(purpose: PromptPurpose, confidence: Confidence, source: Source, reason: String) {
        self.purpose = purpose
        self.confidence = confidence
        self.source = source
        self.reason = reason
    }

    /// The verdict when nothing could classify the prompt: `general` routes mid-tier.
    public static let general = PurposeVerdict(
        purpose: .general, confidence: .low, source: .fallback,
        reason: "No clear purpose signal — using the balanced general lane")

    /// Parse a model layer's strict `LABEL|CONFIDENCE|reason` line (the delegated template's
    /// output contract — see `IntelligenceDelegate`). Nil when the label is off-menu or the
    /// line is garbled, so callers fall back to the heuristic verdict rather than guessing.
    public static func fromModelLine(_ raw: String, source: Source) -> PurposeVerdict? {
        let line = HeuristicTitle.firstModelLine(of: raw)
        guard !line.isEmpty else { return nil }
        let parts = line.split(separator: "|", maxSplits: 2, omittingEmptySubsequences: false)
            .map { $0.trimmingCharacters(in: .whitespaces) }
        guard let purpose = PromptPurpose.fromClassifierToken(parts[0]) else { return nil }
        let confidence = parts.count > 1
            ? (Confidence(rawValue: parts[1].lowercased()) ?? .low) : .low
        return PurposeVerdict(
            purpose: purpose, confidence: confidence, source: source,
            reason: parts.count > 2 ? parts[2] : "")
    }
}

/// Decides — without a model — what a prompt is for, from keyword buckets plus structural
/// signals (pasted-error shape, question form, size). The fallback when Intelligence is
/// off/unavailable, the fast path that skips the model entirely on unambiguous prompts,
/// and the guardrail a low-confidence model verdict is checked against.
///
/// Bias notes (mirrors `HeuristicTurnClassifier`'s discipline): a HIGH-confidence verdict
/// here is accepted outright and skips the model, so `high` requires *exactly one* bucket
/// with at least two independent hits and no structural contradiction — a false HIGH
/// misroutes a whole chat, while a false MEDIUM/LOW merely costs one model call. Ties and
/// thin signals therefore lean toward `low`, which routes through the model layers.
public enum HeuristicPurposeClassifier {
    /// Keyword buckets per purpose. Multi-word phrases are matched as substrings of the
    /// lowered prompt; single words are matched against the tokenized word set, so "api"
    /// can't fire inside "rapid". Listed roughly most- to least-diagnostic.
    static let buckets: [(purpose: PromptPurpose, words: [String], phrases: [String])] = [
        (.planning,
         ["plan", "architecture", "architect", "roadmap", "rfc", "strategy", "spec", "design-doc"],
         ["design doc", "come up with a plan", "migration plan", "how should we structure",
          "high-level design", "break down", "project plan", "implementation plan"]),
        (.backendImpl,
         ["endpoint", "api", "schema", "database", "migration", "server", "backend", "queue",
          "cache", "webhook", "auth", "grpc", "sql", "daemon", "cli", "algorithm", "parser"],
         ["rest api", "background job", "rate limit", "data model", "business logic"]),
        (.frontendImpl,
         ["ui", "view", "component", "css", "swiftui", "layout", "animation", "button",
          "toolbar", "sidebar", "frontend", "styling", "theme", "react", "html", "screen"],
         ["dark mode", "design system", "user interface", "hover state", "empty state"]),
        (.quickFix,
         ["typo", "tweak", "bump", "one-liner", "rename"],
         ["quick fix", "small fix", "one line", "fix typo", "fix the typo",
          "update the version", "change the default", "flip the flag"]),
        (.refactor,
         ["refactor", "restructure", "reorganize", "extract", "dedupe", "cleanup", "simplify",
          "consolidate", "decouple", "modularize"],
         ["clean up", "split out", "pull out", "tidy up", "reduce duplication"]),
        (.debugging,
         ["crash", "bug", "regression", "stacktrace", "traceback", "exception", "segfault",
          "hang", "leak", "flaky", "broken", "error", "fails", "failing"],
         ["stack trace", "doesn't work", "does not work", "why does", "root cause",
          "stopped working", "reproduce the"]),
        (.review,
         ["review", "explain", "understand", "audit", "compare", "evaluate", "assess"],
         ["what does", "how does", "walk me through", "code review", "look over",
          "is it safe", "any issues", "what's the difference"]),
        (.writing,
         ["readme", "docs", "documentation", "changelog", "comment", "docstring", "blog",
          "summarize", "summary", "translate", "reword", "proofread"],
         ["commit message", "release notes", "write up", "doc comment", "pull request description"]),
    ]

    /// Classify `prompt`. Deterministic and cheap (<1 ms) — safe to call synchronously at
    /// send time when the model layers haven't answered yet.
    public static func classify(_ prompt: String) -> PurposeVerdict {
        let trimmed = prompt.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !trimmed.isEmpty else { return .general }
        // Classification signals live in how the user frames the ask, which fronts the
        // prompt; a pasted log/diff tail would swamp the buckets, so weigh the head.
        let head = String(trimmed.prefix(2000)).lowercased()
        let words = Set(
            head.split(whereSeparator: { !$0.isLetter && !$0.isNumber && $0 != "-" })
                .map(String.init))

        var hits: [PromptPurpose: Int] = [:]
        var matched: [PromptPurpose: [String]] = [:]
        for bucket in buckets {
            var count = 0
            var terms: [String] = []
            for word in bucket.words where words.contains(word) {
                count += 1
                terms.append(word)
            }
            for phrase in bucket.phrases where head.contains(phrase) {
                count += 1
                terms.append(phrase)
            }
            if count > 0 {
                hits[bucket.purpose] = count
                matched[bucket.purpose] = terms
            }
        }

        // Structural signals — they nudge or veto, never outweigh a strong bucket alone.
        let wordCount = trimmed.split(whereSeparator: { $0 == " " || $0.isNewline }).count
        let looksLikePastedError =
            head.contains("error:") || head.contains("exception") || head.contains("traceback")
            || head.contains("stack trace") || head.contains("assertion failed")
            || head.range(of: #"at .+:\d+"#, options: .regularExpression) != nil
        if looksLikePastedError { hits[.debugging, default: 0] += 2 }
        // A short imperative micro-edit reads as a quick fix even without bucket words.
        if wordCount <= 8, hits.isEmpty || (hits.count == 1 && hits.keys.first == .quickFix) {
            hits[.quickFix, default: 0] += 1
        }
        // A pure question with no implementation signal reads as review/explain.
        if trimmed.hasSuffix("?"), hits[.backendImpl] == nil, hits[.frontendImpl] == nil,
            hits[.debugging] == nil {
            hits[.review, default: 0] += 1
        }

        let ranked = hits.sorted { lhs, rhs in
            lhs.value != rhs.value ? lhs.value > rhs.value : lhs.key.rawValue < rhs.key.rawValue
        }
        guard let leader = ranked.first else { return .general }
        let runnerUp = ranked.dropFirst().first?.value ?? 0
        let terms = matched[leader.key, default: []].prefix(3).joined(separator: ", ")
        let reason = terms.isEmpty
            ? "Prompt shape suggests \(leader.key.displayName.lowercased())"
            : "Matched \(terms)"

        // HIGH only on an unambiguous single leader: ≥2 hits and a ≥2 lead (see bias note).
        if leader.value >= 2, leader.value - runnerUp >= 2 {
            return PurposeVerdict(
                purpose: leader.key, confidence: .high, source: .heuristic, reason: reason)
        }
        if leader.value > runnerUp {
            return PurposeVerdict(
                purpose: leader.key, confidence: .medium, source: .heuristic, reason: reason)
        }
        return PurposeVerdict(
            purpose: leader.key, confidence: .low, source: .heuristic, reason: reason)
    }
}
