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

package com.example.aiapp

import androidx.compose.foundation.lazy.LazyListItemInfo
import androidx.compose.runtime.Immutable
import androidx.compose.ui.unit.Dp
import androidx.compose.ui.unit.dp

/**
 * One item of the transcript list: a whole row, or one block of a settled reply.
 *
 * The unit of laziness is deliberately smaller than a message. A lazy list pays to compose an item
 * at the moment it scrolls into view, and that cost is proportional to the item -- a reply can be
 * twenty-five screens of markdown, which as one item is a hundred-millisecond frame exactly when
 * the list is moving fastest. A *block* is a paragraph, a fence, a table: bounded, so the worst
 * frame is bounded. This is the piece that was missing when a lazy list was last tried here.
 *
 * Everything else about the row model is unchanged: rows come from [groupToolRuns], and a unit
 * points back at its row. The list draws units; anchors and paging still speak seq.
 */
@Immutable
sealed class TranscriptUnit {
    /** The list identity; must survive pages landing at either end. See [TranscriptRow.key]. */
    abstract val key: Any

    /** Where this unit's row starts in the transcript -- the anchor identity, never the key. */
    abstract val seq: Long

    /**
     * This unit's position within its row, counted from the row's oldest end. What a saved scroll
     * position carries besides the seq: a reply split into forty blocks needs more than "somewhere
     * in this row" to put a reader back where they stopped.
     */
    abstract val ordinal: Int

    /** The gap drawn above this unit -- between rows, or between blocks of one reply. */
    abstract val gap: Dp

    /** A row drawn as itself: a bubble, a tool card, a group -- or the reply still arriving. */
    data class Whole(val row: TranscriptRow, override val gap: Dp) : TranscriptUnit() {
        override val key: Any
            get() = row.key

        override val seq: Long
            get() = row.startSeq

        override val ordinal: Int
            get() = 0
    }

    /** One [Piece] of a settled reply; [text] is the prose it is a piece of. */
    data class Block(
        override val seq: Long,
        override val ordinal: Int,
        val text: String,
        val piece: Piece,
        override val gap: Dp,
    ) : TranscriptUnit() {
        override val key: Any
            get() = "b$seq:$ordinal"
    }

    /**
     * The heading of a message from another agent: who sent it, and the control that opens it.
     *
     * A peer message is the one row whose *opened* size is unbounded -- these are the longest
     * things a transcript holds -- so it is flattened the same way a settled reply is, and for the
     * same reason. Measured on the emulator, opening a 43KB one took the transcript's share of the
     * draw phase from 0.81ms a frame to 3.85ms, and the framework's own per-frame bookkeeping from
     * 0.39ms to 3.15ms.
     *
     * The card is drawn in pieces rather than given up: a filled Material card is elevation zero,
     * so it has no shadow to break, and each piece paints the same fill with only the corners it
     * owns.
     */
    data class PeerHead(
        override val seq: Long,
        val item: TranscriptItem.PeerNote,
        val open: Boolean,
        override val gap: Dp,
    ) : TranscriptUnit() {
        /**
         * The note's own key, so opening and shutting does not change what the list is anchored on
         * -- and so two notes stamped with one turn's seq are still two items.
         */
        override val key: Any
            get() = item.key

        override val ordinal: Int
            get() = 0
    }

    /**
     * One [Piece] of an opened peer message; [last] is the piece that closes the card. Its [gap] is
     * always zero -- the pieces are one card -- so the room between blocks is [spacing], drawn
     * inside the piece where the card's fill covers it.
     */
    data class PeerBlock(
        override val seq: Long,
        override val ordinal: Int,
        val text: String,
        val piece: Piece,
        val last: Boolean,
        val spacing: Dp,
        override val gap: Dp,
        /** The note this block belongs to; its key, not its seq. See [TranscriptItem.PeerNote]. */
        val note: Any,
    ) : TranscriptUnit() {
        override val key: Any
            get() = "p$note:$ordinal"
    }

    /**
     * One slice of a long user message; see [userChunks].
     *
     * A user message is plain text, so cutting it costs a scan rather than a parse -- but the
     * reason is the same as for a settled reply: as one item, a pasted log is a hundred thousand
     * pixels of `Text` whose layout lands in the frame the row scrolls into.
     */
    data class UserChunk(
        override val seq: Long,
        override val ordinal: Int,
        val text: String,
        val first: Boolean,
        val last: Boolean,
        /** The message's attachments, drawn under the words -- so only the last slice has any. */
        val attachments: List<String>,
        override val gap: Dp,
    ) : TranscriptUnit() {
        override val key: Any
            get() = "u$seq:$ordinal"
    }

