package com.example.aiapp

import androidx.compose.runtime.Immutable
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext

/**
 * What the transcript renders: the event stream folded into displayable rows (see [foldEvent]).
 *
 * Events are the only data source, and there is deliberately no second shape for history to drift
 * from: a page fetched backwards, a live frame, and a line read out of this phone's own cache are
 * all the same events through the same parser. [TranscriptCache] stores the server's lines rather
 * than these rows for exactly that reason -- a row is a rendering, and its shape changes whenever
 * this file does.
 */
@Immutable
sealed class TranscriptItem {
    /**
     * The transcript sequence number this row started at, and its identity on screen.
     *
     * The list is drawn newest-first, so every new message is an insertion at index 0 and every
     * page of history is an insertion at the far end. Without an identity that survives both, the
     * list is addressed by position: whatever somebody had scrolled to keeps its index while the
     * content underneath it slides, which reads as the view scrolling on its own.
     *
     * A row built from several events keeps the seq of the first, so it holds still while the rest
     * of it arrives.
     */
    abstract val seq: Long

    /**
     * This item's identity on screen, which is its [seq] for everything that has one of its own.
     *
     * Here rather than in [TranscriptRow.Single] because the two items that need something else are
     * the two that know why. Asking each item what it is called is also what stops the next such
     * item being missed -- a `when` over concrete types would have to gain a case, silently.
     */
    open val key: Any
        get() = seq

    data class UserMsg(
        override val seq: Long,
        val text: String,
        /** Refs of what was attached, drawn inside the bubble. */
        val attachments: List<String> = emptyList(),
    ) : TranscriptItem()

    data class AssistantMsg(
        override val seq: Long,
        val text: String,
        /**
         * Whether this reply is finished: the session has stopped working since its last delta.
         *
         * What it buys is the split. [transcriptUnits] keeps the newest reply whole because a
         * streaming reply's text changes per delta and splitting a changing text is a parse per
         * delta -- but "newest" outlives the turn, so a session that ends on a long reply was
         * drawing it as one item indefinitely. Measured on a Pixel 9 Pro XL: one 34,996px reply on
         * screen put the frame's draw phase at 13.8ms, 79% of it framework bookkeeping.
         *
         * Folded from the status event that ended the turn rather than read off the screen's
         * status, because rows only change through the held-events gate: the split changes the
         * newest row's list identity, and doing that from a status flip while somebody is reading
         * inside that reply would step the list under them.
         */
        val settled: Boolean = false,
    ) : TranscriptItem()

    data class ToolRun(
        override val seq: Long,
        val id: String,
        /**
         * The run of adjacent calls this one belongs to, named once when the call is folded in and
         * never recomputed.
         *
         * Carried rather than derived because a run can gain members at *either* end, so no
         * function of its current members is stable. It is the first call's id at the moment the
         * run started, which is a name rather than a description: [joinPages] hands it to older
         * calls that turn out to belong to the same run.
         */
        val runId: String,
        val tool: String,
        val input: String,
        val output: String,
        val done: Boolean,
        /**
         * The questions this call is waiting on, in the order they were asked.
         *
         * On the call's own row rather than beside it: an ask used to arrive as a second card
         * repeating the input verbatim, so the reader saw the same command twice. The backend says
         * which call a question is about, so this is a fact rather than a match on the input.
         *
         * A list because AskUserQuestion asks up to four at once, and a permission is the case of
         * exactly one rather than a different shape.
         */
        val asks: List<QuestionCard> = emptyList(),
        /**
         * Images this call's result carried, drawn under it. Beside it they had to be paired by
         * position, and position is what a page boundary breaks.
         */
        val images: List<String> = emptyList(),
    ) : TranscriptItem() {
        /** A run, not a seq: see [TranscriptRow.key] for what that identity has to survive. */
        override val key: Any
            get() = runId
    }

    data class QuestionCard(
        override val seq: Long,
        val id: String,
        val prompt: String,
        /** A few words naming what this is about, when the asker offered one. */
        val header: String?,
        val options: List<QuestionOption>,
        /** Whether several options may be chosen at once. */
        val multiSelect: Boolean,
        /** What was chosen, once something was; empty until then. */
        val answers: List<String>,
    ) : TranscriptItem()

