ai-app/app/androidApp/src/main/kotlin/com/example/aiapp/TranscriptItems.kt

package com.example.aiapp

import androidx.compose.runtime.Immutable
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext

/**
 * What the transcript renders: the event stream folded into displayable rows (see [foldEvent]). The
 * stream is the only data source -- opening a session screen replays from seq 0, and a reconnect
 * resumes from the last seq seen, so there is no separate history fetch to drift from it.
 */
@Immutable
sealed class TranscriptItem {
    /**
     * The transcript sequence number this row started at, and its identity on screen.
     *
     * The list is drawn newest-first, so every new message is an insertion at index 0 and every
     * page of history is an insertion at the far end. Without an identity that survives both, the
     * list is addressed by position: whatever somebody had scrolled to keeps its index while the
     * content underneath it slides, which reads as the view scrolling on its own.
     *
     * A seq is the right identity because it is what the transcript itself is ordered by, it never
     * changes, and it is already carried by every event. A row built from several events -- a
     * streaming message, a tool call and its result -- keeps the seq of the first, so it holds
     * still while the rest of it arrives.
     */
    abstract val seq: Long

    /**
     * This item's identity on screen, which is its [seq] for everything that has one of its own.
     *
     * Here rather than in [TranscriptRow.Single] because the two items that need something else are
     * the two that know why: a tool call is named after its run, and a peer note is *sorted* by the
     * turn it started rather than by where it arrived. Asking each item what it is called is also
     * what stops the next such item being missed -- a `when` over concrete types in the row would
     * have to gain a case, silently, and nothing says when it did not.
     */
    open val key: Any
        get() = seq

    data class UserMsg(
        override val seq: Long,
        val text: String,
        /** Refs of what was attached, drawn inside the bubble. */
        val images: List<String> = emptyList(),
    ) : TranscriptItem()

    data class AssistantMsg(
        override val seq: Long,
        val text: String,
        /**
         * Whether this reply is finished: the session has stopped working since its last delta.
         *
         * What it buys is the split. [transcriptUnits] keeps the newest reply whole because a
         * streaming reply's text changes per delta and splitting a changing text is a parse per
         * delta -- but "newest" outlives the turn, so a session that ends on a long reply was
         * drawing it as one item indefinitely, with every node of it alive. Measured on a Pixel 9
         * Pro XL: one 34,996px reply on screen put the frame's draw phase at 13.8ms, 79% of it the
         * framework's own bookkeeping, which grows with alive nodes.
         *
         * Folded from the status event that ended the turn, rather than read off the screen's
         * status, because rows only change through the held-events gate: the split changes the
         * newest row's list identity, and doing that from a status flip while somebody is reading
         * inside that reply would step the list under them. An event has to wait for the reader to
         * be at the newest end; a screen state does not.
         */
        val settled: Boolean = false,
    ) : TranscriptItem()

    data class ToolRun(
        override val seq: Long,
        val id: String,
        /**
         * The run of adjacent calls this one belongs to, named once when the call is folded in and
         * never recomputed.
         *
         * Carried rather than derived because a run can gain members at *either* end -- a new call
         * arriving beside it, or a page of history arriving in front of it -- so no function of its
         * current members is stable. It is the first call's id at the moment the run started, which
         * is a name rather than a description: [joinPages] hands it to older calls that turn out to
         * belong to the same run, instead of renaming the run they joined.
         */
        val runId: String,
        val tool: String,
        val input: String,
        val output: String,
        val done: Boolean,
        /**
         * The questions this call is waiting on, in the order they were asked.
         *
         * On the call's own row rather than beside it: an ask used to arrive as a second card
         * repeating the input verbatim, so the reader saw the same command twice and had to work
         * out that it was one event. The backend says which call a question is about, so this is a
         * fact rather than a match on the input.
         *
         * A list because AskUserQuestion asks up to four at once, and they are one decision to make
         * -- a permission is the case of exactly one, not a different shape.
         */
        val asks: List<QuestionCard> = emptyList(),
        /**
         * Images this call's result carried, drawn under it.
         *
         * Beside it they had to be paired by position, and position is the thing a page boundary
         * breaks -- a screenshot loaded on one page and its call on the next read as unrelated.
         */
        val images: List<String> = emptyList(),
    ) : TranscriptItem() {
        /** A run, not a seq: see [TranscriptRow.key] for what that identity has to survive. */
        override val key: Any
            get() = runId
    }

