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; the
 * block splitting existed only inside the row, where the list could not see it.
 *
 * 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 markdown block of a settled reply. */
    data class Block(
        override val seq: Long,
        override val ordinal: Int,
        val text: String,
        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: as one item, every block of it is composed, measured, placed and kept alive
     * while any part of it is on screen. 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 -- which grows with how many nodes are *alive* -- 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. See [PeerHeadRow] and [PeerBlockRow].
     */
    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. See
         * [TranscriptItem.PeerNote].
         */
        override val key: Any
            get() = item.key

        override val ordinal: Int
            get() = 0
    }

    /** One markdown block of an opened peer message; [last] is the piece that closes the card. */
    data class PeerBlock(
        override val seq: Long,
        override val ordinal: Int,
        val text: String,
        val last: Boolean,
        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 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"
    }
}

/**
 * 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 blocks ([markdownBlocks], via the caches on [replies] so a
 * message is only ever split once), and so is an *opened* peer message -- [openNotes] is which ones
 * those are, which is why the flatten needs it. A shut one is a single heading and cannot be worth
 * splitting. The reply still arriving -- the newest row, until the status event that ends its turn
 * marks it [TranscriptItem.AssistantMsg.settled] -- stays whole: its text changes with every delta,
 * and splitting it here would parse the whole message per delta on whichever thread is composing.
 * [AssistantMessage]'s own streaming path already parses deltas off the main thread and gives the
 * live message a layer per block. 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.blocksOf] are lookups for any text [warm] has
 * seen, and a miss -- the one message that just finished streaming -- costs its split exactly once.
 */
fun transcriptUnits(
    rows: List<TranscriptRow>,
    replies: ParsedReplies,
    openNotes: Set<Long>,
): List<TranscriptUnit> {
    val units = ArrayList<TranscriptUnit>(rows.size)
    rows.forEachIndexed { index, row ->
        val rowGap = if (index == 0) 0.dp else TRANSCRIPT_SPACING
        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 blocks = replies.blocksOf(item.text)
                blocks.forEachIndexed { at, block ->
                    units +=
                        TranscriptUnit.PeerBlock(
                            row.startSeq,
                            at + 1,
                            block,
                            last = at == blocks.lastIndex,
                            gap = 0.dp,
                            note = item.key,
                        )
                }
            }
        } else if (
            item is TranscriptItem.AssistantMsg && (item.settled || index != rows.lastIndex)
        ) {
            var ordinal = 0
            fun gap() = if (ordinal == 0) rowGap else BLOCK_SPACING
            replies.partsOf(item.text).forEach { part ->
                when (part) {
                    is MessagePart.Prose ->
                        replies.blocksOf(part.text).forEach { block ->
                            units += TranscriptUnit.Block(row.startSeq, ordinal, block, gap())
                            ordinal++
                        }
                    is MessagePart.Remembered -> {
                        units += TranscriptUnit.Memory(row.startSeq, ordinal, part, gap())
                        ordinal++
                    }
                }
            }
        } else {
            units += TranscriptUnit.Whole(row, rowGap)
        }
    }
    units.reverse()
    reportDuplicateKeys(units)
    return units
}

/**
 * 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 -> "reply block"
            is TranscriptUnit.PeerHead -> if (open) "peer heading (open)" else "peer heading"
            is TranscriptUnit.PeerBlock -> "peer block"
            is TranscriptUnit.Memory -> "memory note"
            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 }
}