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

package com.example.aiapp

import androidx.compose.foundation.background
import androidx.compose.foundation.horizontalScroll
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.Row
import androidx.compose.foundation.layout.fillMaxWidth
import androidx.compose.foundation.layout.padding
import androidx.compose.foundation.rememberScrollState
import androidx.compose.foundation.shape.RoundedCornerShape
import androidx.compose.material3.MaterialTheme
import androidx.compose.runtime.Composable
import androidx.compose.runtime.CompositionLocalProvider
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.Stable
import androidx.compose.runtime.key
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.draw.drawWithContent
import androidx.compose.ui.graphics.graphicsLayer
import androidx.compose.ui.layout.layout
import androidx.compose.ui.semantics.CollectionInfo
import androidx.compose.ui.semantics.CollectionItemInfo
import androidx.compose.ui.semantics.collectionInfo
import androidx.compose.ui.semantics.collectionItemInfo
import androidx.compose.ui.semantics.heading
import androidx.compose.ui.semantics.semantics
import androidx.compose.ui.text.AnnotatedString
import androidx.compose.ui.text.TextLinkStyles
import androidx.compose.ui.text.TextStyle
import androidx.compose.ui.text.font.FontFamily
import androidx.compose.ui.text.font.FontWeight
import androidx.compose.ui.text.style.TextDecoration
import androidx.compose.ui.unit.TextUnit
import androidx.compose.ui.unit.dp
import com.mikepenz.markdown.compose.LocalImageTransformer
import com.mikepenz.markdown.compose.LocalMarkdownAnimations
import com.mikepenz.markdown.compose.LocalMarkdownColors
import com.mikepenz.markdown.compose.LocalMarkdownComponents
import com.mikepenz.markdown.compose.LocalMarkdownDimens
import com.mikepenz.markdown.compose.LocalMarkdownPadding
import com.mikepenz.markdown.compose.LocalMarkdownTypography
import com.mikepenz.markdown.compose.LocalReferenceLinkHandler
import com.mikepenz.markdown.compose.components.markdownComponents
import com.mikepenz.markdown.compose.elements.MarkdownDivider
import com.mikepenz.markdown.compose.elements.listDepth
import com.mikepenz.markdown.m3.elements.MarkdownCheckBox
import com.mikepenz.markdown.m3.markdownColor
import com.mikepenz.markdown.m3.markdownTypography
import com.mikepenz.markdown.model.NoOpImageTransformerImpl
import com.mikepenz.markdown.model.State
import com.mikepenz.markdown.model.markdownAnimations
import com.mikepenz.markdown.model.markdownDimens
import com.mikepenz.markdown.model.markdownPadding
import com.mikepenz.markdown.model.parseMarkdown
import java.util.concurrent.ConcurrentHashMap
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import org.intellij.markdown.MarkdownTokenTypes
import org.intellij.markdown.ast.ASTNode
import org.intellij.markdown.ast.findChildOfType
import org.intellij.markdown.flavours.gfm.GFMElementTypes
import org.intellij.markdown.flavours.gfm.GFMTokenTypes

/**
 * [text] drawn as its pieces, one under the other; see [Piece].
 *
 * [live] is the reply still arriving, and two things are different for it. Its parse is incremental
 * -- see [LiveParse] -- so a delta costs a parse of the block it landed in rather than of the whole
 * message. And its pieces get a layer each: when drawing is invalidated, only the piece that
 * changed is re-recorded instead of the whole reply, which is worth a great deal while every delta
 * invalidates the message and a finished one can be twenty-five screens tall. It is worth nothing
 * once the message stops changing -- measured on a Pixel 9 Pro XL, whole rows were re-recorded 65
 * times in fifty seconds of reading -- and it is not free: each layer is a layout node and a
 * display list held for the life of the row, and live node count is what the per-frame cost of the
 * transcript scales with.
 */
