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

package com.example.aiapp

import androidx.compose.material3.MaterialTheme
import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.Stable
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.text.TextLinkStyles
import androidx.compose.ui.text.TextStyle
import androidx.compose.ui.text.font.FontFamily
import androidx.compose.ui.text.style.TextDecoration
import androidx.compose.ui.text.style.TextOverflow
import androidx.compose.ui.unit.Dp
import androidx.compose.ui.unit.TextUnit
import androidx.compose.ui.unit.dp
import com.mikepenz.markdown.compose.components.markdownComponents
import com.mikepenz.markdown.compose.elements.MarkdownTable
import com.mikepenz.markdown.compose.elements.MarkdownTableHeader
import com.mikepenz.markdown.compose.elements.MarkdownTableRow
import com.mikepenz.markdown.m3.Markdown
import com.mikepenz.markdown.m3.elements.MarkdownCheckBox
import com.mikepenz.markdown.m3.markdownColor
import com.mikepenz.markdown.m3.markdownTypography
import com.mikepenz.markdown.model.State
import com.mikepenz.markdown.model.markdownDimens
import com.mikepenz.markdown.model.parseMarkdown
import java.util.concurrent.ConcurrentHashMap
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import org.intellij.markdown.ast.ASTNode

/**
 * An assistant's reply, rendered as the markdown it is written in.
 *
 * The parsing is the library's. Markdown is somebody else's specification, and a hand-written
 * subset of one disagrees with it at the edges -- which is where the bug reports come from, one
 * case at a time. This file's whole job is the mapping onto the app's palette and type scale.
 *
 * 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.
 */
@Composable
fun MarkdownText(text: String, replies: ParsedReplies, modifier: Modifier = Modifier) {
    val body = MaterialTheme.typography.bodyLarge
    val parsed = parsedMarkdown(text, replies)
    Markdown(
        parsed,
        colors =
            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,
                inlineCodeBackground = rawSurface,
                // 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,
            ),
        typography =
            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()
                    ),
            ),
        dimens =
            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,
            ),
        components =
            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) },
                table = {
                    MarkdownTable(
                        it.content,
                        it.node,
                        style = it.typography.table,
                        headerBlock = ::WrappingTableHeader,
                        rowBlock = ::WrappingTableRow,
                    )
                },
            ),
        modifier = modifier,
    )
}

/**
 * A table header, and a table row, whose cells wrap rather than being cut off.
 *
 * The renderer draws every cell at `maxLines = 1` with an ellipsis, which on a phone means most of
 * a table is simply not readable: a column is 160dp at its narrowest, so anything past about twenty
 * characters ends in "..." with no way to see the rest. Nothing about the value says it was cut,
 * either -- an elided cell looks like a short one, so a table of measurements reads as a table of
 * plausible shorter measurements.
 *
 * So: as many lines as the cell needs, and [TextOverflow.Clip] rather than an ellipsis, which now
 * never has anything to hide since the height grows to fit. Cells align 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. The renderer already gives each column
 * a 160dp minimum and scrolls the whole table sideways when they do not fit the screen, 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 two work together: the width is what the columns
 * need, and the wrapping is what fills the space that width provides.
 *
 * Two functions rather than one because the renderer's header and row are separate composables --
 * the header is bold and sizes itself to its tallest cell -- and the parameters that matter here
 * are the same three in both.
 */
@Composable
private fun WrappingTableHeader(
    content: String,
    header: ASTNode,
    tableWidth: Dp,
    style: TextStyle,
) {
    MarkdownTableHeader(
        content = content,
        header = header,
        tableWidth = tableWidth,
        style = style,
        verticalAlignment = Alignment.Top,
        maxLines = Int.MAX_VALUE,
        overflow = TextOverflow.Clip,
    )
}

@Composable
private fun WrappingTableRow(content: String, row: ASTNode, tableWidth: Dp, style: TextStyle) {
    MarkdownTableRow(
        content = content,
        header = row,
        tableWidth = tableWidth,
        style = style,
        verticalAlignment = Alignment.Top,
        maxLines = Int.MAX_VALUE,
        overflow = TextOverflow.Clip,
    )
}

/**
 * [text] parsed: on the composing thread the first time this row is drawn, and off it every time
 * 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, and no
 * amount of scroll anchoring can survive a row that lies about its height first.
 *
 * Every parse *after* the first is a different case, and it is the one that was costing: a reply
 * arrives as hundreds of deltas, each one re-parsing the whole message it has grown into. Measured
 * against `/stream 200` on the emulator, that was fifty-eight parses and 78ms of main-thread work
 * in three seconds, with single parses reaching 7ms -- most of a frame at 60Hz and more than one at
 * 120. Those go to a background 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 parsedMarkdown(text: String, replies: ParsedReplies): State {
    // The text each parse came from, so the first composition's is not immediately repeated.
    val parsed = remember { mutableStateOf(text to replies.of(text)) }
    LaunchedEffect(text) {
        if (parsed.value.first == text) return@LaunchedEffect
        // Not through [replies]: this is a reply still arriving, and every delta would leave
        // another copy of a message that is about to be superseded.
        parsed.value =
            text to
                withContext(Dispatchers.Default) {
                    DebugStats.timed("markdown reparsed while streaming") { parseMarkdown(text) }
                }
    }
    return parsed.value.second
}

/**
 * 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 blocks, cached beside the parses of those blocks.
     *
     * Here rather than in a `remember` because the answer is wanted on two threads: by [warm], to
     * know which strings to make ready, and by the row that draws them. Finding it costs a parse of
     * the whole message, so doing it twice would undo what splitting is for.
     */
    private val blocks = ConcurrentHashMap<String, List<String>>()

    /**
     * How each message divides into prose and memory notes, cached for the same reason as
     * [blocksOf]: [transcriptUnits] asks per fold, and the regex scan behind [messageParts] is
     * proportional to the message every time where a lookup is proportional to nothing.
     */
    private val parts = ConcurrentHashMap<String, List<MessagePart>>()

    fun blocksOf(text: String): List<String> = blocks.computeIfAbsent(text) { markdownBlocks(it) }

    fun partsOf(text: String): List<MessagePart> = parts.computeIfAbsent(text) { messageParts(it) }

    /** 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 ->
            parsed.computeIfAbsent(text) {
                DebugStats.timed("markdown warmed") { parseMarkdown(it) }
            }
        }
    }

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