    data class ErrorMsg(override val seq: Long, val message: String) : TranscriptItem()

    /** An image by server-side ref, fetched from the session's files route. */
    data class ImageItem(override val seq: Long, val ref: String) : TranscriptItem()

    /**
     * A message another agent sent this session. Its own row rather than a [UserMsg]: see
     * [PeerMessageRow] for why the voice matters.
     */
    data class PeerNote(
        override val seq: Long,
        val from: String,
        val text: String,
        /**
         * The seq of the event this note came in on, which is what makes it itself.
         *
         * [seq] is where the note *sorts*, and [placePeerNote] sets it to the seq the turn began
         * at. Two messages that arrive during one turn therefore share a seq -- and sharing an
         * identity as well killed the app, because the list refuses two items with one key.
         */
        val arrived: Long = seq,
    ) : TranscriptItem() {
        override val key: Any
            get() = arrived
    }

    /**
     * A task the session started in the background reporting back: a subagent that finished, or a
     * backgrounded command.
     *
     * Its own row rather than an update to the Task call's, which is wherever the call was made --
     * above everything the session has said since, where a reader at the bottom would never see it
     * change. Here it is where it arrived, in front of the turn it caused.
     */
    data class TaskNote(
        override val seq: Long,
        /** The tool call it belongs to; a subagent is named by that id. */
        val about: String,
        /**
         * The subagent's title, or null for a backgrounded command -- see [SessionEvent.TaskNote].
         */
        val title: String?,
        val status: String,
        val summary: String?,
    ) : TranscriptItem()

    /**
     * A command the session ran on itself -- `/compact`, `/rename`. Kept in the transcript rather
     * than only shown while it waits, because it explains what follows: a conversation that
     * suddenly has half the context, or a session with a new name.
     */
    data class CommandRow(override val seq: Long, val text: String) : TranscriptItem()

    /** Placeholder row for events this build can't render (newer kinds). */
    data class Note(override val seq: Long, val text: String) : TranscriptItem()

    /**
     * A clear that happened: everything above it left the session's context and stayed on screen.
     * Carries only its position, because that is all it means.
     */
    data class ClearedNote(override val seq: Long) : TranscriptItem()

    /**
     * A compaction that happened, and what it recovered.
     *
     * In the transcript rather than only in the status line, because the status is gone the moment
     * it finishes and this is the part worth keeping: the explanation for a gap in the
     * conversation.
     *
     * The wire also says what triggered it, and this deliberately does not carry that -- the row
     * says the two sizes and nothing else, so keeping the trigger would be a field nothing can
     * read.
     */
    data class CompactedNote(
        override val seq: Long,
        val preTokens: Long?,
        val postTokens: Long?,
    ) : TranscriptItem()

    /**
     * The account ran out of quota, so the turn stopped here.
     *
     * A divider rather than an error: nothing failed, and what a reader scrolling back needs from
     * it is the same thing a clear or a compaction gives them -- why the conversation stops at this
     * line.
     *
     * [resetsAt] is epoch seconds and null where the session was told nothing, which is a state the
     * row has words for rather than a time it invents.
     */
    data class LimitNote(override val seq: Long, val resetsAt: Double?) : TranscriptItem()
}

/**
 * The run a call joins: the one it lands next to, or a new one named after itself.
 *
 * Only ever consulted when the call is first folded in. That is what makes the name stable -- a run
 * keeps whatever it was called when it started, however many calls arrive at either end afterwards.
 *
 * A question to the reader is in a run of its own, which is what puts it on the transcript as a row
 * rather than inside a collapsed "Called 6 tools" card. Two things follow: it is always visible,
 * since a run of one is drawn as itself; and the calls around it fall into a group before it and a
 * group after it, so where the reader was asked something is legible in the shape of the transcript
 * without opening anything.
 */
private fun runIdFor(items: List<TranscriptItem>, id: String, tool: String): String {
    val previous = items.lastOrNull() as? TranscriptItem.ToolRun ?: return id
    if (tool == ASK_USER_QUESTION || previous.tool == ASK_USER_QUESTION) return id
    return previous.runId
}

