import Foundation

/// Translates xAI Grok's **ACP** (`grok agent stdio`) `session/update` notifications into the
/// normalized `AgentEvent.Kind`s (BACKEND_PROTOCOL §3). The Agent Client Protocol is an open
/// standard, so unlike the previous inferred `streaming-json` vocab these shapes are pinned to
/// the published ACP schema:
///
/// `params = { sessionId, update: { sessionUpdate: "<variant>", … } }`, variants:
///   - `agent_message_chunk { content: ContentBlock, messageId? }`  → `assistantText(partial)`
///   - `agent_thought_chunk { content: ContentBlock, messageId? }`  → `thinking(partial)`
///   - `tool_call        { toolCallId, title, kind, status, rawInput, content, locations }`
///   - `tool_call_update { toolCallId, status, content, locations? }`
///   - `plan` / `available_commands` / `current_mode_update` / `config_option_update` → structural
///
/// ACP describes a tool by `kind` (read|edit|delete|move|search|execute|think|fetch|other) + a
/// human `title`, **not** a raw tool name. We map `kind` onto Nucleic's tool vocabulary
/// (`Bash`/`Edit`/`Read`/…) so the shared `RiskClassifier`, command summaries, and approval UI
/// work across backends — the cross-adapter symmetry the protocol is designed for (ADAPTERS §5).
/// `fileChange`s are synthesized from the tool call's authoritative `locations` (else `rawInput`),
/// emitted only on success; the authoritative diff still comes from `WorktreeManager`
/// (`emitsFileChangeEvents = false`). Lenient by contract: anything unrecognized → `.raw`
/// (the CLI-drift early-warning signal, OBSERVABILITY A.5).
public final class GrokACPDecoder {
    public struct Decoded: Sendable, Equatable {
        public let nativeType: String?
        public let kind: AgentEvent.Kind
        public init(nativeType: String?, kind: AgentEvent.Kind) {
            self.nativeType = nativeType
            self.kind = kind
        }
    }

    /// The raw ACP fields of an announced `tool_call` (merged across `tool_call_update`s). The
    /// backend recovers these by id when a `session/request_permission` carries only a minimal
    /// `ToolCallUpdate` (`toolCallId` alone — ACP permits omitting fields already sent).
    public struct AnnouncedToolCall: Sendable, Equatable {
        public let kind: String
        public let rawInput: JSONValue
        public let locations: [String]
    }

    /// Tool-call ids already announced via `toolCallStarted` (dedup across start/update).
    private var startedToolCallIDs: Set<String> = []
    /// The announced `ToolCall` per id, so a later `tool_call_update` can re-emit it as completed.
    private var toolCalls: [String: ToolCall] = [:]
    /// Raw announced fields per id, merged as updates arrive — the source the backend's approval
    /// handler falls back to, and what `complete()` recomputes the emitted call from.
    private var announced: [String: AnnouncedToolCall] = [:]
    /// Edit-class file changes awaiting a successful completion (synthesized only on success —
    /// a denied/failed edit changed nothing, mirroring the other decoders).
    private var pendingFileChanges: [String: [FileChange]] = [:]

    public init() {}

    /// The (merged) raw fields of an earlier-announced tool call, for the backend's approval
    /// handler to recover `kind`/`rawInput`/`locations` from a minimal permission request.
    public func announcedToolCall(id: String) -> AnnouncedToolCall? { announced[id] }

    public func decode(method: String, params: JSONValue) -> [Decoded] {
        // The backend routes only agent→client *notifications* here. ACP's are all `session/update`;
        // anything else is unexpected chatter → pass through as `.raw` (never a trap).
        guard method == "session/update" else {
            return [Decoded(nativeType: method, kind: .raw(Raw(native: params)))]
        }
        let update = params["update"] ?? params
        switch update["sessionUpdate"]?.stringValue {
        case "agent_message_chunk":
            return [textChunk(update, thinking: false)]
        case "agent_thought_chunk":
            return [textChunk(update, thinking: true)]
        case "tool_call":
            return decodeToolCall(update)
        case "tool_call_update":
            return decodeToolCallUpdate(update)
        case "plan", "available_commands", "current_mode_update", "config_option_update",
            "user_message_chunk":
            // Structural / echo notifications: no first-class AgentEvent. (We inject user turns
            // into the transcript ourselves; plan/modes/commands aren't modeled in v1.)
            return []
        default:
            return [Decoded(nativeType: "session/update", kind: .raw(Raw(native: update)))]
        }
    }

    // MARK: - text

    private func textChunk(_ update: JSONValue, thinking: Bool) -> Decoded {
        let chunk = TextChunk(
            messageID: update["messageId"]?.stringValue ?? "msg-unknown",
            text: Self.contentText(update["content"]),
            isPartial: true)  // ACP streams message/thought content as chunks
        let variant = thinking ? "agent_thought_chunk" : "agent_message_chunk"
        return Decoded(
            nativeType: "session/update:\(variant)",
            kind: thinking ? .thinking(chunk) : .assistantText(chunk))
    }

