import Foundation
import NucleicProtocol

/// One render-ready row, folded from the raw `[AgentEvent]` stream. The phone subscribes at
/// `.full`, so it receives streaming text deltas and every step of a tool call's lifecycle;
/// this projection coalesces those into the same shapes the Mac transcript shows — one bubble
/// per message, one card per tool call, a single contiguous block for a run of consecutive
/// calls, and a subagent spawn carrying its own nested activity — so the view layer stays dumb.
/// (Mirrors the desktop's `TranscriptProjection` at phone fidelity, including its event
/// partitioning, tool-run coalescing, and subagent nesting.)
struct TranscriptItem: Identifiable, Equatable {
    let id: String
    /// The seq of the first event that created this item — stable scroll anchor + ordering.
    let seq: UInt64
    var kind: Kind

    enum Role: Equatable { case user, assistant }

    enum Kind: Equatable {
        case message(role: Role, text: String)
        case thinking(text: String)
        /// A single tool call. A subagent spawn (`Task`/`Agent`) carries its inner activity in
        /// `group.children` and renders as the gold subagent card; every other tool is a plain
        /// collapsible card.
        case tool(ToolGroup)
        /// A run of ≥2 consecutive tool calls coalesced into one contiguous block — the phone
        /// echo of the Mac's `.toolGroup`. A run made entirely of subagent spawns is a
        /// multi-agent fan-out (the gold orchestration card); any other run is a tool block.
        case toolBlock([ToolGroup])
        case sessionStarted(model: String, cwd: String)
        case usage(Usage)
        case rateLimit(RateLimit)
        case turnBoundary
        case approval(toolName: String)
        case runFinished(outcome: RunFinished.Outcome)
        case error(message: String)
        case note(text: String, icon: String?, lockEvent: Bool)
        case raw(type: String, body: String)
    }

    /// A row that draws nothing visible — an empty/redacted `.thinking` block (the agent streams
    /// a finalized empty thinking item whenever the reasoning itself is redacted). Held aside
    /// during tool-run coalescing so it doesn't split an otherwise-contiguous run of calls into
    /// two cards with a gap. Mirrors the desktop projection's `rendersNothing`.
    var rendersNothing: Bool {
        if case .thinking(let text) = kind {
            return text.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty
        }
        return false
    }
}

/// One tool call's coalesced lifecycle: start → input deltas → complete → result → file changes.
struct ToolGroup: Equatable {
    var toolCallID: String
    var name: String
    var input: JSONValue
    var result: JSONValue?
    var isError: Bool = false
    var finished: Bool = false
    /// Paths the call touched (from `fileChange` events tagged with this tool call).
    var fileChanges: [FilePatch] = []
    /// For a subagent spawn (`Task`/`Agent`): the inner activity it produced — its own tool
    /// calls, thinking, and prose — projected the same way and nested here, so the card can show
    /// its live Read/Bash/Edit stream. Empty for every ordinary (non-subagent) tool.
    var children: [TranscriptItem] = []

    struct FilePatch: Equatable { let path: String; let change: FileChange.ChangeKind }

    /// Subagent orchestration (`Task`/`Agent`) gets the gold card treatment, like the Mac.
    var isOrchestration: Bool { name == "Task" || name == "Agent" }
    var isAskUserQuestion: Bool { name == "AskUserQuestion" }
}

enum TranscriptProjection {
    /// How deep subagent nesting is followed before deeper descendants are left unattached — a
    /// safety bound on the recursion (a subagent spawning a subagent spawning a subagent …), far
    /// past any real fan-out depth. Mirrors the desktop projection.
    private static let maxSubagentDepth = 4

    /// Fold the raw event stream into render rows. `showRaw` surfaces unrecognized passthrough
    /// events (debug); `showLockEvents` keeps file-lock lifecycle notes (off = quieter feed).
    ///
    /// Work that happens *inside* a spawned subagent arrives in this same stream tagged with the
    /// parent `Task`'s id (`parentToolCallID`). We partition events by which subagent (if any)
    /// owns them, project the main agent's own events at the top level, and recursively project
    /// each subagent's events into the `children` of its spawn — so a subagent's inner work nests
    /// under its card instead of leaking (and interleaving) into the main transcript.
    static func build(_ events: [AgentEvent], showRaw: Bool, showLockEvents: Bool) -> [TranscriptItem] {
        let (topLevel, byParent) = partition(events)
        return project(topLevel, byParent: byParent, depth: 0, showRaw: showRaw, showLockEvents: showLockEvents)
    }

