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

package com.example.aiapp

import androidx.compose.foundation.layout.Box
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.text.BasicText
import androidx.compose.material3.MaterialTheme
import androidx.compose.material3.Text
import androidx.compose.runtime.Composable
import androidx.compose.runtime.Immutable
import androidx.compose.ui.Modifier
import androidx.compose.ui.semantics.isTraversalGroup
import androidx.compose.ui.semantics.semantics
import androidx.compose.ui.text.TextStyle
import androidx.compose.ui.unit.Dp
import androidx.compose.ui.unit.dp
import com.mikepenz.markdown.compose.LocalMarkdownColors
import com.mikepenz.markdown.compose.LocalMarkdownComponents
import com.mikepenz.markdown.compose.LocalMarkdownPadding
import com.mikepenz.markdown.compose.LocalMarkdownTypography
import com.mikepenz.markdown.compose.MarkdownElement
import com.mikepenz.markdown.compose.components.MarkdownComponentModel
import com.mikepenz.markdown.model.State
import org.intellij.markdown.MarkdownElementTypes
import org.intellij.markdown.MarkdownTokenTypes
import org.intellij.markdown.ast.ASTNode
import org.intellij.markdown.ast.findChildOfType
import org.intellij.markdown.ast.getTextInNode
import org.intellij.markdown.flavours.gfm.GFMTokenTypes

/**
 * One drawable piece of a parsed message: a top-level block, or one item of a top-level list.
 *
 * The point is the draw phase and the lazy list. A reply's display list holds every glyph of it and
 * is re-recorded whenever drawing is invalidated, so one long message costs as much to draw as a
 * hundred short ones; and the list composes an item whole in the frame it scrolls into, so an item
 * has to be bounded for the worst frame to be. Measured on a Pixel 9 Pro XL, the tallest row still
 * being drawn was 36,982px, twenty-five screens in one message. A piece is a paragraph, a fence, a
 * table, one bullet: bounded, so both costs are.
 *
 * Cut where the parser says the blocks are, which is the whole reason this is safe: a fence, a
 * table and a nested list are each one node whatever is inside them, so nothing is ever split down
 * the middle. A list is the one block that is not bounded -- a reply's list of sources can be forty
 * items -- so it is cut once more, into its items, and a nested list stays inside the item that
 * holds it.
 *
 * A piece is an *address* into the message's one parse ([block] indexes the root's children, [item]
 * the list items of that child) rather than a substring of the message. Every piece of a message is
 * drawn from the same tree, so a message is parsed once however many pieces it is drawn as, and a
 * reference definition at its foot still resolves the links above it -- the two costs of cutting a
 * message into strings and parsing each on its own.
 */
@Immutable
data class Piece(val block: Int, val item: Int = WHOLE_BLOCK) {
    companion object {
        const val WHOLE_BLOCK = -1
    }
}

/**
 * The pieces of [parse], in reading order. Blank nodes between blocks -- the parser keeps the
 * newlines -- are not pieces.
 *
 * A parse that failed yields one piece, so [MarkdownPiece] can still say what the message was: a
 * message that drew as nothing would be a hole in the transcript with no sign of what fell out.
 */
fun pieces(parse: State): List<Piece> {
    val success = parse as? State.Success ?: return listOf(Piece(0))
    val out = ArrayList<Piece>()
    success.node.children.forEachIndexed { at, node ->
        when {
            node.getTextInNode(success.content).isBlank() -> {}
            node.isList -> repeat(node.listItems().size) { out += Piece(at, it) }
            else -> out += Piece(at)
        }
    }
    return out
}

/**
 * The room above [piece] when it follows [previous] in the same message: none between two items of
 * one list, whose own padding already separates them, and a block's gap otherwise. The first piece
 * of a message takes the message's gap, which is the caller's to know.
 */
fun gapBefore(previous: Piece?, piece: Piece): Dp =
    if (previous != null && previous.block == piece.block) 0.dp else BLOCK_SPACING

/** The gap between one block of a reply and the next, wherever a reply is drawn in pieces. */
val BLOCK_SPACING: Dp = 6.dp

/**
 * [piece] of [parse], drawn. Must be inside [MarkdownRoot] for the parse, which is what carries the
 * theme, the components and the reference links to the renderer's element composables.
 *
 * A whole block goes to the renderer's own dispatch with this app's component table, so a paragraph
 * or heading is a [LinkedText], a table is [LinkedTableRow]s, and a nested list comes back here
 * through [MarkdownList]. Only the list item is drawn directly, because a list item is the one
 * piece the renderer has no element for.
 */