@Composable
fun MarkdownText(
    text: String,
    replies: ParsedReplies,
    modifier: Modifier = Modifier,
    live: Boolean = false,
) {
    val segments =
        if (live) liveSegments(text)
        else remember(text) { listOf(Segment(text, 0, replies.of(text), replies.piecesOf(text))) }
    Column(modifier.fillMaxWidth()) {
        var previous: Piece? = null
        var previousSegment: Segment? = null
        segments.forEachIndexed { at, segment ->
            val nextContinues = segments.getOrNull(at + 1)?.continues == true
            // Only the tail is still being written; a frozen segment is finished text that
            // happens to sit in a live reply, and it takes its colours now. See [MarkdownRoot].
            MarkdownRoot(segment.parse, replies, streaming = live && at == segments.lastIndex) {
                segment.pieces.forEachIndexed { index, piece ->
                    val gap =
                        when {
                            previousSegment == null -> 0.dp
                            previousSegment !== segment ->
                                if (segment.continues) 0.dp else BLOCK_SPACING
                            else -> gapBefore(previous, piece)
                        }
                    // Keyed by where the piece starts in the message rather than by its position
                    // in this column, so a delta landing in the last block leaves every other
                    // piece's composition alone -- and a block keeps its key when it freezes.
                    key(segment.start, piece) {
                        MarkdownPiece(
                            segment.parse,
                            segment.text,
                            piece,
                            Modifier.padding(top = gap)
                                .then(if (live) Modifier.graphicsLayer() else Modifier)
                                .drawWithContent {
                                    val started = System.nanoTime()
                                    drawContent()
                                    DebugStats.record(
                                        "record: one block",
                                        System.nanoTime() - started,
                                    )
                                },
                            continuesList = segment.continues && index == 0,
                            listContinues = nextContinues && index == segment.pieces.lastIndex,
                        )
                    }
                    previous = piece
                    previousSegment = segment
                }
            }
        }
    }
}

/**
 * A stretch of a message with a parse of its own: the whole of a settled message, or one block, the
 * finished items of one list, or the unfinished tail of a live one. [start] is where [text] begins
 * in the message. [continues] says the first piece is an item of the list the segment before it
 * ended with, so the two draw as one list: no block gap between them, and neither the item above
 * the seam nor the one below it takes the padding of a list's edge.
 */
private class Segment(
    val text: String,
    val start: Int,
    val parse: State,
    val pieces: List<Piece>,
    val continues: Boolean = false,
)

/**
 * The live reply's segments: parsed on the composing thread the first time the row is drawn, and
 * incrementally off it for every delta afterwards.
 *
 * The first parse has to be inline. The renderer's own asynchronous path draws an empty loading
 * slot until its result arrives, so a row is measured at nothing before it is measured at its real
 * height, and the transcript above it collapses and springs back. Seen with five replies on screen
 * at once, every one of them blank, the whole conversation shrunk to fit a single screen; a moment
 * later it was all there again. That is the "skipping up and down" this list must never do.
 *
 * Every parse after the first is off the composing thread, and the row keeps drawing the parse it
 * already has until the new one lands, so there is never a frame without a height. What is on
 * screen is always a real prefix of the reply rather than a guess at it; it is simply one parse
 * behind.
 */
@Composable
private fun liveSegments(text: String): List<Segment> {
    val parsed = remember {
        mutableStateOf(
            DebugStats.timed("markdown parsed while composing") { LiveParse.whole(text) }
        )
    }
    LaunchedEffect(text) {
        if (parsed.value.text == text) return@LaunchedEffect
        val previous = parsed.value
        parsed.value =
            withContext(Dispatchers.Default) {
                DebugStats.timed("markdown reparsed while streaming") { previous.advanceTo(text) }
            }
    }
    return parsed.value.segments
}