    // MARK: - Subagent partitioning

    /// Split a scope's events into the main agent's own (`topLevel`) and each subagent's, keyed by
    /// the spawning `Task`'s id.
    private static func partition(_ events: [AgentEvent])
        -> (topLevel: [AgentEvent], byParent: [String: [AgentEvent]])
    {
        let parentOf = toolParentMap(events)
        var topLevel: [AgentEvent] = []
        var byParent: [String: [AgentEvent]] = [:]
        for event in events {
            if let owner = subagentOwner(of: event, parentOf: parentOf) {
                byParent[owner, default: []].append(event)
            } else {
                topLevel.append(event)
            }
        }
        return (topLevel, byParent)
    }

    /// Maps each tool-call id to its parent subagent's id, built from the call start/complete
    /// events. A tool *result* or *file change* names only a tool id, so it inherits its subagent
    /// scope from the call it belongs to via this map.
    private static func toolParentMap(_ events: [AgentEvent]) -> [String: String] {
        var map: [String: String] = [:]
        for event in events {
            switch event.kind {
            case .toolCallStarted(let call), .toolCallCompleted(let call):
                if let parent = call.parentToolCallID { map[call.toolCallID] = parent }
            default:
                break
            }
        }
        return map
    }

    /// The id of the subagent whose scope this event belongs to, or nil for the main agent's own
    /// turn. Tool calls and the subagent's prose/thinking carry the parent link directly; a tool
    /// result or file change inherits it from the call it references.
    private static func subagentOwner(of event: AgentEvent, parentOf: [String: String]) -> String? {
        switch event.kind {
        case .toolCallStarted(let call), .toolCallCompleted(let call):
            return call.parentToolCallID
        case .toolResult(let result):
            return parentOf[result.toolCallID]
        case .fileChange(let change):
            return change.toolCallID.flatMap { parentOf[$0] }
        case .assistantText(let chunk), .thinking(let chunk):
            return chunk.parentToolCallID
        default:
            return nil
        }
    }

    /// Projects one scope's events into items, coalesces runs of tool calls, then nests each
    /// subagent spawn's own activity (drawn from `byParent`) as its `children`, recursing for
    /// subagents-within-subagents up to `maxSubagentDepth`.
    private static func project(
        _ events: [AgentEvent], byParent: [String: [AgentEvent]], depth: Int,
        showRaw: Bool, showLockEvents: Bool
    ) -> [TranscriptItem] {
        let items = coalesceToolRuns(flatItems(events, showRaw: showRaw, showLockEvents: showLockEvents))
        guard depth < maxSubagentDepth else { return items }
        return items.map {
            attachSubagentChildren($0, byParent: byParent, depth: depth,
                                   showRaw: showRaw, showLockEvents: showLockEvents)
        }
    }

    /// For a subagent spawn (`Task`/`Agent`) — lone or inside a fan-out block — projects the
    /// events it owns into its `children`. Leaves every ordinary tool untouched.
    private static func attachSubagentChildren(
        _ item: TranscriptItem, byParent: [String: [AgentEvent]], depth: Int,
        showRaw: Bool, showLockEvents: Bool
    ) -> TranscriptItem {
        func childrenFor(_ group: ToolGroup) -> ToolGroup {
            guard group.isOrchestration else { return group }
            var g = group
            g.children = project(byParent[group.toolCallID] ?? [], byParent: byParent, depth: depth + 1,
                                 showRaw: showRaw, showLockEvents: showLockEvents)
            return g
        }
        switch item.kind {
        case .tool(let group):
            guard group.isOrchestration else { return item }
            return TranscriptItem(id: item.id, seq: item.seq, kind: .tool(childrenFor(group)))
        case .toolBlock(let groups):
            return TranscriptItem(id: item.id, seq: item.seq, kind: .toolBlock(groups.map(childrenFor)))
        default:
            return item
        }
    }

