import Foundation

/// Decoding + answer-encoding for Claude Code's built-in `AskUserQuestion` tool.
///
/// `AskUserQuestion` is not a permission gate — the CLI calls it to *collect* the
/// user's choices and expects the answers back inside the permission reply's
/// `updatedInput`. If we reply `allow` with the input unchanged, the CLI records
/// no selection and reports "the user did not answer the questions". So we parse
/// the questions, surface a picker (see `AskUserQuestionBar`), and fold the
/// selections into `updatedInput` here via ``updatedInput(original:selections:)``.
///
/// Wire format (Agent SDK "Handle approvals and user input"): `updatedInput`
/// carries the original `questions` array unchanged plus an `answers` object whose
/// keys are the exact question text and whose values are the chosen option
/// `label`(s) — a string for single-select, an array for `multiSelect`.
public enum AskUserQuestion {
    public static let toolName = "AskUserQuestion"

    public struct Option: Sendable, Equatable, Identifiable {
        public let label: String
        public let description: String?
        public var id: String { label }
    }

    public struct Question: Sendable, Equatable, Identifiable {
        public let header: String?
        public let question: String
        public let multiSelect: Bool
        public let options: [Option]
        public var id: String { question }
    }

    /// Parse the tool input into questions, or `nil` if the payload doesn't carry
    /// a `questions` array (i.e. it isn't an AskUserQuestion call we understand).
    public static func questions(from input: JSONValue) -> [Question]? {
        guard let raw = input["questions"]?.arrayValue, !raw.isEmpty else { return nil }
        return raw.map { q in
            Question(
                header: q["header"]?.stringValue,
                question: q["question"]?.stringValue ?? "",
                multiSelect: q["multiSelect"]?.boolValue ?? false,
                options: (q["options"]?.arrayValue ?? []).map { o in
                    Option(
                        label: o["label"]?.stringValue ?? "",
                        description: o["description"]?.stringValue)
                })
        }
    }

    /// Build the `updatedInput` payload that carries the user's selections back to
    /// the CLI. `selections` maps each question's text to the chosen label(s);
    /// free-text answers are passed through verbatim as their own label.
    public static func updatedInput(original: JSONValue, selections: [String: [String]])
        -> JSONValue
    {
        var answers: [String: JSONValue] = [:]
        for question in questions(from: original) ?? [] {
            let picked = selections[question.question] ?? []
            answers[question.question] =
                question.multiSelect
                ? .array(picked.map(JSONValue.string))
                : .string(picked.first ?? "")
        }
        var object = original.objectValue ?? [:]
        object["answers"] = .object(answers)
        return .object(object)
    }

    /// Recover the user's recorded choices from the tool *result* the CLI emits once the
    /// questions are answered. The streamed/persisted `tool_use` input carries only the
    /// questions — the selections we fold into `updatedInput` never reappear there — so the
    /// answer for each question lives in the result text, which has the shape:
    ///
    ///     Your questions have been answered: "<question>"="<answer>", "<q2>"="<a2>". You can now continue…
    ///
    /// For each known question we read the answer between the `="…"` quotes that follow its
    /// exact text (matching the literal text, so punctuation inside a question doesn't trip a
    /// pattern). A `multiSelect` answer is several labels the CLI joins with a bare comma and
    /// no space (`"A,B"`), so it's split back into the picked labels — matched against the
    /// question's known options, so a label that itself contains a comma stays intact.
    /// Questions with no match — or an empty answer — are omitted, so an empty map means "no
    /// answers found" (still pending, or an unexpected result shape).
    public static func answers(in resultText: String, for questions: [Question]) -> [String: [String]] {
        var out: [String: [String]] = [:]
        for question in questions {
            guard let raw = answer(in: resultText, question: question.question), !raw.isEmpty
            else { continue }
            out[question.question] = selectedLabels(from: raw, options: question.options.map(\.label))
        }
        return out
    }

    /// The answer recorded for `question` in `resultText`, or nil when absent: locate
    /// `"<question>"="` and read up to the *closing* quote of the value.
    ///
    /// The CLI echoes a free-text "Other" answer verbatim, so the value can itself contain a
    /// `"` (`Use the "fast" path`). Stopping at the first inner quote would cut the answer off,
    /// so we take the closing quote to be the one followed by the CLI's separator (`, ` before
    /// the next answer) or terminator (`.` before the trailing sentence, or end of string) —
    /// the quotes that wrap an option label never appear in that position.
    private static func answer(in resultText: String, question: String) -> String? {
        guard let opened = resultText.range(of: "\"\(question)\"=\"") else { return nil }
        let rest = resultText[opened.upperBound...]
        var searchStart = rest.startIndex
        while let close = rest[searchStart...].firstIndex(of: "\"") {
            let tail = rest[rest.index(after: close)...]
            if tail.isEmpty || tail.hasPrefix(", ") || tail.hasPrefix(".") {
                return String(rest[..<close])
            }
            searchStart = rest.index(after: close)
        }
        return nil
    }

    /// Split a recorded answer back into the individual option labels the user picked. The CLI
    /// joins a multi-select answer with a bare comma (`"A,B"`, occasionally `"A, B"`), so we
    /// peel off known option labels left-to-right — preferring the longest match, and tolerating
    /// an optional space after the comma — which keeps a label that itself contains a comma
    /// whole. Anything that doesn't decompose cleanly (a free-text "Other", or an unexpected
    /// shape) is returned verbatim as a single answer rather than chopped at the wrong commas.
    static func selectedLabels(from raw: String, options: [String]) -> [String] {
        if options.contains(raw) { return [raw] }  // exact single pick (incl. a label with commas)
        var remaining = Substring(raw)
        var picked: [String] = []
        while !remaining.isEmpty {
            let match = options
                .filter { remaining.hasPrefix($0) }
                .filter { label in  // the label must end at a comma or the string's end
                    let after = remaining.index(remaining.startIndex, offsetBy: label.count)
                    return after == remaining.endIndex || remaining[after] == ","
                }
                .max { $0.count < $1.count }
            guard let match else { return [raw] }  // can't decompose cleanly — keep verbatim
            picked.append(match)
            remaining = remaining.dropFirst(match.count)
            if remaining.first == "," { remaining = remaining.dropFirst() }
            while remaining.first == " " { remaining = remaining.dropFirst() }
        }
        return picked
    }
}
