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.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. 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 what makes it safe: a fence, a table and a
 * nested list are each one node whatever is inside them. 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.
 *
 * A piece is an *address* into the message's one parse rather than a substring of it. Every piece
 * 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.
 */
@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 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 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. Only the list
 * item is drawn directly, because a list item is the one piece the renderer has no element for.
 *
 * [continuesList] and [listContinues] are for a list cut across the segments of a live reply: an
 * item that is the first or last of its own parse but not of the list the reader sees keeps an
 * inner item's padding, so nothing moves when the seam between segments does.
 */
@Composable
fun MarkdownPiece(
    parse: State,
    text: String,
    piece: Piece,
    modifier: Modifier = Modifier,
    continuesList: Boolean = false,
    listContinues: Boolean = false,
) {
    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 && !continuesList,
            last = piece.item == items.lastIndex && !listContinues,
            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.
 */
@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 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(BULLETS[depth % BULLETS.size], 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 marker in [listMarkerColor]; the renderer's styles carry no colour of their own. */
@Composable
private fun Marker(text: String, style: TextStyle) {
    BasicText(text, style = style.copy(color = listMarkerColor))
}

/**
 * The bullet at each depth, cycling past the third: a disc, a ring, a square -- the ladder a
 * browser draws, so a nested list is told from its parent by the glyph as well as by the indent.
 * Checked on the emulator's system fonts; a glyph the platform lacks draws as a box, and that check
 * is the price of adding one here.
 */
private val BULLETS = listOf("• ", "◦ ", "▪ ")

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

internal 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