/**
 * A reply still arriving, parsed a block at a time.
 *
 * Reparsing the whole message per delta was fine for a short reply and not for a long one: a
 * twenty-five-screen reply parses in tens of milliseconds, hundreds of times, and although that ran
 * off the composing thread it was every core busy while the frame's own thread waited for one.
 * Markdown's blocks make the cut safe: a top-level block that another block has started *after* is
 * finished -- nothing appended later can reach back into it, since a paragraph ends at the blank
 * line or the block that interrupts it, a fence at its closing fence, a list at the first line that
 * is neither an item nor indented under one. So every block but the last is [frozen] with the parse
 * that finished it, and only the tail -- the last block and whatever has arrived since -- is parsed
 * again.
 *
 * A list is cut once more, at its last item, by the same reasoning one level down: an item is
 * finished once the next item has begun, since a line can only continue the item it is indented
 * under or start a new one. Without this a reply that is one long list -- forty sources -- parsed
 * the whole list per delta, and a list streams as forty paragraphs would. The item the cut lands on
 * has to have begun in earnest: a bare `-` is an empty item now and the first character of a
 * paragraph line once `-x` arrives, and cutting on it would draw that line as a new item.
 *
 * What the cut gives up is one thing: a reference definition arriving later than a link that uses
 * it, since the frozen block's parse never sees it. The link draws as its brackets until the reply
 * settles and is parsed whole by [warm], which is the same moment every other transient of
 * streaming is put right.
 */
private class LiveParse(
    val text: String,
    private val frozen: List<Segment>,
    /** How much of [text] the frozen segments cover; the tail starts here. */
    private val consumed: Int,
    private val tail: Segment,
) {
    val segments: List<Segment>
        get() = frozen + tail

    fun advanceTo(next: String): LiveParse {
        // Anything but an append to what was frozen -- a message replaced, a stream reset --
        // starts over.
        if (!next.regionMatches(0, text, 0, consumed)) return whole(next)
        val tailText = next.substring(consumed)
        val parse = parseMarkdown(tailText)
        val all = pieces(parse)
        val open = (parse as? State.Success)?.let { openPiece(it, all) }
        if (open == null) {
            return LiveParse(
                next,
                frozen,
                consumed,
                Segment(tailText, consumed, parse, all, tail.continues),
            )
        }
        val done =
            all.subList(0, all.indexOf(open))
                .groupBy { it.block }
                .values
                .mapIndexed { at, pieces ->
                    Segment(
                        tailText,
                        consumed,
                        parse,
                        pieces,
                        continues = at == 0 && tail.continues,
                    )
                }
        // Cut at the start of the open piece's line rather than at the piece, so an indented item
        // or block keeps the indentation the parse of the rest reads its nesting from.
        val node =
            parse.node.children[open.block].let {
                if (open.item == Piece.WHOLE_BLOCK) it else it.listItems()[open.item]
            }
        val cut = tailText.lastIndexOf('\n', node.startOffset) + 1
        val rest = tailText.substring(cut)
        val restParse = parseMarkdown(rest)
        return LiveParse(
            next,
            frozen + done,
            consumed + cut,
            Segment(rest, consumed + cut, restParse, pieces(restParse), continues = open.item > 0),
        )
    }

    /**
     * The piece of the tail still being written: the last item of a list of several, or the first
     * piece of the last block when there is more than one block. Null when nothing before it is
     * finished, so the tail stays whole.
     */
    private fun openPiece(parse: State.Success, all: List<Piece>): Piece? {
        val last = all.lastOrNull() ?: return null
        val lastBlockStart = all.indexOfFirst { it.block == last.block }
        return when {
            last.item > 0 && parse.node.children[last.block].listItems()[last.item].hasBegun -> last
            lastBlockStart > 0 -> all[lastBlockStart]
            else -> null
        }
    }

    /** Whether a list item holds anything beyond its marker yet. */
    private val ASTNode.hasBegun: Boolean
        get() = children.any { it.type !in MARKER_TOKENS }

    companion object {
        fun whole(text: String): LiveParse {
            val parse = parseMarkdown(text)
            return LiveParse(text, emptyList(), 0, Segment(text, 0, parse, pieces(parse)))
        }

        private val MARKER_TOKENS =
            setOf(
                MarkdownTokenTypes.LIST_BULLET,
                MarkdownTokenTypes.LIST_NUMBER,
                MarkdownTokenTypes.WHITE_SPACE,
                MarkdownTokenTypes.EOL,
            )
    }
}