    data class QuestionCard(
        override val seq: Long,
        val id: String,
        val prompt: String,
        /** A few words naming what this is about, when the asker offered one. */
        val header: String?,
        val options: List<QuestionOption>,
        /** Whether several options may be chosen at once. */
        val multiSelect: Boolean,
        /** What was chosen, once something was; empty until then. */
        val answers: List<String>,
    ) : TranscriptItem()

    data class ErrorMsg(override val seq: Long, val message: String) : TranscriptItem()

    /** An image by server-side ref, fetched from the session's files route. */
    data class ImageItem(override val seq: Long, val ref: String) : TranscriptItem()

    /**
     * A message another agent sent this session.
     *
     * Its own row rather than a [UserMsg]: see [PeerMessageRow] for why the voice matters.
     */
    data class PeerNote(
        override val seq: Long,
        val from: String,
        val text: String,
        /**
         * The seq of the event this note came in on, which is what makes it itself.
         *
         * [seq] is where the note *sorts*, and [placePeerNote] sets it to the seq the turn began at
         * so the note is drawn above the reply it caused. Two messages that arrive during one turn
         * therefore share a seq -- and sharing an identity as well killed the app, because the
         * transcript list refuses two items with one key. Two agents writing to a session mid-turn
         * is an ordinary afternoon, not a corner.
         */
        val arrived: Long = seq,
    ) : TranscriptItem() {
        override val key: Any
            get() = arrived
    }

    /**
     * A command the session ran on itself -- `/compact`, `/rename`.
     *
     * Kept in the transcript rather than only shown while it waits, because it explains what
     * follows: a conversation that suddenly has half the context, or a session with a new name.
     */
    data class CommandRow(override val seq: Long, val text: String) : TranscriptItem()

    /** Placeholder row for events this build can't render (newer kinds). */
    data class Note(override val seq: Long, val text: String) : TranscriptItem()

    /**
     * A clear that happened: everything above it left the session's context and stayed on screen.
     *
     * Carries only its position, because that is all it means.
     */
    data class ClearedNote(override val seq: Long) : TranscriptItem()

    /**
     * A compaction that happened, and what it recovered.
     *
     * In the transcript rather than only in the status line, because the status is gone the moment
     * it finishes and this is the part worth keeping: it is the explanation for a gap in the
     * conversation, and for a minute or two in which the session was busy with nothing to show.
     *
     * The wire also says what triggered it, and this deliberately does not carry that: the row says
     * the two sizes and nothing else (see [compactionSummary]), so keeping the trigger here would
     * be a field nothing can read.
     */
    data class CompactedNote(
        override val seq: Long,
        val preTokens: Long?,
        val postTokens: Long?,
    ) : TranscriptItem()
}

/**
 * The run a call joins: the one it lands next to, or a new one named after itself.
 *
 * Only ever consulted when the call is first folded in. That is what makes the name stable -- a run
 * keeps whatever it was called when it started, however many calls arrive at either end of it
 * afterwards.
 *
 * A question to the reader is in a run of its own, which is what puts it on the transcript as a row
 * rather than inside a collapsed "Called 6 tools" card. Two things follow from being alone: it is
 * always visible, since a run of one is drawn as itself rather than as a group; and the calls
 * around it fall into a group before it and a group after it, so where the reader was asked
 * something is legible in the shape of the transcript without opening anything. It ends the run
 * before it as well as starting a fresh one after -- the moment somebody was asked is a boundary in
 * the work, not a gap in the middle of one run.
 */
private fun runIdFor(items: List<TranscriptItem>, id: String, tool: String): String {
    val previous = items.lastOrNull() as? TranscriptItem.ToolRun ?: return id
    if (tool == ASK_USER_QUESTION || previous.tool == ASK_USER_QUESTION) return id
    return previous.runId
}

