package com.example.aiapp

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"
    }

    /** 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). The reply still arriving -- the last row -- 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.
 *
 * 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): 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.AssistantMsg && 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()
    return units
}

/**
 * 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 }
}