    /**
     * The "Show all N lines" under a message drawn only as far as [TextCap.MESSAGE_LINES].
     *
     * Its own unit rather than something inside the row above it, because the row above it is a
     * *bounded* item now and this is what says so -- and because a control that lives inside the
     * thing it reveals moves the moment it is pressed.
     */
    data class ShowAll(
        override val seq: Long,
        override val ordinal: Int,
        /** The row this belongs to; what goes into the set of rows shown whole. */
        val row: Any,
        /** The line count of the whole message, which is what the offer says. */
        val lines: Int,
        override val gap: Dp,
    ) : TranscriptUnit() {
        override val key: Any
            get() = "s$row"
    }

    /** One memory note of a settled reply; see [MemoryNote]. */
    data class Memory(
        override val seq: Long,
        override val ordinal: Int,
        val part: MessagePart.Remembered,
        override val gap: Dp,
    ) : TranscriptUnit() {
        override val key: Any
            get() = "m$seq:$ordinal"
    }
}

/**
 * A message row cut to [TextCap]'s worth of itself, and the line count of the whole of it.
 *
 * The cut happens **before** the flatten below decides how to draw the row, so everything after it
 * -- pieces, chunks, warming -- sees a shorter message and needs to know nothing about caps. The
 * shortened row keeps its key and its seq, so the list's identity and every saved scroll anchor are
 * untouched by a reader opening or closing one.
 *
 * A reply still arriving is never capped: it grows by deltas, and a row that stopped growing at two
 * hundred lines while the model was plainly still writing would read as the stream having died.
 * `iris`'s `row::build_row` states the same rule for the same reason.
 */
private fun capRow(row: TranscriptRow, shownWhole: Set<Any>): Pair<TranscriptRow, Int?> {
    val item = (row as? TranscriptRow.Single)?.item ?: return row to null
    if (row.key in shownWhole) return row to null
    val cut =
        when {
            item is TranscriptItem.UserMsg ->
                cutText(item.text, TextCap.MESSAGE_LINES, TextCap.MESSAGE_BYTES)?.let {
                    it to TranscriptRow.Single(item.copy(text = it.shown))
                }
            item is TranscriptItem.AssistantMsg && item.settled ->
                cutText(item.text, TextCap.MESSAGE_LINES, TextCap.MESSAGE_BYTES)?.let {
                    it to TranscriptRow.Single(item.copy(text = it.shown))
                }
            else -> null
        } ?: return row to null
    return cut.second to cut.first.lines
}

/**
 * The rows flattened into list units, newest first -- index zero is the item at the bottom of the
 * screen, which is what a reversed lazy list calls the start.
 *
 * Every settled reply is cut into its pieces (via the caches on [replies] so a message is only ever
 * cut once), and so is an *opened* peer message -- [openNotes] is which ones those are. A shut one
 * is a single heading and cannot be worth splitting. The reply still arriving stays whole: its text
 * changes with every delta, and splitting it here would parse the whole message per delta on
 * whichever thread is composing. Once settled it splits like every other reply, which is what
 * bounds the newest row's cost after a session ends on a long one.
 *
 * Runs per fold, so it must stay proportional to what is loaded with no parsing in it on the warm
 * path: [ParsedReplies.partsOf] and [ParsedReplies.piecesOf] are lookups for any text [warm] has
 * seen, and a miss -- the one message that just finished streaming -- costs its parse exactly once.
 */