/**
 * Puts a page of older items in front of the ones already loaded, healing whatever the page
 * boundary cut in two.
 *
 * Two things straddle a boundary: a tool call separated from its result, and a message separated
 * from the rest of itself. Both were one thing before the transcript was cut into pages.
 *
 * A boundary lands wherever it lands, and roughly half the time that is between a call and its
 * result. The newer page then holds a `ToolEnd` whose start it never saw, which [foldEvent] draws
 * as a row of its own -- correctly, because a call that renders as nothing is indistinguishable
 * from one that never happened. When the older page arrives it brings the real `ToolStart`, and
 * concatenating the two lists left *both*: the same call twice.
 *
 * Merged by the call's own id rather than by position, because position is exactly what a page
 * boundary destroys. The older row wins on what a start knows and the newer on what an end knows,
 * which is the only way round that loses nothing.
 *
 * The third thing is the *run*, and it is the one that used to be missed. Every page ends up here,
 * but [adoptRun] only ran on the path where a split call had been found -- so the boundary that
 * falls cleanly between two finished calls, which is most of them, left the older page's calls
 * under the run name they were folded with. On screen: one run of tool calls drawn as two groups,
 * with the seam wherever the reader happened to have paged.
 */
fun joinPages(earlier: List<TranscriptItem>, later: List<TranscriptItem>): List<TranscriptItem> {
    val (older, newer) = healSplitMessage(earlier, later)
    val startedEarlier =
        older.filterIsInstance<TranscriptItem.ToolRun>().mapTo(mutableSetOf()) { it.id }
    val endedLater =
        newer
            .filterIsInstance<TranscriptItem.ToolRun>()
            .associateBy { it.id }
            .filterKeys { it in startedEarlier }
    val healed = older.map { row ->
        val half = (row as? TranscriptItem.ToolRun)?.let { endedLater[it.id] }
        if (row is TranscriptItem.ToolRun && half != null) {
            row.copy(
                output = half.output,
                done = half.done,
                // Kept from both halves: a question or an image can be attached to either,
                // depending on which side of the boundary its event fell.
                asks = row.asks + half.asks,
                images = row.images + half.images,
            )
        } else {
            row
        }
    }
    val kept = newer.filterNot { it is TranscriptItem.ToolRun && it.id in endedLater }
    return adoptRun(healed, kept) + kept
}

/**
 * Rejoins a message the page boundary cut, and hands back the two pages to concatenate.
 *
 * [foldEvent] never leaves an *unfinished* assistant message with another behind it inside one
 * page, so an unsettled one at a join is always the far half of the reply the boundary cut, and
 * leaving the two apart drew a single answer as two with a paragraph break through the middle of a
 * sentence. Two settled replies meeting there are two turns and stay two.
 *
 * The newer half keeps its identity, for the reason [adoptRun] gives. It grows by what the older
 * half brings, which is safe here and nowhere else -- the join is at the oldest end of what is
 * loaded, so the growth extends off the top of the screen.
 */
private fun healSplitMessage(
    earlier: List<TranscriptItem>,
    later: List<TranscriptItem>,
): Pair<List<TranscriptItem>, List<TranscriptItem>> {
    val head = earlier.lastOrNull()
    val tail = later.firstOrNull()
    if (head !is TranscriptItem.AssistantMsg || tail !is TranscriptItem.AssistantMsg) {
        return earlier to later
    }
    // A settled reply is a whole turn, so the two are two answers that happen to meet at the
    // boundary rather than one cut in half -- the same distinction the fold makes, and joining them
    // here would put back exactly the run-together paragraph it stops.
    if (head.settled) return earlier to later
    return earlier.dropLast(1) to (listOf(tail.copy(text = head.text + tail.text)) + later.drop(1))
}

/**
 * Hands the older calls at the join the name of the run they are joining.
 *
 * The two pages were folded separately, so a run split by the boundary came back as two runs with
 * two names. Naming the joined run after the *older* half would be the obvious way round and is
 * wrong: the newer half is the part already on screen, and renaming it is renaming the row the
 * reader is looking at, which is how a list loses its anchor.
 */
