import Foundation

/// Strips the agentic scaffolding Apple Foundation Models wraps free-form replies in.
///
/// The model shipped in macOS 27 beta 4 (build 26A5388g) answers a plain `respond(to:)` as if it
/// were driving a tool-calling harness: instead of the bare line the instructions ask for, it
/// returns a JSON object (`{"task": "Explore Repository Contents"}`), a fenced JSON block
/// (```` ```json\n{"summary": "…"}\n``` ````), a pseudo tool-call (`tool_call: {name: "search_terms",
/// …}`), `[start_x]`/`[end_x]` section markers, or a persona preamble ("I'm an Apple foundation
/// model.") ahead of the real answer. Guided generation (`@Generable`) is unaffected — the schema
/// still binds — so only the free-form paths need this.
///
/// Every caller downstream takes the reply's first non-empty line (`HeuristicTitle.firstModelLine`,
/// `sanitizeModelTitle`, the keyword/triage splitters), so unwrapped scaffolding would land
/// verbatim in chat titles and collapsed summaries — "```json" as a Bash summary, "I'm an Apple
/// foundation model." as a session name. This runs once at each provider's text choke point, before
/// any of that parsing, and yields "" when nothing but scaffolding was returned so the caller falls
/// back to its deterministic heuristic rather than showing the wrapper.
public enum ModelOutput {
    /// Keys whose value carries the real answer, most-specific first. `tool_calls`/`parameters`
    /// and friends are deliberately absent — they hold the harness's own bookkeeping.
    private static let payloadKeys = [
        "summary", "title", "text", "line", "answer", "response", "output",
        "result", "content", "message", "description", "task", "value",
    ]

    /// Keys that never hold prose, so their value must not be mistaken for the answer. `status`
    /// earns its place from a captured run: the seed answered a Bash-summary prompt with
    /// `{"status": "completed"}`, and taking that as the summary would label the collapsed row
    /// "completed" instead of falling back to the command itself.
    private static let ignoredKeys: Set<String> = [
        "tool_calls", "tool_call", "toolcalls", "function_call", "functions",
        "parameters", "params", "args", "arguments", "name", "type", "role", "id",
        "status", "state", "action", "command", "query", "kind", "tool", "term", "tokens",
    ]

    /// `unwrap`, folded to nil when the reply held no usable text — the same "no result" every
    /// caller already handles by falling back to its heuristic.
    public static func usableText(_ raw: String?) -> String? {
        guard let raw else { return nil }
        let text = unwrap(raw)
        return text.isEmpty ? nil : text
    }

    /// The model's actual answer with any harness scaffolding removed. Returns "" when the reply
    /// was scaffolding all the way down. Plain replies pass through untouched (only trimmed), so
    /// this is safe to run unconditionally on every free-form generation.
    public static func unwrap(_ raw: String) -> String {
        let defenced = fencedBody(of: raw)
        // A JSON envelope carries its answer in a value, not on a line, so decode it before any
        // line-level scrubbing — which would otherwise leave the raw braces and key/value debris
        // as the "answer". Once the text is known to BE an envelope, its prose is the only thing
        // worth keeping: an envelope holding nothing but machine slots unwraps to "" (caller falls
        // back) rather than leaking the JSON downstream.
        if let object = jsonObject(of: defenced) { return prose(in: object) ?? "" }
        return scrubScaffoldingLines(defenced)
    }

    /// The contents of the first ``` fenced block, or the input unchanged when there is none.
    /// Every prompt these back explicitly forbids code fences, so a fence is always the harness's
    /// wrapper and its body is always the answer — including the preamble the fence follows.
    private static func fencedBody(of raw: String) -> String {
        let lines = raw.components(separatedBy: .newlines)
        guard let open = lines.firstIndex(where: { $0.trimmingCharacters(in: .whitespaces).hasPrefix("```") })
        else { return raw.trimmingCharacters(in: .whitespacesAndNewlines) }
        let rest = lines[(open + 1)...]
        let close = rest.firstIndex(where: { $0.trimmingCharacters(in: .whitespaces).hasPrefix("```") }) ?? lines.endIndex
        return lines[(open + 1)..<close].joined(separator: "\n")
            .trimmingCharacters(in: .whitespacesAndNewlines)
    }

    /// The decoded JSON envelope, or nil when the text isn't one (the overwhelmingly common case —
    /// a plain reply never starts with a brace).
    private static func jsonObject(of raw: String) -> Any? {
        let trimmed = raw.trimmingCharacters(in: .whitespacesAndNewlines)
        guard trimmed.hasPrefix("{") || trimmed.hasPrefix("["),
              let data = trimmed.data(using: .utf8)
        else { return nil }
        return try? JSONSerialization.jsonObject(with: data, options: [.fragmentsAllowed])
    }