/**
 * Puts a page of older items in front of the ones already loaded, healing whatever the page
 * boundary cut in two.
 *
 * Two things straddle a boundary: a tool call separated from its result, and a message separated
 * from the rest of itself. Both were one thing before the transcript was cut into pages, and both
 * have to be one thing again -- a reply drawn as two messages is the same defect as a call drawn
 * twice, arriving from the same cause.
 *
 * A boundary lands wherever it lands, and roughly half the time that is between a call and its
 * result. The newer page then holds a `ToolEnd` whose start it never saw, which [foldEvent] draws
 * as a row of its own -- correctly, because a call that renders as nothing is indistinguishable
 * from one that never happened. When the older page arrives it brings the real `ToolStart`, and
 * concatenating the two lists left *both*: the same call twice, once as a proper card and once as a
 * nameless placeholder. Visible as a run of four calls reporting "Called 5 tools", and worse than
 * the miscount -- the extra row is at the join, so it also moves everything the reader was looking
 * at.
 *
 * Merged by the call's own id rather than by position, because position is exactly what a page
 * boundary destroys. The older row wins on what a start knows (the tool's name, its input) and the
 * newer on what an end knows (the output, and whether it finished), which is the only way round
 * that loses nothing.
 *
 * The third thing is the *run*, and it is the one that used to be missed. Every page ends up here,
 * but [adoptRun] only ran on the path where a split call had been found -- so the boundary that
 * falls cleanly between two finished calls, which is most of them, went straight to concatenation
 * and left the older page's calls under the run name they were folded with. On screen: one run of
 * tool calls drawn as two groups, with the seam wherever the reader happened to have paged. The two
 * early returns were an optimisation on a list the size of one page, and they were skipping work
 * rather than saving it.
 */
fun joinPages(earlier: List<TranscriptItem>, later: List<TranscriptItem>): List<TranscriptItem> {
    val (older, newer) = healSplitMessage(earlier, later)
    val startedEarlier =
        older.filterIsInstance<TranscriptItem.ToolRun>().mapTo(mutableSetOf()) { it.id }
    val endedLater =
        newer
            .filterIsInstance<TranscriptItem.ToolRun>()
            .associateBy { it.id }
            .filterKeys { it in startedEarlier }
    val healed = older.map { row ->
        val half = (row as? TranscriptItem.ToolRun)?.let { endedLater[it.id] }
        if (row is TranscriptItem.ToolRun && half != null) {
            row.copy(
                output = half.output,
                done = half.done,
                // Kept from both halves: a question or an image can be attached to either,
                // depending on which side of the boundary its event fell.
                asks = row.asks + half.asks,
                images = row.images + half.images,
            )
        } else {
            row
        }
    }
    val kept = newer.filterNot { it is TranscriptItem.ToolRun && it.id in endedLater }
    return adoptRun(healed, kept) + kept
}

/**
 * Rejoins a message the page boundary cut, and hands back the two pages to concatenate.
 *
 * [foldEvent] never leaves two assistant messages next to each other inside one page -- deltas
 * accumulate into the message before them -- so two meeting at a join are always the two halves of
 * one reply, and leaving them apart drew a single answer as two, with a paragraph break through the
 * middle of a sentence.
 *
 * The newer half keeps its identity, for the reason [adoptRun] gives: it is the row already on
 * screen, and renaming that is how the list loses its anchor. It grows by what the older half
 * brings, which is safe here and nowhere else -- the join is at the oldest end of what is loaded,
 * so the growth extends off the top of the screen, away from the row the list anchors to.
 */
private fun healSplitMessage(
    earlier: List<TranscriptItem>,
    later: List<TranscriptItem>,
): Pair<List<TranscriptItem>, List<TranscriptItem>> {
    val head = earlier.lastOrNull()
    val tail = later.firstOrNull()
    if (head !is TranscriptItem.AssistantMsg || tail !is TranscriptItem.AssistantMsg) {
        return earlier to later
    }
    return earlier.dropLast(1) to (listOf(tail.copy(text = head.text + tail.text)) + later.drop(1))
}

/**
 * Hands the older calls at the join the name of the run they are joining.
 *
 * The two pages were folded separately, so a run split by the boundary came back as two runs with
 * two names. Naming the joined run after the *older* half would be the obvious way round and is the
 * wrong one: the newer half is the part already on screen, and renaming it is renaming the row the
 * reader is looking at, which is how a list loses its anchor and steps under them. So the arriving
 * calls take the name of the ones already there, and nothing visible changes identity.
 */