/** One [piece] of [text], drawn on its own -- a unit of the transcript list. */
@Composable
fun MarkdownPiece(
    text: String,
    piece: Piece,
    replies: ParsedReplies,
    modifier: Modifier = Modifier,
) {
    // Remembered so a message the flatten drew before [warm] reached it is parsed once here, not
    // once per composition.
    val parse = remember(text) { replies.of(text) }
    MarkdownRoot(parse, replies) { MarkdownPiece(parse, text, piece, modifier) }
}

/**
 * The renderer's own environment -- its colours, type scale, dimensions, component table and
 * reference links -- around whatever draws pieces of [parse].
 *
 * The parsing is the library's. Markdown is somebody else's specification, and a hand-written
 * parser would get the edge cases wrong one case at a time. So is the environment: the element
 * composables its dispatch reaches read these locals, and providing them once here is what lets a
 * piece be drawn anywhere -- in a message's column, or as one item of the transcript list.
 * Everything below this is the mapping onto the app's palette and type scale.
 *
 * The locals are provided directly rather than through the renderer's `Markdown()` composable,
 * which was the last of its composables on the hot path and was here only to provide them. What
 * that buys is that nothing between a piece and the screen is the library's but the leaf
 * composables named in the component table, so a different parser could stand behind [State]
 * without the renderer's entry point being involved.
 *
 * Colours come from the theme rather than from the renderer's defaults, so code, links and rules
 * are the same Catppuccin values the rest of the app uses. Nothing here picks a colour of its own.
 *
 * [streaming] says this parse is the part of a reply still being written, which only the fences
 * care about: lexing is proportional to how much code there is, and a fence still arriving is
 * re-lexed at every delta on the composing thread. Measured streaming a two-hundred-line Kotlin
 * fence: **13.7 seconds** of lexing across the turn, 211 of them, the worst 177ms -- for colours on
 * text that was being replaced as fast as they were computed. So a fence still being written is
 * drawn plain and takes its colours when the block freezes, which is the same bargain [LiveParse]
 * already makes for a reference link defined at the foot of a message.
 */