fun transcriptUnits(
    rows: List<TranscriptRow>,
    replies: ParsedReplies,
    openNotes: Set<Long>,
    shownWhole: Set<Any> = emptySet(),
): List<TranscriptUnit> {
    val started = System.nanoTime()
    val units = ArrayList<TranscriptUnit>(rows.size)
    rows.forEachIndexed { index, whole ->
        val rowGap = if (index == 0) 0.dp else TRANSCRIPT_SPACING
        val (row, hidden) = capRow(whole, shownWhole)
        val rowStart = units.size
        val item = (row as? TranscriptRow.Single)?.item
        if (item is TranscriptItem.PeerNote) {
            val open = item.seq in openNotes
            units += TranscriptUnit.PeerHead(row.startSeq, item, open, rowGap)
            // No gap between the pieces: they are one card, and a card with a stripe through it is
            // what any spacing here would draw.
            if (open) {
                val pieces = replies.piecesOf(item.text)
                var previous: Piece? = null
                pieces.forEachIndexed { at, piece ->
                    units +=
                        TranscriptUnit.PeerBlock(
                            row.startSeq,
                            at + 1,
                            item.text,
                            piece,
                            last = at == pieces.lastIndex,
                            spacing = gapBefore(previous, piece),
                            gap = 0.dp,
                            note = item.key,
                        )
                    previous = piece
                }
            }
        } else if (item is TranscriptItem.UserMsg && item.text.length > USER_SPLIT_CHARS) {
            // A scan, not a parse, so it is cheap enough for the fold path -- and cached like the
            // markdown splits so the scan happens once per message rather than once per fold.
            val chunks = replies.chunksOf(item.text)
            chunks.forEachIndexed { at, chunk ->
                units +=
                    TranscriptUnit.UserChunk(
                        row.startSeq,
                        at,
                        chunk,
                        first = at == 0,
                        last = at == chunks.lastIndex,
                        attachments = if (at == chunks.lastIndex) item.attachments else emptyList(),
                        gap = if (at == 0) rowGap else 0.dp,
                    )
            }
        } else if (
            item is TranscriptItem.AssistantMsg &&
                splitWanted(item, index, rows.lastIndex) &&
                replies.splitReady(item.text)
        ) {
            var ordinal = 0
            fun gap(within: Dp) = if (ordinal == 0) rowGap else within
            replies.partsOf(item.text).forEach { part ->
                when (part) {
                    is MessagePart.Prose -> {
                        var previous: Piece? = null
                        replies.piecesOf(part.text).forEach { piece ->
                            units +=
                                TranscriptUnit.Block(
                                    row.startSeq,
                                    ordinal,
                                    part.text,
                                    piece,
                                    gap(gapBefore(previous, piece)),
                                )
                            ordinal++
                            previous = piece
                        }
                    }
                    is MessagePart.Remembered -> {
                        units +=
                            TranscriptUnit.Memory(row.startSeq, ordinal, part, gap(BLOCK_SPACING))
                        ordinal++
                    }
                }
            }
        } else {
            units += TranscriptUnit.Whole(row, rowGap)
        }
        if (hidden != null) {
            units +=
                TranscriptUnit.ShowAll(
                    row.startSeq,
                    units.size - rowStart,
                    row.key,
                    hidden,
                    BLOCK_SPACING,
                )
        }
    }
    units.reverse()
    reportDuplicateKeys(units)
    // Timed because this runs per fold on the composing thread: "loading messages feels bumpy" is
    // this number growing, and it was invisible until it was written down.
    DebugStats.record("units flattened", System.nanoTime() - started)
    return units
}

/**
 * Whether this reply should be drawn as blocks: settled, or anywhere but the newest row.
 *
 * Wanting is not being ready -- the flatten also asks [ParsedReplies.splitReady], and the two are
 * answered by different things: this one by the fold, the other by whether [warm] has run.
 */
private fun splitWanted(item: TranscriptItem.AssistantMsg, index: Int, lastIndex: Int) =
    item.settled || index != lastIndex

/**
 * The replies among [rows] that should draw as blocks but whose parses are not made yet.
 *
 * Normally empty: every page's rows are warmed before the fold lands. The one row that can be cold
 * is the reply that just finished streaming -- nothing warms live deltas. The session screen warms
 * what this returns off-thread and re-flattens, so the whole-to-blocks swap always composes against
 * ready parses.
 */
fun unwarmedReplies(
    rows: List<TranscriptRow>,
    replies: ParsedReplies,
    shownWhole: Set<Any> = emptySet(),
): List<TranscriptItem> = rows.mapIndexedNotNull { index, whole ->
    // The *capped* row's text, since that is what the flatten will draw and so what has to be
    // ready: a capped row draws its head, which is a different string from the message and so a
    // different cache entry.
    val row = capRow(whole, shownWhole).first
    val item = (row as? TranscriptRow.Single)?.item as? TranscriptItem.AssistantMsg
    item?.takeIf { splitWanted(it, index, rows.lastIndex) && !replies.splitReady(it.text) }
}

/**
 * Above this many characters, a user message is drawn in slices rather than as one bubble.
 *
 * Not zero, because a bubble's width wraps its content: slices have to fill the row to look like
 * one bubble, and forcing that on a short message would visibly widen it. A message past this
 * length has lines that wrap, so its bubble is at the full width already and the slices match it
 * exactly. Below it, one item of at most a few screens is nothing the list minds composing.
 */
const val USER_SPLIT_CHARS = 4000

/**
 * Roughly how much text one slice holds -- bounded, like a markdown block, is the whole point.
 *
 * About one viewport of wrapped text: a slice is composed whole in the frame it scrolls into, so
 * its size is a frame-budget decision, and one screenful keeps that to a few milliseconds on the
 * phone. Smaller buys nothing -- the seams are free -- but the units multiply.
 */
private const val USER_CHUNK_CHARS = 1000

/**
 * A long user message cut at line starts into slices of roughly [USER_CHUNK_CHARS].
 *
 * At newlines only, never mid-line: text layout runs per line, so slices that own whole lines stack
 * back into exactly the lines the single `Text` drew, and a cut inside one would reflow it. The
 * newline at each cut is dropped -- the boundary between two stacked slices *is* that line break. A
 * single line longer than a slice (minified JSON, a base64 blob) stays whole in its slice, so a
 * slice is bounded by the longest line rather than absolutely.
 */