    // MARK: - Tool-run coalescing

    /// Collapses maximal runs of adjacent `.tool` items (length ≥2) into one `.toolBlock`.
    ///
    /// Items that draw nothing (empty `.thinking` blocks — see `rendersNothing`) never break a
    /// run: an invisible row landing between two tool calls must not split them into two cards
    /// with a gap. Such items are held aside and re-emitted right after the block, so they still
    /// render in place while the calls stay one contiguous block. Mirrors the desktop projection.
    static func coalesceToolRuns(_ flat: [TranscriptItem]) -> [TranscriptItem] {
        var out: [TranscriptItem] = []
        var run: [TranscriptItem] = []          // consecutive `.tool` items
        var held: [TranscriptItem] = []         // non-rendering rows seen mid-run, held so they don't split it

        func flush() {
            if run.count == 1 {
                out.append(run[0])
            } else if run.count >= 2 {
                let groups = run.compactMap { item -> ToolGroup? in
                    if case .tool(let g) = item.kind { return g } else { return nil }
                }
                out.append(TranscriptItem(id: "toolblock-\(groups[0].toolCallID)",
                                          seq: run[0].seq, kind: .toolBlock(groups)))
            }
            run.removeAll(keepingCapacity: true)
            out.append(contentsOf: held)
            held.removeAll(keepingCapacity: true)
        }

        for item in flat {
            if case .tool = item.kind {
                run.append(item)
            } else if item.rendersNothing, !run.isEmpty {
                held.append(item)
            } else {
                flush()
                out.append(item)
            }
        }
        flush()
        return out
    }

    // MARK: - Flat projection (one scope)

    /// Fold one scope's raw events into interleaved, display-ready rows: streaming assistant /
    /// thinking deltas accumulate into one growing row, and a tool call shows the moment it starts
    /// (`toolCallStarted`), with its full input/result/file-changes filled in as they arrive.
    private static func flatItems(
        _ events: [AgentEvent], showRaw: Bool, showLockEvents: Bool
    ) -> [TranscriptItem] {
        var items: [TranscriptItem] = []
        var messageIndex: [String: Int] = [:]   // messageID → items index (text coalescing)
        var thinkingIndex: [String: Int] = [:]
        var toolIndex: [String: Int] = [:]       // toolCallID → items index

        for event in events {
            switch event.kind {
            case .userText(let chunk):
                coalesceText(.user, chunk, seq: event.seq, into: &items, index: &messageIndex)
            case .assistantText(let chunk):
                coalesceText(.assistant, chunk, seq: event.seq, into: &items, index: &messageIndex)
            case .thinking(let chunk):
                coalesceThinking(chunk, seq: event.seq, into: &items, index: &thinkingIndex)

            case .toolCallStarted(let call), .toolCallCompleted(let call):
                let finished: Bool = { if case .toolCallCompleted = event.kind { return true } else { return false } }()
                upsertTool(call, finished: finished, seq: event.seq, into: &items, index: &toolIndex)
            case .toolCallInputDelta:
                break  // the final `toolCallCompleted` carries the assembled input
            case .toolResult(let result):
                attachResult(result, into: &items, index: toolIndex)
            case .fileChange(let change):
                attachFileChange(change, seq: event.seq, into: &items, index: toolIndex)

            case .sessionStarted(let started):
                items.append(.init(id: "start-\(event.seq)", seq: event.seq,
                                   kind: .sessionStarted(model: started.model, cwd: started.cwd)))
            case .usage(let usage):
                items.append(.init(id: "usage-\(event.seq)", seq: event.seq, kind: .usage(usage)))
            case .rateLimit(let limit):
                items.append(.init(id: "rate-\(event.seq)", seq: event.seq, kind: .rateLimit(limit)))
            case .turnCompleted:
                items.append(.init(id: "turn-\(event.seq)", seq: event.seq, kind: .turnBoundary))
            case .approvalRequested(let req):
                items.append(.init(id: "appr-\(event.seq)", seq: event.seq, kind: .approval(toolName: req.toolName)))
            case .approvalResolved:
                break  // dismissal is reflected in the live approval card, not the transcript
            case .runFinished(let finished):
                items.append(.init(id: "fin-\(event.seq)", seq: event.seq,
                                   kind: .runFinished(outcome: finished.outcome)))
            case .error(let err):
                items.append(.init(id: "err-\(event.seq)", seq: event.seq, kind: .error(message: err.message)))
            case .note(let note):
                if note.lockEvent && !showLockEvents { break }
                items.append(.init(id: "note-\(event.seq)", seq: event.seq,
                                   kind: .note(text: note.text, icon: note.icon, lockEvent: note.lockEvent)))
            case .raw(let raw):
                guard showRaw else { break }
                items.append(.init(id: "raw-\(event.seq)", seq: event.seq,
                                   kind: .raw(type: event.nativeType ?? "raw", body: raw.native.compactSummary)))
            }
        }
        return items
    }