@Composable
private fun MarkdownRoot(
    parse: State,
    replies: ParsedReplies,
    streaming: Boolean = false,
    content: @Composable () -> Unit,
) {
    if (parse !is State.Success) {
        // Nothing below needs the environment; [MarkdownPiece] draws the words plainly.
        content()
        return
    }
    val body = MaterialTheme.typography.bodyLarge
    CompositionLocalProvider(
        LocalReferenceLinkHandler provides parse.referenceLinkHandler,
        LocalMarkdownPadding provides markdownPadding(),
        // Read by the renderer's own text composable, which no paragraph reaches any more, and
        // by its checkbox. Provided so a path that does reach them draws no image rather than
        // failing to compose.
        LocalImageTransformer provides remember { NoOpImageTransformerImpl() },
        LocalMarkdownAnimations provides markdownAnimations(),
        LocalMarkdownColors provides
            markdownColor(
                text = MaterialTheme.colorScheme.onSurface,
                dividerColor = MaterialTheme.colorScheme.outlineVariant,
                // The dark surface every verbatim thing in this app sits on -- see [rawSurface],
                // and the tool call above this reply, which now matches. `surfaceVariant` was
                // exactly a card's own fill, so a fenced block inside a tool call had no
                // background at all and one in a reply read as a step *up* out of the page.
                codeBackground = rawSurface,
                // The same colour, but let through. An inline span's background is part of the
                // *text's* own drawing and the selection rectangle is drawn underneath it, so an
                // opaque chip hides the selection completely: selecting a sentence highlighted
                // every word of it except the ones in backticks, which is a difference in
                // appearance the reader has no way to account for. Translucent, the selection
                // shows through and the chip still reads as one step down from the page -- there
                // is no way to draw it over the selection instead, since the order is the text
                // node's.
                inlineCodeBackground = rawSurface.copy(alpha = INLINE_CODE_ALPHA),
                // The same tint a code block gets, rather than the renderer's 2%-alpha default:
                // two adjacent tints that differ by a fiftieth read as one flat block on a phone,
                // so the table would have had a border-less grid and nothing saying where it began.
                tableBackground = MaterialTheme.colorScheme.surfaceVariant,
            ),
        LocalMarkdownTypography provides
            markdownTypography(
                // A ladder that starts near the body text and descends, because these are headings
                // inside a chat message rather than the top of a document. The renderer's defaults
                // are the Material *display* styles -- `#` came out at 57sp and `##` at 45sp, which
                // is bigger than this app's own screen titles and reads as the reply shouting.
                //
                // Every step is a different size, so two levels of nesting never draw the same:
                // one clear step per level is the whole job of a heading.
                h1 = MaterialTheme.typography.headlineSmall,
                h2 = MaterialTheme.typography.titleLarge,
                h3 = MaterialTheme.typography.titleMedium,
                h4 = MaterialTheme.typography.titleSmall,
                h5 = MaterialTheme.typography.labelMedium,
                h6 = MaterialTheme.typography.labelSmall,
                // Body text at the size everything else in the transcript uses.
                text = body,
                paragraph = body,
                ordered = body,
                bullet = body,
                list = body,
                table = body,
                // Code in a monospace face, in the ordinary text colour. The face and the tinted
                // background are what say "this is code"; colour is not, and it used to be green
                // -- the palette's colour for a *literal*. A block of code is not a literal, it
                // is text that happens to be code, and painting all of it green said the whole
                // block was one. Where a literal really does appear inside code, the thing that
                // should colour it is a syntax highlighter looking at the code, which is exactly
                // what a tool call's input already gets from `catppuccinSyntax`.
                //
                // The colour rides on the style here rather than in `markdownColor`, which
                // stopped carrying `codeText`/`inlineCodeText`/`linkText` when the renderer moved
                // them onto the typography.
                code =
                    MaterialTheme.typography.bodyMedium.copy(
                        fontFamily = FontFamily.Monospace,
                        color = MaterialTheme.colorScheme.onSurface,
                    ),
                inlineCode =
                    body.copy(
                        fontFamily = FontFamily.Monospace,
                        // Unspecified so an inline span keeps the size of the line it sits in.
                        fontSize = TextUnit.Unspecified,
                        color = MaterialTheme.colorScheme.onSurface,
                    ),
                textLink =
                    TextLinkStyles(
                        style =
                            body
                                .copy(
                                    color = linkColor,
                                    textDecoration = TextDecoration.Underline,
                                )
                                .toSpanStyle()
                    ),
            ),
        LocalMarkdownDimens provides
            markdownDimens(
                // Half the renderer's 16dp. Padding is charged on both sides of every cell, so at
                // the default a fifth of the narrowest column went on space rather than on words
                // -- and the narrowest column is where the wrapping below has the least room.
                tableCellPadding = 8.dp,
                // What a column narrows to before the table starts scrolling sideways instead. It
                // is the floor, not the width: a table with room to spare spreads across it.
                //
                // Down from the renderer's 160dp, and the number is a measurement rather than a
                // taste. A phone is about 410-450dp wide and a card takes some of that, so 160dp
                // makes even a three-column table -- the commonest shape there is -- scroll, while
                // 136dp fits three across the phone this app is read on. Four and up still scroll,
                // which is the right answer for genuinely too many columns: squeezing six columns
                // into a phone would give every cell one word per line.
                //
                // Narrower would fit more, and stop being readable. This is the widest minimum
                // that keeps three columns on screen, which is the trade the number is making.
                tableCellWidth = 136.dp,
            ),
        LocalMarkdownComponents provides
            markdownComponents(
                // The m3 renderer's own default, restored: supplying `components` at all replaces
                // the whole set, and this is the only member of it the Material layer overrides.
                checkbox = { MarkdownCheckBox(it.content, it.node, it.typography.text) },
                // Everything that draws a run of text, so a link is a span rather than a node --
                // see [LinkedText]. Setext headings take the same styles as `#` and `##`, which
                // is the renderer's own pairing.
                text = { LinkedText(it, it.typography.text) },
                paragraph = { LinkedText(it, it.typography.paragraph) },
                heading1 = { LinkedHeading(it, it.typography.h1) },
                heading2 = { LinkedHeading(it, it.typography.h2) },
                heading3 = { LinkedHeading(it, it.typography.h3) },
                heading4 = { LinkedHeading(it, it.typography.h4) },
                heading5 = { LinkedHeading(it, it.typography.h5) },
                heading6 = { LinkedHeading(it, it.typography.h6) },
                setextHeading1 = { LinkedHeading(it, it.typography.h1) },
                setextHeading2 = { LinkedHeading(it, it.typography.h2) },
                // Lists are ours wherever the renderer's dispatch meets one -- inside a quote --
                // so they draw like the top-level ones the transcript cuts into items.
                orderedList = { MarkdownList(it.content, it.node, it.listDepth) },
                unorderedList = { MarkdownList(it.content, it.node, it.listDepth) },
                table = { LinkedTable(it.content, it.node, it.typography.table) },
                // Code is highlighted the way a tool call's input is; see [CodeFence].
                codeFence = {
                    CodeFence(it.content, it.node, it.typography.code, replies, streaming)
                },
                codeBlock = {
                    CodeBlock(it.content, it.node, it.typography.code, replies, streaming)
                },
            ),
        content = content,
    )
}