private fun adoptRun(
    earlier: List<TranscriptItem>,
    later: List<TranscriptItem>,
): List<TranscriptItem> {
    val first = later.firstOrNull() as? TranscriptItem.ToolRun ?: return earlier
    // A question is in a run of its own on both sides of the join, the same as it would be had the
    // two pages been folded as one. Without this the heal would merge a group straight through the
    // row the reader was asked something on.
    if (first.tool == ASK_USER_QUESTION) return earlier
    val joining = first.runId
    val tail = earlier.takeLastWhile {
        it is TranscriptItem.ToolRun && it.tool != ASK_USER_QUESTION
    }
    if (tail.isEmpty()) return earlier
    return earlier.dropLast(tail.size) +
        tail.map { (it as TranscriptItem.ToolRun).copy(runId = joining) }
}

/**
 * A peer message goes above the turn it started, not where it happened to arrive.
 *
 * The live Claude Code path cannot record it in place: the CLI says nothing about a peer message
 * until the turn's `result`, so the event lands below the whole reply it caused. The server stamps
 * it with where that turn began and the note takes that seq.
 *
 * Taking the turn's opening seq as its own is also what keeps the list sorted, which anchors and
 * paging both depend on. It is only a *position*, though, and the note keeps its own arrival seq as
 * its identity ([TranscriptItem.PeerNote.arrived]). The argument for sharing was that the turn's
 * seq belongs to a status change and a status draws no row -- true, and it answered the wrong
 * question: what two notes stamped with the same turn collide with is each other.
 *
 * Without a stamp -- a message replayed out of a session file -- it stays where it arrived.
 */
private fun placePeerNote(
    items: List<TranscriptItem>,
    seq: Long,
    event: SessionEvent.PeerMessage,
): List<TranscriptItem> {
    val at = event.turnStart ?: return items + TranscriptItem.PeerNote(seq, event.from, event.text)
    val note = TranscriptItem.PeerNote(at, event.from, event.text, arrived = seq)
    val index = items.indexOfFirst { it.seq > at }
    if (index < 0) return items + note
    val behind = (items.getOrNull(index - 1) as? TranscriptItem.ToolRun)?.runId
    return items.subList(0, index) + note + splitRun(items.subList(index, items.size), behind)
}

/**
 * The calls the note now sits in front of, renamed if they were sharing a run with the calls behind
 * it.
 *
 * A run is named from what a call landed next to, and nothing there knows about turns -- so a turn
 * opening with a tool call, straight after one that ended with one, folds them into a single run.
 * Left alone, [groupToolRuns] would flush at the note and hand both halves the same name: two rows
 * with one key, which a keyed list cannot draw at all.
 *
 * The later half is the one renamed, which is the opposite of a page join ([adoptRun]) and right
 * for the opposite reason: there the two halves were always one run, here they were never one
 * turn's work.
 */
private fun splitRun(tail: List<TranscriptItem>, behind: String?): List<TranscriptItem> {
    val first = tail.firstOrNull() as? TranscriptItem.ToolRun ?: return tail
    if (behind == null || first.runId != behind) return tail
    val run = tail.takeWhile { it is TranscriptItem.ToolRun && it.runId == behind }
    return run.map { (it as TranscriptItem.ToolRun).copy(runId = first.id) } + tail.drop(run.size)
}

