import Foundation

/// Translates Claude Code `stream-json` stdout lines into normalized event kinds
/// (ADAPTERS §1.4). Lenient by contract: anything unrecognized becomes `.raw`,
/// never a trap — that is the CLI-drift early-warning signal (OBSERVABILITY A.5).
///
/// Stateful only for partial-message assembly when `--include-partial-messages`
/// is on: `stream_event` deltas reference content blocks by index, and the full
/// `assistant` message that follows must not double-emit `toolCallStarted`.
public final class ClaudeStreamDecoder {
    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
        }
    }

    private struct StreamingBlock {
        enum BlockKind { case text, thinking, toolUse(id: String, name: String) }
        let kind: BlockKind
        var accumulatedJSON: String = ""
    }

    /// message id of the in-flight streamed assistant message.
    private var currentMessageID: String?
    /// `parent_tool_use_id` of the in-flight streamed message, set when it belongs to a
    /// subagent — so partial text/thinking deltas can be tagged with their owning `Task`,
    /// matching the whole-message path. nil for the main agent's own turns.
    private var currentParentToolCallID: String?
    /// content blocks of the in-flight streamed message, by index.
    private var streamingBlocks: [Int: StreamingBlock] = [:]
    /// tool_use ids already announced via toolCallStarted (stream or full message).
    private var startedToolCallIDs: Set<String> = []
    /// Whether `system/init` has been surfaced (it repeats per streaming turn).
    private var sawInit = false
    /// Editing-tool calls awaiting their result; fileChange is synthesized only
    /// once the tool_result confirms success (a denied/failed Write changed nothing).
    private var pendingFileChanges: [String: FileChange] = [:]

    public init() {}

    public func decode(line: Data) -> [Decoded] {
        guard !line.isEmpty else { return [] }
        guard let root = try? JSONValue(parsing: line), root.objectValue != nil else {
            // Not JSON — pass the raw text through.
            let text = String(decoding: line, as: UTF8.self)
            return [Decoded(nativeType: nil, kind: .raw(Raw(native: .string(text))))]
        }

        let type = root["type"]?.stringValue
        switch type {
        case "system":
            return decodeSystem(root)
        case "assistant":
            return decodeAssistant(root)
        case "user":
            return decodeUser(root)
        case "stream_event":
            return decodeStreamEvent(root)
        case "result":
            return decodeResult(root)
        case "rate_limit_event":
            return decodeRateLimit(root)
        default:
            return [Decoded(nativeType: type, kind: .raw(Raw(native: root)))]
        }
    }

    // MARK: system

    private func decodeSystem(_ root: JSONValue) -> [Decoded] {
        let subtype = root["subtype"]?.stringValue
        let nativeType = "system/\(subtype ?? "?")"
        switch subtype {
        case "init":
            // Claude re-emits `system/init` at the start of every turn in streaming
            // mode (verified against 2.1.167). Surface it once; suppress repeats so a
            // multi-turn chat transcript isn't littered with "session started".
            guard !sawInit else { return [] }
            sawInit = true
            let started = SessionStarted(
                backendSessionID: root["session_id"]?.stringValue ?? "",
                model: root["model"]?.stringValue ?? "",
                cwd: root["cwd"]?.stringValue ?? "",
                toolNames: root["tools"]?.arrayValue?.compactMap(\.stringValue) ?? [],
                nativeTranscriptPath: nil)
            return [Decoded(nativeType: nativeType, kind: .sessionStarted(started))]
        case "api_retry":
            let message = root["message"]?.stringValue ?? "API retry"
            return [
                Decoded(
                    nativeType: nativeType,
                    kind: .error(AgentError(recoverable: true, message: message, native: root)))
            ]
        default:
            return [Decoded(nativeType: nativeType, kind: .raw(Raw(native: root)))]
        }
    }

    // MARK: assistant / user (whole messages)

    private func decodeAssistant(_ root: JSONValue) -> [Decoded] {
        let message = root["message"]
        let messageID = message?["id"]?.stringValue ?? "msg-unknown"
        let parentToolCallID = root["parent_tool_use_id"]?.stringValue
        var out: [Decoded] = []

        for block in message?["content"]?.arrayValue ?? [] {
            switch block["type"]?.stringValue {
            case "text":
                let text = block["text"]?.stringValue ?? ""
                out.append(
                    Decoded(
                        nativeType: "assistant",
                        kind: .assistantText(
                            TextChunk(messageID: messageID, text: text, isPartial: false,
                                      parentToolCallID: parentToolCallID))))
            case "thinking":
                let text = block["thinking"]?.stringValue ?? block["text"]?.stringValue ?? ""
                out.append(
                    Decoded(
                        nativeType: "assistant",
                        kind: .thinking(TextChunk(messageID: messageID, text: text, isPartial: false,
                                                  parentToolCallID: parentToolCallID))))
            case "tool_use":
                let call = ToolCall(
                    toolCallID: block["id"]?.stringValue ?? "toolu-unknown",
                    name: block["name"]?.stringValue ?? "unknown",
                    input: block["input"] ?? .object([:]),
                    parentToolCallID: parentToolCallID)
                if startedToolCallIDs.insert(call.toolCallID).inserted {
                    out.append(Decoded(nativeType: "assistant", kind: .toolCallStarted(call)))
                }
                out.append(Decoded(nativeType: "assistant", kind: .toolCallCompleted(call)))
                if let change = inferredFileChange(for: call) {
                    pendingFileChanges[call.toolCallID] = change
                }
            default:
                out.append(Decoded(nativeType: "assistant", kind: .raw(Raw(native: block))))
            }
        }
        return out
    }

    private func decodeUser(_ root: JSONValue) -> [Decoded] {
        var out: [Decoded] = []
        let content = root["message"]?["content"]
        // String-form content (echo of stdin input) and non-tool_result blocks are
        // passed through as raw: the protocol has no user-text kind — user turns are
        // injected synthetically by the SessionController (RUNTIME §2.1).
        guard let blocks = content?.arrayValue else {
            return [Decoded(nativeType: "user", kind: .raw(Raw(native: root)))]
        }
        for block in blocks {
            if block["type"]?.stringValue == "tool_result" {
                let result = ToolResult(
                    toolCallID: block["tool_use_id"]?.stringValue ?? "toolu-unknown",
                    content: block["content"] ?? .null,
                    isError: block["is_error"]?.boolValue ?? false)
                out.append(Decoded(nativeType: "user", kind: .toolResult(result)))
                if let change = pendingFileChanges.removeValue(forKey: result.toolCallID),
                    !result.isError
                {
                    out.append(Decoded(nativeType: "user", kind: .fileChange(change)))
                }
            } else {
                out.append(Decoded(nativeType: "user", kind: .raw(Raw(native: block))))
            }
        }
        return out
    }

    /// Claude has no native file_change events; infer them from editing tools
    /// (`emitsFileChangeEvents = false` marks them non-authoritative).
    private func inferredFileChange(for call: ToolCall) -> FileChange? {
        let path: String?
        let kind: FileChange.ChangeKind
        switch call.name {
        case "Write":
            path = call.input["file_path"]?.stringValue
            kind = .add
        case "Edit", "MultiEdit":
            path = call.input["file_path"]?.stringValue
            kind = .update
        case "NotebookEdit":
            path = call.input["notebook_path"]?.stringValue
            kind = .update
        default:
            return nil
        }
        guard let path else { return nil }
        return FileChange(path: path, kind: kind, toolCallID: call.toolCallID, diffHint: nil)
    }

    // MARK: stream_event (partials, --include-partial-messages)

    private func decodeStreamEvent(_ root: JSONValue) -> [Decoded] {
        let event = root["event"]
        let eventType = event?["type"]?.stringValue
        let nativeType = "stream_event/\(eventType ?? "?")"

        switch eventType {
        case "message_start":
            currentMessageID = event?["message"]?["id"]?.stringValue
            // A subagent's streamed message carries `parent_tool_use_id` on the wrapper, the
            // same place the whole-message and tool_use paths read it; remember it for this
            // message's text/thinking deltas.
            currentParentToolCallID = root["parent_tool_use_id"]?.stringValue
            streamingBlocks.removeAll()
            return []

        case "content_block_start":
            guard let index = event?["index"]?.intValue,
                let blockType = event?["content_block"]?["type"]?.stringValue
            else { return [Decoded(nativeType: nativeType, kind: .raw(Raw(native: root)))] }
            switch blockType {
            case "text":
                streamingBlocks[index] = StreamingBlock(kind: .text)
                return []
            case "thinking":
                streamingBlocks[index] = StreamingBlock(kind: .thinking)
                return []
            case "tool_use":
                let id = event?["content_block"]?["id"]?.stringValue ?? "toolu-unknown"
                let name = event?["content_block"]?["name"]?.stringValue ?? "unknown"
                streamingBlocks[index] = StreamingBlock(kind: .toolUse(id: id, name: name))
                var out: [Decoded] = []
                if startedToolCallIDs.insert(id).inserted {
                    let call = ToolCall(
                        toolCallID: id, name: name,
                        input: event?["content_block"]?["input"] ?? .object([:]),
                        parentToolCallID: root["parent_tool_use_id"]?.stringValue)
                    out.append(Decoded(nativeType: nativeType, kind: .toolCallStarted(call)))
                }
                return out
            default:
                return [Decoded(nativeType: nativeType, kind: .raw(Raw(native: root)))]
            }

        case "content_block_delta":
            guard let delta = event?["delta"], let deltaType = delta["type"]?.stringValue
            else { return [Decoded(nativeType: nativeType, kind: .raw(Raw(native: root)))] }
            let messageID = currentMessageID ?? "msg-unknown"
            switch deltaType {
            case "text_delta":
                let text = delta["text"]?.stringValue ?? ""
                return [
                    Decoded(
                        nativeType: nativeType,
                        kind: .assistantText(TextChunk(messageID: messageID, text: text, isPartial: true,
                                                       parentToolCallID: currentParentToolCallID)))
                ]
            case "thinking_delta":
                let text = delta["thinking"]?.stringValue ?? ""
                return [
                    Decoded(
                        nativeType: nativeType,
                        kind: .thinking(TextChunk(messageID: messageID, text: text, isPartial: true,
                                                  parentToolCallID: currentParentToolCallID)))
                ]
            case "input_json_delta":
                guard let index = event?["index"]?.intValue,
                    var block = streamingBlocks[index],
                    case .toolUse(let id, _) = block.kind
                else { return [] }
                let partial = delta["partial_json"]?.stringValue ?? ""
                block.accumulatedJSON += partial
                streamingBlocks[index] = block
                return [
                    Decoded(
                        nativeType: nativeType,
                        kind: .toolCallInputDelta(ToolInputDelta(toolCallID: id, partialJSON: partial)))
                ]
            case "signature_delta":
                return []
            default:
                return [Decoded(nativeType: nativeType, kind: .raw(Raw(native: root)))]
            }

        case "content_block_stop":
            // Authoritative finalized input arrives in the subsequent whole
            // `assistant` message, which emits toolCallCompleted (dedup via
            // startedToolCallIDs). Nothing to emit here.
            if let index = event?["index"]?.intValue { streamingBlocks.removeValue(forKey: index) }
            return []

        case "message_delta", "message_stop":
            return []

        default:
            return [Decoded(nativeType: nativeType, kind: .raw(Raw(native: root)))]
        }
    }

    // MARK: rate limit (M0 finding 6)

    private func decodeRateLimit(_ root: JSONValue) -> [Decoded] {
        let info = root["rate_limit_info"]
        let resetsAt = info?["resetsAt"]?.intValue.map {
            Date(timeIntervalSince1970: TimeInterval($0))
        }
        let rateLimit = RateLimit(
            status: info?["status"]?.stringValue,
            rateLimitType: info?["rateLimitType"]?.stringValue,
            resetsAt: resetsAt,
            isUsingOverage: info?["isUsingOverage"]?.boolValue,
            overageStatus: info?["overageStatus"]?.stringValue)
        return [Decoded(nativeType: "rate_limit_event", kind: .rateLimit(rateLimit))]
    }

    // MARK: result

    private func decodeResult(_ root: JSONValue) -> [Decoded] {
        var out: [Decoded] = []
        var usage: Usage?
        if let usageObject = root["usage"] {
            // Context-window occupancy is the *final* model call's input, not the turn's
            // cumulative tokens. The result's `iterations` array is one entry per call;
            // its last entry is the final context. Fall back to the top-level usage for
            // single-call turns (or backends that omit `iterations`).
            let lastIteration = usageObject["iterations"]?.arrayValue?.last
            let contextTokens = Self.singleCallContext(lastIteration)
                ?? Self.singleCallContext(usageObject)
            usage = Usage(
                inputTokens: usageObject["input_tokens"]?.intValue,
                cachedInputTokens: usageObject["cache_read_input_tokens"]?.intValue,
                cacheCreationInputTokens: usageObject["cache_creation_input_tokens"]?.intValue,
                outputTokens: usageObject["output_tokens"]?.intValue,
                reasoningTokens: nil,
                costUSD: root["total_cost_usd"]?.numberValue,
                contextInputTokens: contextTokens)
            out.append(Decoded(nativeType: "result", kind: .usage(usage!)))
        }

        // In streaming mode `result` marks the end of a TURN, not the run: the same
        // session keeps accepting input while stdin stays open (verified, 2.1.167).
        // A successful turn → turnCompleted (→ awaitingInput); the run's terminal
        // runFinished is synthesized by the adapter when the process exits. Only an
        // error / max-turns result actually ends the run.
        let subtype = root["subtype"]?.stringValue
        let isError = root["is_error"]?.boolValue ?? false
        if subtype == "success" && !isError {
            out.append(
                Decoded(
                    nativeType: "result",
                    kind: .turnCompleted(
                        TurnCompleted(stopReason: root["stop_reason"]?.stringValue, usage: usage))))
        } else {
            let outcome: RunFinished.Outcome = subtype == "error_max_turns" ? .maxTurns : .errored
            out.append(
                Decoded(
                    nativeType: "result",
                    kind: .runFinished(
                        RunFinished(
                            outcome: outcome,
                            finalText: root["result"]?.stringValue,
                            totalUsage: usage,
                            durationMs: root["duration_ms"]?.intValue))))
        }
        return out
    }

    /// Total input of a single model call (regular + cache-read + cache-creation) —
    /// the context size for that call. Nil if the object is missing or sums to zero.
    private static func singleCallContext(_ usage: JSONValue?) -> Int? {
        guard let usage else { return nil }
        let total = (usage["input_tokens"]?.intValue ?? 0)
            + (usage["cache_read_input_tokens"]?.intValue ?? 0)
            + (usage["cache_creation_input_tokens"]?.intValue ?? 0)
        return total > 0 ? total : nil
    }
}
