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.Modifier
import androidx.compose.ui.text.TextLinkStyles
import androidx.compose.ui.text.font.FontFamily
import androidx.compose.ui.text.style.TextDecoration
import androidx.compose.ui.unit.TextUnit
import com.mikepenz.markdown.m3.Markdown
import com.mikepenz.markdown.m3.markdownColor
import com.mikepenz.markdown.m3.markdownTypography
import com.mikepenz.markdown.model.State
import com.mikepenz.markdown.model.parseMarkdown
import java.util.concurrent.ConcurrentHashMap
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext

/**
 * 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()
                    ),
            ),
        modifier = modifier,
    )
}

/**
 * [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()
    }
}