fun foldEvent(items: List<TranscriptItem>, entry: SeqEvent): List<TranscriptItem> =
    when (val event = entry.event) {
        is SessionEvent.UserMessage ->
            items + TranscriptItem.UserMsg(entry.seq, event.text, event.attachments)
        is SessionEvent.AssistantText -> {
            // Deltas accumulate into the message they're streaming, which keeps the seq of the
            // first of them: a row whose identity changed with every delta would be a new row on
            // every frame, and the list would jump for the whole of a streamed answer.
            val last = items.lastOrNull()
            // Only into a reply that is still arriving. A settled one is a turn that ended, and
            // text after it belongs to the next turn -- a separate message, drawn as its own row.
            // Growing it instead ran two answers together with not even a space between them,
            // which is what happens whenever a turn starts with nothing recorded in front of it:
            // a subagent reporting back, or a peer message the CLI only owns up to at the end.
            if (last is TranscriptItem.AssistantMsg && !last.settled) {
                items.dropLast(1) + last.copy(text = last.text + event.delta)
            } else {
                items + TranscriptItem.AssistantMsg(entry.seq, event.delta)
            }
        }
        is SessionEvent.ToolStart ->
            items +
                TranscriptItem.ToolRun(
                    entry.seq,
                    event.id,
                    runIdFor(items, event.id, event.tool),
                    event.tool,
                    event.input,
                    "",
                    done = false,
                )
        is SessionEvent.ToolUpdate -> updateTool(items, event.id) { it.copy(output = event.output) }
        is SessionEvent.ToolEnd ->
            // Created when its start is not here, rather than dropped. A fold that only ever
            // *updates* loses the whole call when the start fell outside the loaded window, and a
            // tool call that renders as nothing is indistinguishable from one that never happened.
            // Loading the page before this one replaces the row with the real thing.
            if (items.any { it is TranscriptItem.ToolRun && it.id == event.id }) {
                updateTool(items, event.id) { it.copy(output = event.output, done = true) }
            } else {
                items +
                    TranscriptItem.ToolRun(
                        entry.seq,
                        event.id,
                        // The name is not known from an end alone, so a call that was an ask cannot
                        // be recognised as one here; the page before this replaces the row.
                        runIdFor(items, event.id, "tool"),
                        "tool",
                        "",
                        event.output,
                        done = true,
                    )
            }
        is SessionEvent.Question -> {
            val card =
                TranscriptItem.QuestionCard(
                    entry.seq,
                    event.id,
                    event.prompt,
                    event.header,
                    event.options,
                    event.multiSelect,
                    emptyList(),
                )
            // A question with no tool behind it -- AskUserQuestion, or an ask whose call fell
            // outside the loaded window -- is a card of its own.
            if (
                event.about != null &&
                    items.any { it is TranscriptItem.ToolRun && it.id == event.about }
            ) {
                updateTool(items, event.about) { it.copy(asks = it.asks + card) }
            } else {
                items + card
            }
        }
        is SessionEvent.Answered ->
            // Resolved wherever it is drawn: a card of its own, or a tool row's ask. Missing the
            // second left an Allow/Deny pair live on a question already answered from another
            // device.
            items.map {
                when {
                    it is TranscriptItem.QuestionCard && it.id == event.id ->
                        it.copy(answers = event.answers)
                    it is TranscriptItem.ToolRun && it.asks.any { ask -> ask.id == event.id } ->
                        it.copy(
                            asks =
                                it.asks.map { ask ->
                                    if (ask.id == event.id) ask.copy(answers = event.answers)
                                    else ask
                                }
                        )
                    else -> it
                }
            }
        is SessionEvent.PeerMessage -> placePeerNote(items, entry.seq, event)
        is SessionEvent.TaskNote ->
            items +
                TranscriptItem.TaskNote(
                    entry.seq,
                    event.about,
                    event.title,
                    event.status,
                    event.summary,
                )
        is SessionEvent.CommandSent -> items + TranscriptItem.CommandRow(entry.seq, event.text)
        // Screen-level state, not transcript rows -- see SessionScreen.
        is SessionEvent.CommandQueued -> items
        // No row of its own: a message that is still waiting is drawn as a pending bubble below the
        // transcript, and becomes an ordinary one where the session read it.
        is SessionEvent.MessageQueued -> items
        // The bubble goes away and nothing takes its place: the message was never read, so there is
        // nothing it belongs above.
        is SessionEvent.MessageDropped -> items
        is SessionEvent.Settings -> items
        is SessionEvent.Status -> settleReply(items, event.state)
        is SessionEvent.Error -> items + TranscriptItem.ErrorMsg(entry.seq, event.message)
        is SessionEvent.Image ->
            // Under the call that produced it when there is one, and a row of its own when there is
            // not -- a person's own attachment belongs to no call, and neither does one whose call
            // fell outside the loaded window.
            if (
                event.about != null &&
                    items.any { it is TranscriptItem.ToolRun && it.id == event.about }
            ) {
                updateTool(items, event.about) { it.copy(images = it.images + event.ref) }
            } else {
                items + TranscriptItem.ImageItem(entry.seq, event.ref)
            }
        is SessionEvent.LimitReached -> items + TranscriptItem.LimitNote(entry.seq, event.resetsAt)
        is SessionEvent.Cleared -> items + TranscriptItem.ClearedNote(entry.seq)
        is SessionEvent.Compacted ->
            items + TranscriptItem.CompactedNote(entry.seq, event.preTokens, event.postTokens)
        is SessionEvent.Unknown -> items + TranscriptItem.Note(entry.seq, "[${event.type}]")
        // Screen-level state, not transcript rows -- see SessionScreen.
        is SessionEvent.UsageDelta -> items
    }