private fun adoptRun(
    earlier: List<TranscriptItem>,
    later: List<TranscriptItem>,
): List<TranscriptItem> {
    val first = later.firstOrNull() as? TranscriptItem.ToolRun ?: return earlier
    // A question is in a run of its own on both sides of the join, the same as it would be had
    // the two pages been folded as one -- see `runIdFor`. Without this the heal would merge a
    // group straight through the row the reader was asked something on.
    if (first.tool == ASK_USER_QUESTION) return earlier
    val joining = first.runId
    val tail = earlier.takeLastWhile {
        it is TranscriptItem.ToolRun && it.tool != ASK_USER_QUESTION
    }
    if (tail.isEmpty()) return earlier
    return earlier.dropLast(tail.size) +
        tail.map { (it as TranscriptItem.ToolRun).copy(runId = joining) }
}

/**
 * A peer message goes above the turn it started, not where it happened to arrive.
 *
 * The live Claude Code path cannot record it in place: the CLI says nothing about a peer message
 * until the turn's `result`, so the event lands below the whole reply it caused -- the answer
 * printed above the question. The server stamps it with where that turn began
 * ([SessionEvent.PeerMessage.turnStart]) and the note takes that seq, so it sorts into the list
 * where it belongs rather than being drawn out of order at the end.
 *
 * Taking the turn's opening seq as its own is also what keeps the list sorted, which anchors and
 * paging both depend on. It is only a *position*, though, and the note keeps its own arrival seq as
 * its identity ([TranscriptItem.PeerNote.arrived]). The argument for sharing was that the turn's
 * seq belongs to a status change and a status draws no row -- true, and it answered the wrong
 * question: what two notes stamped with the same turn collide with is each other.
 *
 * Without a stamp -- a message replayed out of a session file, which is already in the right place
 * -- it stays where it arrived.
 */
private fun placePeerNote(
    items: List<TranscriptItem>,
    seq: Long,
    event: SessionEvent.PeerMessage,
): List<TranscriptItem> {
    val at = event.turnStart ?: return items + TranscriptItem.PeerNote(seq, event.from, event.text)
    val note = TranscriptItem.PeerNote(at, event.from, event.text, arrived = seq)
    val index = items.indexOfFirst { it.seq > at }
    if (index < 0) return items + note
    val behind = (items.getOrNull(index - 1) as? TranscriptItem.ToolRun)?.runId
    return items.subList(0, index) + note + splitRun(items.subList(index, items.size), behind)
}

/**
 * The calls the note now sits in front of, renamed if they were sharing a run with the calls behind
 * it.
 *
 * A run is named from what a call landed next to (see [runIdFor]), and nothing there knows about
 * turns -- so a turn opening with a tool call, straight after one that ended with one, folds them
 * into a single run. Left alone, [groupToolRuns] would flush at the note and hand both halves the
 * same name: two rows with one key, which a keyed list cannot draw at all.
 *
 * The later half is the one renamed, which is the opposite of a page join ([adoptRun]) and right
 * for the opposite reason. There the two halves were always one run and the newer was already on
 * screen; here they were never one turn's work, and both halves change appearance at the same
 * moment the note appears between them.
 */
private fun splitRun(tail: List<TranscriptItem>, behind: String?): List<TranscriptItem> {
    val first = tail.firstOrNull() as? TranscriptItem.ToolRun ?: return tail
    if (behind == null || first.runId != behind) return tail
    val run = tail.takeWhile { it is TranscriptItem.ToolRun && it.runId == behind }
    return run.map { (it as TranscriptItem.ToolRun).copy(runId = first.id) } + tail.drop(run.size)
}