fun userChunks(text: String): List<String> {
    val chunks = ArrayList<String>()
    var start = 0
    while (start < text.length) {
        if (text.length - start <= USER_CHUNK_CHARS) {
            chunks += text.substring(start)
            break
        }
        var cut = text.lastIndexOf('\n', start + USER_CHUNK_CHARS)
        if (cut <= start) cut = text.indexOf('\n', start + USER_CHUNK_CHARS)
        if (cut < 0) {
            chunks += text.substring(start)
            break
        }
        chunks += text.substring(start, cut)
        start = cut + 1
    }
    return chunks
}

/**
 * Says which two units share a key, before the list dies of it.
 *
 * A duplicate key is fatal -- `LazyColumn` throws, and the app goes down in the middle of somebody
 * reading a conversation -- and all the framework's message carries is the key. When that key is a
 * seq it names neither row, and there is no way back from it to how the two came to share one: it
 * took an afternoon and a fixture that could reproduce it. Two lines here answered it immediately,
 * naming both rows and the field they had in common ([TranscriptItem.PeerNote.arrived]).
 *
 * Always on, for the same reason [DebugStats] is: an instrument that is only in the build nobody is
 * holding when it breaks is not an instrument. It costs one map over the units that were just
 * built, beside a loop that already allocates one entry per unit.
 */
private fun reportDuplicateKeys(units: List<TranscriptUnit>) {
    val seen = HashMap<Any, TranscriptUnit>()
    units.forEach { unit ->
        val had = seen.put(unit.key, unit)
        if (had != null) {
            android.util.Log.w("ai-app", "duplicate unit key ${unit.key}: $had  AND  $unit")
        }
    }
}

/**
 * What is on screen right now, a unit at a time: what each one is and how tall it is.
 *
 * For the render report, and it is the line every "it is slow here" report has needed. The
 * framework's own per-frame cost grows with how many nodes are *alive* rather than how many are on
 * screen, so a screen holding one enormous item is slow in a way that no counter of ours
 * distinguishes from a screen holding twenty ordinary ones -- and "2 units visible" says one of
 * them is enormous without saying which. This says which.
 *
 * [first] is the index the list gave the first *unit*: the list also holds the waiting-messages
 * slot at index zero and the history spinner past the end, and both are named here rather than
 * silently reported as whichever unit is nearest.
 */
fun visibleUnits(units: List<TranscriptUnit>, visible: List<LazyListItemInfo>, first: Int): String =
    if (visible.isEmpty()) "  nothing on screen"
    else
        "  on screen: " +
            visible.joinToString(", ") { info ->
                "${units.getOrNull(info.index - first).kind} ${info.size}px"
            }

/** What a unit is, in a word, for [visibleUnits]. Null is one of the list's own non-unit items. */
private val TranscriptUnit?.kind: String
    get() =
        when (this) {
            null -> "the list's own"
            is TranscriptUnit.Block ->
                if (piece.item == Piece.WHOLE_BLOCK) "reply block" else "list item"
            is TranscriptUnit.PeerHead -> if (open) "peer heading (open)" else "peer heading"
            is TranscriptUnit.PeerBlock -> "peer block"
            is TranscriptUnit.UserChunk -> "user slice"
            is TranscriptUnit.Memory -> "memory note"
            is TranscriptUnit.ShowAll -> "show all"
            is TranscriptUnit.Whole ->
                when (val row = row) {
                    is TranscriptRow.Tools -> "tool group"
                    // The class name rather than a word per kind: this is a diagnostic, and a
                    // `when` here would be one more place that has to gain a case whenever the
                    // transcript does -- silently naming a new row after an old one until somebody
                    // noticed.
                    is TranscriptRow.Single -> row.item::class.simpleName.orEmpty()
                }
        }

/**
 * Where the unit named by a saved position sits in [units], or null if its row is not loaded.
 *
 * The row is found by [seq] and the unit within it by [ordinal], settling for the nearest older
 * unit when the exact one is gone -- a reply regrouped by a page boundary can split into a
 * different number of blocks than it had when the position was saved, and "a little above where
 * they stopped" loses less than the newest end does.
 */
fun unitIndexFor(units: List<TranscriptUnit>, seq: Long, ordinal: Int): Int? {
    var best: Int? = null
    var bestOrdinal = -1
    units.forEachIndexed { index, unit ->
        if (unit.seq == seq && unit.ordinal <= ordinal && unit.ordinal > bestOrdinal) {
            best = index
            bestOrdinal = unit.ordinal
        }
    }
    return best ?: units.indexOfFirst { it.seq == seq }.takeIf { it >= 0 }
}