/**
 * A status saying the session stopped working is the moment its newest reply is finished.
 *
 * See [TranscriptItem.AssistantMsg.settled]. Status changes are transcript events with seqs of
 * their own, so a replayed session settles its replies the same way a live one does.
 */
private fun settleReply(items: List<TranscriptItem>, state: String): List<TranscriptItem> {
    if (sessionWorking(state)) return items
    val last = items.lastOrNull() as? TranscriptItem.AssistantMsg ?: return items
    if (last.settled) return items
    return items.dropLast(1) + last.copy(settled = true)
}

private fun updateTool(
    items: List<TranscriptItem>,
    id: String,
    change: (TranscriptItem.ToolRun) -> TranscriptItem.ToolRun,
): List<TranscriptItem> = items.map {
    if (it is TranscriptItem.ToolRun && it.id == id) change(it) else it
}

/**
 * Where markdown is parsed ahead of being drawn: two threads, never all of them.
 *
 * The default dispatcher sizes itself to the machine, which is right for work somebody is waiting
 * on and wrong for work nobody is. A page of history is hundreds of parses arriving at once, and
 * taking every core for them leaves the thread that draws the frame queueing behind one -- measured
 * on a Pixel 9 Pro XL as 21ms of `waited` at the 90th percentile.
 */
@OptIn(kotlinx.coroutines.ExperimentalCoroutinesApi::class)
private val parsingThreads = Dispatchers.Default.limitedParallelism(2)

/**
 * Parses the markdown among [rows], off whatever thread is drawing.
 *
 * Called where a page of transcript is folded rather than where a row is composed, which is the
 * whole point: the work happens seconds before the reader reaches the rows it was done for.
 *
 * What is warmed mirrors what the rows draw -- each prose part of a reply, a memory note, a peer
 * message -- because a string warmed under a key no row ever looks up is a miss that nothing
 * reports; see [transcriptUnits], which is the flatten this has to agree with. It reads the same
 * [ParsedReplies.partsOf] cache the flatten does, so a message is scanned once however many pages
 * hand it back through here.
 *
 * Every kind of row that draws markdown belongs in the `when` below. That is the rule the peer
 * message was missing: this used to filter for assistant replies alone, so the one row type nobody
 * had thought about paid its whole parse in the frame it appeared in.
 */
suspend fun warm(replies: ParsedReplies, rows: List<TranscriptItem>) {
    withContext(parsingThreads) {
        val texts = rows.flatMap { row ->
            when (row) {
                is TranscriptItem.AssistantMsg -> replies.partsOf(row.text).map { it.text }
                // A message from another agent is markdown too, and it is the longest thing in a
                // transcript often enough that leaving it out was the whole of why one cost a fifth
                // of a second to open.
                is TranscriptItem.PeerNote -> listOf(row.text)
                else -> emptyList()
            }
        }
        if (texts.isNotEmpty()) replies.warm(texts)
        // After the parses exist, not before: [ParsedReplies.splitReady] is the flatten's licence
        // to draw these as blocks on the composing thread.
        rows.forEach { if (it is TranscriptItem.AssistantMsg) replies.markSplitReady(it.text) }
    }
}