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

package com.example.aiapp

import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.fillMaxWidth
import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.ui.Modifier
import androidx.compose.ui.draw.drawWithContent
import androidx.compose.ui.graphics.graphicsLayer
import androidx.compose.ui.unit.dp
import com.mikepenz.markdown.model.State
import com.mikepenz.markdown.model.parseMarkdown
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext

/**
 * A message's top-level markdown blocks, cut where the parser says the blocks are.
 *
 * The point is the draw phase. A reply's display list holds every glyph of it, and it is
 * re-recorded whenever drawing is invalidated -- so one long message is as expensive to draw as a
 * hundred short ones, and skipping the rows around it cannot help while it is the one on screen.
 * Measured on a Pixel 9 Pro XL: 97% of rows correctly skipped, and the tallest row still being
 * drawn was 36,982px, about twenty-five screens in a single message. Cut into blocks, only the
 * screen or two actually being read is ever recorded.
 *
 * Cut at the parser's own boundaries rather than at blank lines, which is the whole reason this is
 * safe: a heading, a fenced code block, a table and a list are each one node whatever is inside
 * them, so a loose list does not become five one-item lists and a fence is never split down the
 * middle. Guessing at block boundaries with a line scanner gets all three of those wrong.
 *
 * It also bounds parsing, which was the other symptom: one message took **1.4 seconds** to parse as
 * a single unit, and a block is a paragraph.
 */
fun markdownBlocks(text: String): List<String> {
    // A reference definition sits at the foot of a message and is used by links above it. Parsed on
    // its own each block would lose the definition, and the link would draw as literal brackets --
    // so a message carrying one is kept whole. Rare enough to be worth giving up the split for.
    if (REFERENCE_DEFINITION.containsMatchIn(text)) return listOf(text)
    val parsed = parseMarkdown(text) as? State.Success ?: return listOf(text)
    val blocks =
        parsed.node.children
            .map { text.substring(it.startOffset, it.endOffset) }
            .filter { it.isNotBlank() }
    return if (blocks.size <= 1) listOf(text) else blocks
}

/** `[label]: https://…` at the start of a line -- see [markdownBlocks]. */
private val REFERENCE_DEFINITION = Regex("""^ {0,3}\[[^\]]+]:\s""", RegexOption.MULTILINE)

/**
 * A reply drawn a block at a time.
 *
 * Each block keeps its composition and its layout whichever way it is scrolled -- that is what
 * stops a message being rebuilt when somebody comes back to it. The heights come from the blocks
 * themselves as they are measured, so the running total is the same arrangement the list uses one
 * level up.
 *
 * [live] is the message currently arriving, and it is the only one that gets a layer per block. A
 * layer buys one thing here: when drawing is invalidated, only the block that changed is
 * re-recorded instead of the whole reply. That is worth a great deal while a reply is streaming,
 * because 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 BlockedMarkdown(
    text: String,
    replies: ParsedReplies,
    modifier: Modifier = Modifier,
    live: Boolean = false,
) {
    // The first split is inline for the reason [parsedMarkdown]'s first parse is: the row must
    // have its real height in its first frame. Every split after that is a delta landing, and it
    // runs off the composing thread with the message drawing the split it already has until the
    // new one arrives -- when this recomputed wherever composition ran, one streamed reply cost
    // 815 whole-message parses and 2.9 seconds of them, a few milliseconds per delta, on the
    // thread that draws. Not through [replies]: a reply mid-stream is a different text per
    // delta, and each would leave a cache entry nothing reads again.
    val split = remember { mutableStateOf(text to markdownBlocks(text)) }
    LaunchedEffect(text) {
        if (split.value.first == text) return@LaunchedEffect
        split.value =
            text to
                withContext(Dispatchers.Default) {
                    DebugStats.timed("blocks split while streaming") { markdownBlocks(text) }
                }
    }
    val blocks = split.value.second
    if (blocks.size == 1) {
        MarkdownText(blocks.first(), replies, modifier)
        return
    }
    Column(modifier.fillMaxWidth(), verticalArrangement = Arrangement.spacedBy(BLOCK_SPACING)) {
        blocks.forEach { block ->
            MarkdownText(
                block,
                replies,
                Modifier.fillMaxWidth()
                    .then(if (live) Modifier.graphicsLayer() else Modifier)
                    .drawWithContent {
                        val started = System.nanoTime()
                        drawContent()
                        DebugStats.record("record: one block", System.nanoTime() - started)
                    },
            )
        }
    }
}

/** The gap between one block of a reply and the next, here and in [transcriptUnits]. */
val BLOCK_SPACING = 6.dp