fun foldEvent(items: List<TranscriptItem>, entry: SeqEvent): List<TranscriptItem> =
    when (val event = entry.event) {
        is SessionEvent.UserMessage ->
            items + TranscriptItem.UserMsg(entry.seq, event.text, event.images)
        is SessionEvent.AssistantText -> {
            // Deltas accumulate into the message they're streaming, which keeps the seq of the
            // first of them: a row whose identity changed with every delta would be a new row on
            // every frame, and the list would jump for the whole of a streamed answer.
            val last = items.lastOrNull()
            if (last is TranscriptItem.AssistantMsg) {
                // A message growing again is not finished, whatever a status said in between.
                items.dropLast(1) + last.copy(text = last.text + event.delta, settled = false)
            } else {
                items + TranscriptItem.AssistantMsg(entry.seq, event.delta)
            }
        }
        is SessionEvent.ToolStart ->
            items +
                TranscriptItem.ToolRun(
                    entry.seq,
                    event.id,
                    runIdFor(items, event.id, event.tool),
                    event.tool,
                    event.input,
                    "",
                    done = false,
                )
        is SessionEvent.ToolUpdate -> updateTool(items, event.id) { it.copy(output = event.output) }
        is SessionEvent.ToolEnd ->
            // Created when its start is not here, rather than dropped. A
            // fold that only ever *updates* loses the whole call when the
            // start fell outside the loaded window, and a tool call that
            // renders as nothing is indistinguishable from one that never
            // happened. The name is unknown from an end alone; loading the
            // page before this one replaces the row with the real thing.
            if (items.any { it is TranscriptItem.ToolRun && it.id == event.id }) {
                updateTool(items, event.id) { it.copy(output = event.output, done = true) }
            } else {
                items +
                    TranscriptItem.ToolRun(
                        entry.seq,
                        event.id,
                        // The name is not known from an end alone, so a call that was an ask
                        // cannot be recognised as one here; loading the page before this
                        // replaces the row with the real thing, which is when it splits out.
                        runIdFor(items, event.id, "tool"),
                        "tool",
                        "",
                        event.output,
                        done = true,
                    )
            }
        is SessionEvent.Question -> {
            val card =
                TranscriptItem.QuestionCard(
                    entry.seq,
                    event.id,
                    event.prompt,
                    event.header,
                    event.options,
                    event.multiSelect,
                    emptyList(),
                )
            // A question with no tool behind it -- AskUserQuestion, or an ask
            // whose call fell outside the loaded window -- is a card of its
            // own, which is what every question was before this.
            if (
                event.about != null &&
                    items.any { it is TranscriptItem.ToolRun && it.id == event.about }
            ) {
                updateTool(items, event.about) { it.copy(asks = it.asks + card) }
            } else {
                items + card
            }
        }
        is SessionEvent.Answered ->
            // Resolved wherever it is drawn: a card of its own, or a tool
            // row's ask. Missing the second left an Allow/Deny pair live on
            // a question already answered from another device.
            items.map {
                when {
                    it is TranscriptItem.QuestionCard && it.id == event.id ->
                        it.copy(answers = event.answers)
                    it is TranscriptItem.ToolRun && it.asks.any { ask -> ask.id == event.id } ->
                        it.copy(
                            asks =
                                it.asks.map { ask ->
                                    if (ask.id == event.id) ask.copy(answers = event.answers)
                                    else ask
                                }
                        )
                    else -> it
                }
            }
        is SessionEvent.PeerMessage -> placePeerNote(items, entry.seq, event)
        is SessionEvent.CommandSent -> items + TranscriptItem.CommandRow(entry.seq, event.text)
        // Screen-level state, not transcript rows -- see SessionScreen.
        is SessionEvent.CommandQueued -> items
        // No row of its own: a message that is still waiting is drawn as a pending bubble below
        // the transcript, and becomes an ordinary one where the session read it.
        is SessionEvent.MessageQueued -> items
        // The bubble goes away and nothing takes its place: the message was never read, so there
        // is nothing it belongs above.
        is SessionEvent.MessageDropped -> items
        is SessionEvent.Settings -> items
        is SessionEvent.Status -> settleReply(items, event.state)
        is SessionEvent.Error -> items + TranscriptItem.ErrorMsg(entry.seq, event.message)
        is SessionEvent.Image ->
            // Under the call that produced it when there is one, and a row of
            // its own when there is not -- a person's own attachment belongs
            // to no call, and neither does one whose call fell outside the
            // loaded window.
            if (
                event.about != null &&
                    items.any { it is TranscriptItem.ToolRun && it.id == event.about }
            ) {
                updateTool(items, event.about) { it.copy(images = it.images + event.ref) }
            } else {
                items + TranscriptItem.ImageItem(entry.seq, event.ref)
            }
        is SessionEvent.Cleared -> items + TranscriptItem.ClearedNote(entry.seq)
        is SessionEvent.Compacted ->
            items + TranscriptItem.CompactedNote(entry.seq, event.preTokens, event.postTokens)
        is SessionEvent.Unknown -> items + TranscriptItem.Note(entry.seq, "[${event.type}]")
        // Screen-level state, not transcript rows -- see SessionScreen.
        is SessionEvent.UsageDelta -> items
    }