    // MARK: - tools

    private func decodeToolCall(_ update: JSONValue) -> [Decoded] {
        let id = update["toolCallId"]?.stringValue ?? "call-unknown"
        let kindStr = update["kind"]?.stringValue ?? "other"
        let rawInput = update["rawInput"] ?? .object([:])
        let locations = Self.locationPaths(update["locations"])
        announced[id] = AnnouncedToolCall(kind: kindStr, rawInput: rawInput, locations: locations)
        let call = ToolCall(
            toolCallID: id,
            name: Self.toolName(forKind: kindStr),
            input: Self.normalizedInput(forKind: kindStr, rawInput: rawInput, locations: locations))
        toolCalls[id] = call

        var out: [Decoded] = []
        if startedToolCallIDs.insert(id).inserted {
            out.append(Decoded(nativeType: "session/update:tool_call", kind: .toolCallStarted(call)))
        }
        // Stage any file changes for emission once the call completes successfully.
        if let changes = Self.fileChanges(forKind: kindStr, id: id, rawInput: rawInput, locations: locations) {
            pendingFileChanges[id] = changes
        }
        // A tool_call that already carries a terminal status completes inline (no separate update).
        if Self.isTerminal(update["status"]?.stringValue) {
            out += complete(id: id, update: update)
        }
        return out
    }

    private func decodeToolCallUpdate(_ update: JSONValue) -> [Decoded] {
        let id = update["toolCallId"]?.stringValue ?? "call-unknown"
        // Fold any changed fields (kind/rawInput/locations) into the announced state so the
        // completed call + file changes reflect the latest values, not the initial announce.
        mergeUpdate(update, into: id)
        guard Self.isTerminal(update["status"]?.stringValue) else {
            return []  // pending / in_progress are advisory; authoritative content lands at terminal
        }
        return complete(id: id, update: update)
    }

    /// Deep-merge a `tool_call_update`'s present fields into the announced state for `id`, and
    /// recompute the mapped `ToolCall` + staged file changes from the merged result. No-op when
    /// the id was never announced (an orphan update keeps the original toolResult-only behavior).
    private func mergeUpdate(_ update: JSONValue, into id: String) {
        guard let prior = announced[id] else { return }
        let kind = update["kind"]?.stringValue ?? prior.kind
        let rawInput = Self.mergedObject(prior.rawInput, update["rawInput"])
        let newLocations = Self.locationPaths(update["locations"])
        let locations = newLocations.isEmpty ? prior.locations : newLocations
        announced[id] = AnnouncedToolCall(kind: kind, rawInput: rawInput, locations: locations)
        toolCalls[id] = ToolCall(
            toolCallID: id,
            name: Self.toolName(forKind: kind),
            input: Self.normalizedInput(forKind: kind, rawInput: rawInput, locations: locations))
        if let changes = Self.fileChanges(forKind: kind, id: id, rawInput: rawInput, locations: locations) {
            pendingFileChanges[id] = changes
        }
    }

    /// Emit the completion trio for a terminal tool call: `toolCallCompleted` (re-emitting the
    /// announced call), `toolResult` (isError from a `failed` status), and — on success — any
    /// staged `fileChange`s.
    private func complete(id: String, update: JSONValue) -> [Decoded] {
        var out: [Decoded] = []
        if let call = toolCalls[id] {
            out.append(Decoded(nativeType: "session/update:tool_call_update", kind: .toolCallCompleted(call)))
        }
        let isError = update["status"]?.stringValue == "failed"
        out.append(
            Decoded(
                nativeType: "session/update:tool_call_update",
                kind: .toolResult(
                    ToolResult(
                        toolCallID: id,
                        content: Self.toolOutput(update["content"]),
                        isError: isError))))
        if !isError, let changes = pendingFileChanges.removeValue(forKey: id) {
            for change in changes {
                out.append(Decoded(nativeType: "session/update:tool_call_update", kind: .fileChange(change)))
            }
        } else if isError {
            pendingFileChanges.removeValue(forKey: id)
        }
        return out
    }

    // MARK: - shared kind → Nucleic mapping (reused by the backend's approval handler)

    /// Map an ACP tool `kind` onto Nucleic's tool vocabulary so the shared `RiskClassifier`,
    /// command summaries, and approval UI work for Grok exactly as for Claude/Codex.
    public static func toolName(forKind kind: String) -> String {
        switch kind {
        case "execute": return "Bash"
        case "read": return "Read"
        case "edit": return "Edit"
        case "delete": return "Edit"  // a worktree deletion; classified as a write
        case "move": return "Edit"    // rename
        case "search": return "Grep"
        case "fetch": return "WebFetch"
        case "think": return "Think"
        default: return "Tool"  // ACP "other"/unknown → fail-safe: surfaces for approval
        }
    }