/**
 * A table: its rows, on the renderer's tinted, rounded background, as wide as its columns need.
 *
 * Each column has a floor ([markdownDimens]'s `tableCellWidth`), so the table is at least
 * columns-times-floor wide; narrower than the room it has, it spreads to fill it, and wider, it
 * scrolls sideways rather than squeezing. The renderer decided that with a `BoxWithConstraints`,
 * which is a subcomposition; here it is one layout modifier, and the trick is where it sits.
 * `fillMaxWidth` fixes the minimum width to the room available, the horizontal scroll passes that
 * minimum through to its content while lifting the maximum to unbounded, and the modifier after it
 * reads the minimum back as the room and sizes the rows to the larger of that and the floor. The
 * scroll then has exactly the overflow to scroll, which is none when the table fits.
 */
@Composable
private fun LinkedTable(content: String, node: ASTNode, style: TextStyle) {
    val dimens = LocalMarkdownDimens.current
    val colors = LocalMarkdownColors.current
    val columns =
        remember(node) {
            node.findChildOfType(GFMElementTypes.HEADER)?.children?.count {
                it.type == GFMTokenTypes.CELL
            } ?: 0
        }
    val rows = remember(node) { node.children.count { it.type == GFMElementTypes.ROW } + 1 }
    val floor = dimens.tableCellWidth * columns
    Column(
        Modifier.background(colors.tableBackground, RoundedCornerShape(dimens.tableCornerSize))
            .semantics { collectionInfo = CollectionInfo(rowCount = rows, columnCount = columns) }
            .fillMaxWidth()
            .horizontalScroll(rememberScrollState())
            .layout { measurable, constraints ->
                val width = maxOf(constraints.minWidth, floor.roundToPx())
                val placeable =
                    measurable.measure(constraints.copy(minWidth = width, maxWidth = width))
                layout(width, placeable.height) { placeable.place(0, 0) }
            }
    ) {
        var rowIndex = 1
        node.children.forEach { child ->
            when (child.type) {
                GFMElementTypes.HEADER -> LinkedTableRow(content, child, style, rowIndex = 0)
                GFMElementTypes.ROW -> LinkedTableRow(content, child, style, rowIndex = rowIndex++)
                GFMTokenTypes.TABLE_SEPARATOR -> MarkdownDivider()
            }
        }
    }
}