@Composable
fun MarkdownPiece(parse: State, text: String, piece: Piece, modifier: Modifier = Modifier) {
    if (parse !is State.Success) {
        // The parser threw. Nothing else in the app has seen this happen; if it does, the words
        // are still worth more than a blank.
        Text(text, modifier, style = MaterialTheme.typography.bodyLarge)
        return
    }
    val node = parse.node.children[piece.block]
    if (piece.item == Piece.WHOLE_BLOCK) {
        Box(modifier) {
            MarkdownElement(
                node,
                LocalMarkdownComponents.current,
                parse.content,
                includeSpacer = false,
            )
        }
    } else {
        val items = node.listItems()
        MarkdownListItem(
            content = parse.content,
            list = node,
            item = items[piece.item],
            index = piece.item,
            first = piece.item == 0,
            last = piece.item == items.lastIndex,
            depth = 0,
            modifier = modifier,
        )
    }
}

/**
 * A whole list, for the places the renderer's dispatch reaches one it cannot hand to a piece: a
 * list inside a quote, and the nested lists an item holds. Top-level lists never come here; they
 * are drawn an item at a time as pieces.
 */
@Composable
fun MarkdownList(content: String, list: ASTNode, depth: Int, modifier: Modifier = Modifier) {
    val items = list.listItems()
    Column(modifier) {
        items.forEachIndexed { index, item ->
            MarkdownListItem(
                content,
                list,
                item,
                index,
                first = index == 0,
                last = index == items.lastIndex,
                depth = depth,
            )
        }
    }
}

/**
 * One item: its marker beside its content, laid out the way the renderer's own list does so that a
 * list drawn as pieces looks exactly like one drawn whole. The list's own padding goes on its first
 * and last items, since there is no list column to carry it.
 *
 * The marker is the renderer's bullet and number, and a checkbox for a task item. It is drawn here
 * rather than by a handler because it is the thing a reader might one day want styled -- a
 * different glyph per depth, a colour -- and this is the one place it is drawn.
 */
@Composable
private fun MarkdownListItem(
    content: String,
    list: ASTNode,
    item: ASTNode,
    index: Int,
    first: Boolean,
    last: Boolean,
    depth: Int,
    modifier: Modifier = Modifier,
) {
    val padding = LocalMarkdownPadding.current
    val typography = LocalMarkdownTypography.current
    val components = LocalMarkdownComponents.current
    // A task item's box sits right after the bullet: `- [ ] text`.
    val checkbox = item.children.getOrNull(1)?.takeIf { it.type == GFMTokenTypes.CHECK_BOX }
    Row(
        modifier
            .semantics { isTraversalGroup = true }
            .fillMaxWidth()
            .padding(
                start = padding.listIndent * depth,
                top = padding.listItemTop + if (first) padding.list else 0.dp,
                bottom = padding.listItemBottom + if (last) padding.list else 0.dp,
            )
    ) {
        if (checkbox != null) {
            components.checkbox(MarkdownComponentModel(content, checkbox, typography))
        } else if (list.type == MarkdownElementTypes.ORDERED_LIST) {
            Marker("${list.startNumber(content) + index}. ", typography.ordered)
        } else {
            Marker("• ", typography.bullet)
        }
        Column {
            item.children.forEach { child ->
                when (child.type) {
                    MarkdownTokenTypes.LIST_BULLET,
                    MarkdownTokenTypes.LIST_NUMBER,
                    GFMTokenTypes.CHECK_BOX -> {}
                    MarkdownElementTypes.ORDERED_LIST,
                    MarkdownElementTypes.UNORDERED_LIST -> MarkdownList(content, child, depth + 1)
                    else -> MarkdownElement(child, components, content, includeSpacer = false)
                }
            }
        }
    }
}

/** The renderer's text colour on the marker; its styles carry none of their own. */
@Composable
private fun Marker(text: String, style: TextStyle) {
    BasicText(text, style = style.copy(color = LocalMarkdownColors.current.text))
}

private val ASTNode.isList: Boolean
    get() = type == MarkdownElementTypes.ORDERED_LIST || type == MarkdownElementTypes.UNORDERED_LIST

private fun ASTNode.listItems(): List<ASTNode> = children.filter {
    it.type == MarkdownElementTypes.LIST_ITEM
}

/** Where an ordered list counts from: the number its first item was written with. */
private fun ASTNode.startNumber(content: String): Int =
    findChildOfType(MarkdownElementTypes.LIST_ITEM)
        ?.findChildOfType(MarkdownTokenTypes.LIST_NUMBER)
        ?.getTextInNode(content)
        ?.takeWhile(Char::isDigit)
        ?.toString()
        ?.toIntOrNull() ?: 1