    /// Normalize an ACP `rawInput` into the key shape `RiskClassifier`/titles expect for the
    /// mapped tool name: `command` for execute, `file_path` for edit/delete/move/read,
    /// `url` for fetch, `query` for search. Preserves the original raw fields too.
    public static func normalizedInput(
        forKind kind: String, rawInput: JSONValue, locations: [String]
    ) -> JSONValue {
        var fields = rawInput.objectValue ?? [:]
        switch kind {
        case "execute":
            if fields["command"] == nil {
                fields["command"] = .string(
                    rawInput["command"]?.stringValue ?? rawInput["cmd"]?.stringValue
                        ?? rawInput.stringValue ?? "")
            }
        case "edit", "delete", "move", "read":
            if fields["file_path"] == nil, let path = canonicalPath(rawInput: rawInput, locations: locations) {
                fields["file_path"] = .string(path)
            }
        case "fetch":
            if fields["url"] == nil {
                fields["url"] = .string(rawInput["url"]?.stringValue ?? rawInput["uri"]?.stringValue ?? "")
            }
        case "search":
            if fields["query"] == nil {
                fields["query"] = .string(
                    rawInput["query"]?.stringValue ?? rawInput["pattern"]?.stringValue ?? "")
            }
        default:
            break
        }
        return .object(fields)
    }

    /// The repo-relative file path(s) an ACP edit-class call touches, for conflict detection —
    /// derived from the authoritative `locations` when present, else the `rawInput` path. Empty
    /// for non-editing kinds (read/search/execute/fetch/think/other).
    public static func editPaths(forKind kind: String, rawInput: JSONValue, locations: [String]) -> [String] {
        switch kind {
        case "edit", "delete":
            if !locations.isEmpty { return locations }
            return canonicalPath(rawInput: rawInput, locations: []).map { [$0] } ?? []
        case "move":
            // A rename touches BOTH endpoints — lock both. Union locations with the common
            // source/destination spellings (order-stable, deduped) so neither side is missed.
            var paths = locations
            for key in ["source", "destination", "old_path", "new_path", "from", "to", "file_path", "path"] {
                if let p = rawInput[key]?.stringValue, !paths.contains(p) { paths.append(p) }
            }
            return paths
        default:
            return []
        }
    }

    // MARK: - helpers

    /// Synthesize the (inferred) file changes for a successful edit-class call.
    private static func fileChanges(
        forKind kind: String, id: String, rawInput: JSONValue, locations: [String]
    ) -> [FileChange]? {
        let changeKind: FileChange.ChangeKind
        switch kind {
        case "edit": changeKind = .update
        case "delete": changeKind = .delete
        case "move": changeKind = .rename
        default: return nil
        }
        let paths = editPaths(forKind: kind, rawInput: rawInput, locations: locations)
        guard !paths.isEmpty else { return nil }
        return paths.map { FileChange(path: $0, kind: changeKind, toolCallID: id, diffHint: nil) }
    }

    /// Shallow-merge `new`'s object keys over `old` (used to fold `tool_call_update.rawInput` into
    /// the announced `rawInput`). If `new` isn't an object, it replaces `old`; if absent, keep `old`.
    private static func mergedObject(_ old: JSONValue, _ new: JSONValue?) -> JSONValue {
        guard let new else { return old }
        guard let oldObj = old.objectValue, let newObj = new.objectValue else { return new }
        var merged = oldObj
        for (key, value) in newObj { merged[key] = value }
        return .object(merged)
    }

    private static func canonicalPath(rawInput: JSONValue, locations: [String]) -> String? {
        rawInput["file_path"]?.stringValue
            ?? rawInput["path"]?.stringValue
            ?? rawInput["filePath"]?.stringValue
            ?? rawInput["notebook_path"]?.stringValue
            ?? locations.first
    }

    static func locationPaths(_ locations: JSONValue?) -> [String] {
        (locations?.arrayValue ?? []).compactMap { $0["path"]?.stringValue }
    }

    private static func isTerminal(_ status: String?) -> Bool {
        status == "completed" || status == "failed"
    }

    /// Text of an ACP `ContentBlock` (or an array of them) — used for message/thought chunks.
    private static func contentText(_ content: JSONValue?) -> String {
        guard let content else { return "" }
        if let arr = content.arrayValue {
            return arr.compactMap { $0["text"]?.stringValue }.joined()
        }
        return content["text"]?.stringValue ?? content.stringValue ?? ""
    }

    /// Best-effort textual output from a tool call's `content` (array of `ToolCallContent`,
    /// each `{type:"content", content:{type:"text",text}}`). Falls back to the raw value.
    private static func toolOutput(_ content: JSONValue?) -> JSONValue {
        guard let content else { return .null }
        if let arr = content.arrayValue {
            let text = arr.compactMap { $0["content"]?["text"]?.stringValue ?? $0["text"]?.stringValue }
                .joined(separator: "\n")
            return text.isEmpty ? content : .string(text)
        }
        return content
    }
}