/**
 * A status saying the session stopped working is the moment its newest reply is finished.
 *
 * See [TranscriptItem.AssistantMsg.settled] for what the mark buys and why it is made here in the
 * fold. Status changes are transcript events with seqs of their own, so a replayed session settles
 * its replies the same way a live one does.
 */
private fun settleReply(items: List<TranscriptItem>, state: String): List<TranscriptItem> {
    if (sessionWorking(state)) return items
    val last = items.lastOrNull() as? TranscriptItem.AssistantMsg ?: return items
    if (last.settled) return items
    return items.dropLast(1) + last.copy(settled = true)
}

private fun updateTool(
    items: List<TranscriptItem>,
    id: String,
    change: (TranscriptItem.ToolRun) -> TranscriptItem.ToolRun,
): List<TranscriptItem> = items.map {
    if (it is TranscriptItem.ToolRun && it.id == id) change(it) else it
}

/**
 * Where markdown is parsed ahead of being drawn: two threads, never all of them.
 *
 * The default dispatcher sizes itself to the machine, which is right for work somebody is waiting
 * on and wrong for work nobody is. A page of history is hundreds of parses arriving at once, and
 * taking every core for them leaves the thread that draws the frame queueing behind one -- measured
 * on a Pixel 9 Pro XL as 21ms of `waited` at the 90th percentile, which is the frame failing to
 * *start* rather than taking too long once it had.
 */
@OptIn(kotlinx.coroutines.ExperimentalCoroutinesApi::class)
private val parsingThreads = Dispatchers.Default.limitedParallelism(2)

/**
 * Parses the markdown among [rows], off whatever thread is drawing.
 *
 * Called where a page of transcript is folded rather than where a row is composed, which is the
 * whole point: the work happens seconds before the reader reaches the rows it was done for. See
 * [ParsedReplies].
 *
 * What is warmed mirrors what the rows draw, unit by unit -- prose split into its blocks, a memory
 * note whole, a peer message split the same way prose is -- because a string warmed under a key no
 * row ever looks up is a miss that nothing reports; see [transcriptUnits], which is the flatten
 * this has to agree with. It reads the same [ParsedReplies.partsOf] and [ParsedReplies.blocksOf]
 * caches the flatten does, so a message is scanned once however many pages hand it back through
 * here, while the whole loaded transcript crosses this on every page.
 *
 * Every kind of row that draws markdown belongs in the `when` below. That is the rule the peer
 * message was missing: this used to filter for assistant replies alone, so the one row type nobody
 * had thought about paid its whole parse in the frame it appeared in, with no counter saying which
 * row it was.
 */
suspend fun warm(replies: ParsedReplies, rows: List<TranscriptItem>) {
    withContext(parsingThreads) {
        val texts = rows.flatMap { row ->
            when (row) {
                is TranscriptItem.AssistantMsg ->
                    replies.partsOf(row.text).flatMap { part ->
                        when (part) {
                            is MessagePart.Prose -> replies.blocksOf(part.text)
                            // Drawn as one MarkdownText, so its whole text is the key
                            // looked up.
                            is MessagePart.Remembered -> listOf(part.text)
                        }
                    }
                // A message from another agent is markdown too, and it is the longest thing
                // in a transcript often enough that leaving it out was the whole of why one
                // cost a fifth of a second to open: it was the only markdown in the app
                // parsed on the thread that draws. Its blocks, not its text, because
                // [PeerMessageRow] draws it a block at a time.
                is TranscriptItem.PeerNote -> replies.blocksOf(row.text)
                else -> emptyList()
            }
        }
        if (texts.isNotEmpty()) replies.warm(texts)
        // After the parses exist, not before: [ParsedReplies.splitReady] is the flatten's
        // licence to draw these as blocks on the composing thread.
        rows.forEach { if (it is TranscriptItem.AssistantMsg) replies.markSplitReady(it.text) }
    }
}