/**
 * One row of a table -- the header when [rowIndex] is zero -- with every cell a [LinkedText].
 *
 * The renderer's own rows draw each cell at `maxLines = 1` with an ellipsis, which on a phone means
 * most of a table is simply not readable: anything past about twenty characters ends in "..." with
 * no way to see the rest, and an elided cell looks like a short one, so a table of measurements
 * reads as a table of plausible shorter measurements. And they draw a link in a cell as its own
 * layout node, the cost [LinkedText] exists to avoid.
 *
 * So: as many lines as the cell needs, cells aligned to the top of the row, because a two-line cell
 * beside a one-line one centred the short one against the middle of the tall one and lost the line
 * the reader was reading across. What the wrapping does *not* do is make a wide table fit;
 * [LinkedTable] scrolls it instead, which is the right answer for too many columns -- wrapping a
 * six-column table into the width of a phone would give every cell one word per line.
 *
 * The semantics are the renderer's: each cell is an item of the table's collection, and a header
 * cell is a heading.
 */
@Composable
private fun LinkedTableRow(content: String, row: ASTNode, style: TextStyle, rowIndex: Int) {
    val padding = LocalMarkdownDimens.current.tableCellPadding
    val header = rowIndex == 0
    val cellStyle = if (header) style.copy(fontWeight = FontWeight.Bold) else style
    Row(verticalAlignment = Alignment.Top, modifier = Modifier.fillMaxWidth()) {
        row.children
            .filter { it.type == GFMTokenTypes.CELL }
            .forEachIndexed { column, cell ->
                LinkedText(
                    content,
                    cell,
                    cellStyle,
                    Modifier.padding(padding).weight(1f).semantics {
                        if (header) heading()
                        collectionItemInfo =
                            CollectionItemInfo(
                                rowIndex = rowIndex,
                                rowSpan = 1,
                                columnIndex = column,
                                columnSpan = 1,
                            )
                    },
                )
            }
    }
}

/**
 * Replies parsed before the row that draws them is composed.
 *
 * Parsing is the expensive half of drawing a reply, and it is expensive in proportion to how much
 * was written. Measured against a real Claude Code transcript on the emulator, one message took
 * **51ms** and several took 10-25ms, against 4.6ms for the short synthetic replies this was first
 * tuned on -- so a page of history landing composed several rows that each stalled the frame they
 * appeared in. That is the lag when a block loads.
 *
 * Nothing here changes what a row does when it has no answer waiting: it parses inline, on the
 * composing thread, because a row measured at nothing before it is measured at its real height
 * collapses the transcript above it. The point is only that by the time the reader scrolls to a
 * row, the answer is usually already made -- [warm] runs on a background thread as each page of
 * history arrives, which is seconds before anybody reaches the rows it brought.
 *
 * A miss is not stored, and that is what bounds this: the map holds one entry per message a page
 * warmed and nothing else, so a reply still streaming cannot fill it with hundreds of copies of
 * itself on the way to being finished. It is dropped with the screen, and emptied by the stream
 * reset that drops the rows it describes.
 */
@Stable
class ParsedReplies {
    private val parsed = ConcurrentHashMap<String, State>()

    /**
     * How each message divides into pieces, cached beside its parse: [transcriptUnits] asks per
     * fold, and walking the tree again each time is proportional to the message where a lookup is
     * proportional to nothing.
     */
    private val pieces = ConcurrentHashMap<String, List<Piece>>()

    /**
     * How each message divides into prose and memory notes, cached for the same reason as
     * [piecesOf]: the regex scan behind [messageParts] is proportional to the message.
     */
    private val parts = ConcurrentHashMap<String, List<MessagePart>>()

    private val chunks = ConcurrentHashMap<String, List<String>>()

    /**
     * Each fence's coloured text, keyed by its language and code.
     *
     * Beside the parses for the same reason and at the same cost: lexing is proportional to how
     * much code was written -- a two-hundred-line Kotlin fence measured 174ms on the emulator --
     * and a lazy list drops the composition of a block that scrolls away, so a `remember` inside
     * the fence paid that again every time the reader came back to it. Six times in one scroll,
     * measured. [warm] fills this off the drawing thread before the row is reached.
     */
    private val highlights = ConcurrentHashMap<String, AnnotatedString>()