    // MARK: - Coalescing helpers

    private static func coalesceText(
        _ role: TranscriptItem.Role, _ chunk: TextChunk, seq: UInt64,
        into items: inout [TranscriptItem], index: inout [String: Int]
    ) {
        if let i = index[chunk.messageID], case .message(let r, let existing) = items[i].kind {
            // A non-partial chunk is the authoritative full text; partials accumulate.
            let text = chunk.isPartial ? existing + chunk.text : chunk.text
            items[i].kind = .message(role: r, text: text)
        } else {
            index[chunk.messageID] = items.count
            items.append(.init(id: "msg-\(chunk.messageID)", seq: seq,
                               kind: .message(role: role, text: chunk.text)))
        }
    }

    private static func coalesceThinking(
        _ chunk: TextChunk, seq: UInt64, into items: inout [TranscriptItem], index: inout [String: Int]
    ) {
        if let i = index[chunk.messageID], case .thinking(let existing) = items[i].kind {
            items[i].kind = .thinking(text: chunk.isPartial ? existing + chunk.text : chunk.text)
        } else {
            index[chunk.messageID] = items.count
            items.append(.init(id: "think-\(chunk.messageID)", seq: seq, kind: .thinking(text: chunk.text)))
        }
    }

    private static func upsertTool(
        _ call: ToolCall, finished: Bool, seq: UInt64,
        into items: inout [TranscriptItem], index: inout [String: Int]
    ) {
        if let i = index[call.toolCallID], var group = currentGroup(items[i]) {
            group.name = call.name
            if !call.input.isEmptyValue { group.input = call.input }
            group.finished = group.finished || finished
            items[i].kind = .tool(group)
        } else {
            index[call.toolCallID] = items.count
            let group = ToolGroup(toolCallID: call.toolCallID, name: call.name, input: call.input, finished: finished)
            items.append(.init(id: "tool-\(call.toolCallID)", seq: seq, kind: .tool(group)))
        }
    }

    private static func attachResult(
        _ result: ToolResult, into items: inout [TranscriptItem], index: [String: Int]
    ) {
        guard let i = index[result.toolCallID], var group = currentGroup(items[i]) else { return }
        group.result = result.content
        group.isError = result.isError
        group.finished = true
        items[i].kind = .tool(group)
    }

    private static func attachFileChange(
        _ change: FileChange, seq: UInt64, into items: inout [TranscriptItem], index: [String: Int]
    ) {
        if let id = change.toolCallID, let i = index[id], var group = currentGroup(items[i]) {
            group.fileChanges.append(.init(path: change.path, change: change.kind))
            items[i].kind = .tool(group)
        }
        // Untagged file changes are folded into the diff stat, not the transcript.
    }

    /// Extract a tool group from a `.tool` item (the only kind `flatItems` produces for a call;
    /// blocks/children are formed later, after coalescing).
    private static func currentGroup(_ item: TranscriptItem) -> ToolGroup? {
        if case .tool(let g) = item.kind { return g }
        return nil
    }
}

extension JSONValue {
    /// Whether this value carries nothing worth keeping (so a later, fuller input wins).
    var isEmptyValue: Bool {
        switch self {
        case .null: return true
        case .object(let o): return o.isEmpty
        case .string(let s): return s.isEmpty
        default: return false
        }
    }
}