    /// Depth-first search for the answer inside a decoded JSON envelope: the first payload key
    /// that holds usable prose, else any non-bookkeeping value that does.
    private static func prose(in object: Any) -> String? {
        if let string = object as? String { return isProse(string) ? clean(string) : nil }
        if let array = object as? [Any] { return array.lazy.compactMap(prose).first }
        guard let dict = object as? [String: Any] else { return nil }
        // Case-insensitive lookup: the model varies "Summary"/"summary" between generations.
        let byLowerKey = Dictionary(dict.map { ($0.key.lowercased(), $0.value) }, uniquingKeysWith: { a, _ in a })
        for key in payloadKeys {
            if let value = byLowerKey[key], let found = prose(in: value) { return found }
        }
        // No known key — fall back to any remaining value, preferring the longest prose so a
        // multi-word answer wins over an incidental one-word field.
        return byLowerKey
            .filter { !ignoredKeys.contains($0.key) }
            .compactMap { prose(in: $0.value) }
            .max(by: { $0.count < $1.count })
    }

    /// Whether a JSON string value is the model's answer rather than a machine token. Natural
    /// language has spaces; a spaceless `snake_case` value ("login_button_response") is the
    /// harness naming a slot, and reads as gibberish in a title, so it is rejected — while a
    /// legitimate one-word answer ("AWAITING", "CONSISTENT") is kept.
    private static func isProse(_ value: String) -> Bool {
        let trimmed = value.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !trimmed.isEmpty else { return false }
        let hasSpace = trimmed.contains(where: { $0 == " " || $0 == "\t" || $0.isNewline })
        return hasSpace || !trimmed.contains("_")
    }

    private static func clean(_ value: String) -> String {
        value.trimmingCharacters(in: .whitespacesAndNewlines)
    }

    /// Drops whole lines that are harness scaffolding, keeping everything else in order. Operates
    /// per line rather than on the reply as a whole so a preamble in front of a real answer costs
    /// only the preamble — the keyword list after "Here are 5 …:" survives.
    private static func scrubScaffoldingLines(_ raw: String) -> String {
        let kept = raw.components(separatedBy: .newlines).filter { !isScaffolding($0) }
        return kept.joined(separator: "\n").trimmingCharacters(in: .whitespacesAndNewlines)
    }

    private static func isScaffolding(_ line: String) -> Bool {
        let trimmed = line.trimmingCharacters(in: .whitespaces)
        guard !trimmed.isEmpty else { return false }
        let lowered = trimmed.lowercased()
        // `tool_call: {…}` / `Tool call: {…}` / `function_call = {…}`.
        if lowered.range(of: "^(tool|function)[ _]?calls?\\s*[:=]", options: .regularExpression) != nil { return true }
        // The tagged variant, where the call is narrated after a bracketed label rather than in a
        // literal — `[Tool call] Retrieve search results for: "login"`, `[Tool Call 2] …`. The whole
        // line is the harness describing a call it never made, so none of it is the answer.
        if lowered.range(of: "^\\[\\s*(tool|function)[ _]?calls?(\\s+\\d+)?\\s*\\]", options: .regularExpression) != nil { return true }
        // A line that is nothing but a bracketed note is the harness talking about itself, never
        // the answer: `[start_coding_session]`, `[end_coding_session]`, `[No tools needed]`.
        if lowered.range(of: "^\\[[^\\]]*\\]$", options: .regularExpression) != nil { return true }
        // A bare label line the seed prefixes its real answer with ("SUMMARY:"), which would
        // otherwise become the summary itself once `firstModelLine` takes line one. Anchored to
        // all-caps + colon so the requested "**Done:**" Markdown headings are untouched.
        if trimmed.range(of: "^[A-Z][A-Z ]*:$", options: .regularExpression) != nil { return true }
        // A leftover bare fence, when the reply opened one but never closed it.
        if trimmed.hasPrefix("```") { return true }
        // "I'm an Apple foundation model." / "I am a large language model." — the model
        // introducing itself before (or instead of) answering.
        if lowered.range(of: "^i(\\s+am|'m|’m)\\s+(a|an|the)\\b.*\\b(language model|foundation model|ai model|assistant|apple)\\b",
                         options: .regularExpression) != nil { return true }
        // "Here are 5 lowercase search keywords:" — a lead-in to the answer on the next line.
        if lowered.range(of: "^(sure[,.!]?\\s*)?here (are|is)\\b.*:$", options: .regularExpression) != nil { return true }
        return false
    }
}