    private val ready = ConcurrentHashMap.newKeySet<String>()

    /** The pieces of [text], from its parse -- made now if [warm] has not made it. */
    fun piecesOf(text: String): List<Piece> =
        pieces.computeIfAbsent(text) {
            DebugStats.timed("markdown cut into pieces") { pieces(of(it)) }
        }

    /** How a long user message divides into slices; cached for the same reason as [piecesOf]. */
    fun chunksOf(text: String): List<String> =
        chunks.computeIfAbsent(text) {
            DebugStats.timed("user message cut into slices") { userChunks(it) }
        }

    /**
     * Whether [warm] has made everything drawing [text] as pieces will look up.
     *
     * What the flatten asks before drawing a reply that way. Cutting costs a parse of the whole
     * message and the flatten runs on the composing thread -- so a reply not marked yet stays
     * whole, drawing the parse it already has, until the screen has warmed it and re-flattens. An
     * explicit mark rather than a peek into the parse cache, because a message with memory notes is
     * warmed as its *parts*: nothing ever parses its full text, and inferring readiness from the
     * cache left exactly that message unsplittable forever, re-warmed on every fold.
     */
    fun splitReady(text: String): Boolean = text in ready

    /** The other half of [splitReady]; [warm] calls it once a message's parses exist. */
    fun markSplitReady(text: String) {
        ready.add(text)
    }

    fun partsOf(text: String): List<MessagePart> =
        parts.computeIfAbsent(text) {
            DebugStats.timed("message cut into parts") { messageParts(it) }
        }

    /**
     * [code] coloured for [language] -- the answer made ahead, or one made now.
     *
     * The key carries the language, because the same code lexes differently under two of them.
     */
    fun highlighted(code: String, language: Language?): AnnotatedString =
        if (language == null) AnnotatedString(code)
        else highlights.computeIfAbsent("$language\n$code") { highlight(code, language) }

    /** The parse of [text] -- the one made ahead, or one made now. */
    fun of(text: String): State =
        parsed[text]?.also { DebugStats.count("markdown ready") }
            ?: DebugStats.timed("markdown parsed while composing") { parseMarkdown(text) }

    /**
     * Parses whatever is not held yet. Call off the composing thread; that is the whole point.
     *
     * Suspending, and yielding between messages, because "off the composing thread" is not the same
     * as "free". A page of history arrives as hundreds of parses at once -- 1.5 seconds of them in
     * a twelve second scroll, measured on a Pixel 9 Pro XL -- and on the default dispatcher that is
     * every core busy, with the frame's own thread waiting for one. That showed up as 21ms of
     * `waited` at the 90th percentile: the frame could not start, rather than taking too long.
     */
    suspend fun warm(texts: List<String>) {
        texts.forEach { text ->
            val parse =
                parsed.computeIfAbsent(text) {
                    DebugStats.timed("markdown warmed") { parseMarkdown(it) }
                }
            // The fences too, and here rather than in a pass of its own: they are found in the
            // parse this just made, and lexing one is the same kind of cost as parsing the
            // message it is in -- proportional to what was written, and charged to the frame
            // that first draws it if nobody paid it earlier.
            fences(parse).forEach { (code, language) -> highlighted(code, language) }
        }
    }

    /** Everything these described is gone; see [ParsedReplies]. */
    fun clear() {
        parsed.clear()
        pieces.clear()
        parts.clear()
        chunks.clear()
        highlights.clear()
        ready.clear()
    }
}

/**
 * How much of the inline-code chip's fill is its own colour, the rest being whatever it sits on.
 *
 * High enough that the chip is still a clear step down from the page, low enough that a selection
 * under it changes what the chip looks like. Both halves are the point: at 1.0 the chip was the
 * only part of a selected sentence that did not look selected.
 */
private const val INLINE_CODE_ALPHA = 0.6f