ai-app/app/androidApp/src/test/kotlin/com/example/aiapp/ToolRowsTest.kt

package com.example.aiapp

import kotlin.test.Test
import kotlin.test.assertEquals
import kotlin.test.assertTrue

/**
 * How a run of tool calls is cut into rows: the call still running, the last call in the
 * transcript, and one held out because the reader has it open are drawn on their own, and every
 * piece the cut leaves behind still has a key of its own -- two rows sharing one key take the app
 * down, and a key that moves takes the reader's place with it.
 */
class ToolRowsTest {
    private var seq = 0L

    private fun call(id: String, runId: String = id, done: Boolean = true) =
        TranscriptItem.ToolRun(
            seq = ++seq,
            id = id,
            runId = runId,
            tool = "Bash",
            input = "{}",
            output = if (done) "ok" else "",
            done = done,
        )

    /** Something that is not a tool call, to put behind the run so its last call folds in. */
    private fun reply() = TranscriptItem.AssistantMsg(seq = ++seq, text = "done")

    private fun shape(rows: List<TranscriptRow>) = rows.map { row ->
        when (row) {
            is TranscriptRow.Tools -> row.calls.map { it.id }
            is TranscriptRow.Single -> listOf((row.item as? TranscriptItem.ToolRun)?.id ?: "reply")
        }
    }

    private fun assertKeysDistinct(rows: List<TranscriptRow>) =
        assertEquals(rows.size, rows.map { it.key }.toSet().size, "$rows")

    @Test
    fun the_call_still_running_is_a_row_of_its_own() {
        val rows =
            groupToolRuns(
                listOf(
                    call("a"),
                    call("b", runId = "a"),
                    call("c", runId = "a", done = false),
                    call("d", runId = "a"),
                    reply(),
                )
            )
        assertEquals(
            listOf(listOf("a", "b"), listOf("c"), listOf("d"), listOf("reply")),
            shape(rows),
        )
        assertKeysDistinct(rows)
    }

    @Test
    fun a_call_running_in_the_middle_of_its_run_splits_the_group_in_two() {
        val rows =
            groupToolRuns(
                listOf(
                    call("a"),
                    call("b", runId = "a", done = false),
                    call("c", runId = "a"),
                    call("d", runId = "a"),
                    reply(),
                )
            )
        assertEquals(
            listOf(listOf("a"), listOf("b"), listOf("c", "d"), listOf("reply")),
            shape(rows),
        )
        assertKeysDistinct(rows)
    }

    /**
     * A call held out is one the reader opened while it stood on its own; being overtaken while
     * they read it does not fold it away, and closing it hands it back to its run.
     */
    @Test
    fun a_held_out_call_stays_out_of_its_group() {
        val calls =
            listOf(
                call("a"),
                call("b", runId = "a"),
                call("c", runId = "a"),
                call("d", runId = "a"),
                reply(),
            )

        val whileHeld = groupToolRuns(calls, heldOut = setOf("d"))
        val afterItCloses = groupToolRuns(calls)

        assertEquals(listOf(listOf("a", "b", "c"), listOf("d"), listOf("reply")), shape(whileHeld))
        assertTrue(whileHeld[1] is TranscriptRow.Single, "$whileHeld")
        assertKeysDistinct(whileHeld)
        assertEquals(listOf(listOf("a", "b", "c", "d"), listOf("reply")), shape(afterItCloses))
        assertTrue(afterItCloses.first() is TranscriptRow.Tools, "$afterItCloses")
    }

    /**
     * The one case where the run's name is a call that is not in the run's first row: a page of
     * history joined onto a run whose own first call is still going ([joinPages] renames the older
     * calls to the newer run's name). Both rows would key on that name.
     */
    @Test
    fun the_run_keeps_its_name_even_when_the_call_it_is_named_after_is_the_one_running() {
        val rows = groupToolRuns(listOf(call("a", runId = "b"), call("b", done = false)))
        assertEquals(listOf(listOf("a"), listOf("b")), shape(rows))
        assertKeysDistinct(rows)
        assertEquals("b", rows.first().key)
    }

    @Test
    fun a_run_that_reappears_after_another_row_keeps_distinct_keys() {
        val rows =
            groupToolRuns(
                listOf(
                    call("older", runId = "exec-1"),
                    call("older-2", runId = "exec-1"),
                    reply(),
                    call("exec-1", runId = "exec-1"),
                    call("newer", runId = "exec-1"),
                    reply(),
                )
            )

        assertTrue(rows[0] is TranscriptRow.Tools, "$rows")
        assertTrue(rows[2] is TranscriptRow.Tools, "$rows")
        assertKeysDistinct(rows)
        assertEquals("exec-1", rows[0].key)
        assertEquals("exec-1/exec-1", rows[2].key)
    }

    /**
     * Finishing is not what folds a call back in -- being overtaken is. A session that has run its
     * last command and is writing its reply leaves that command standing until the reply starts.
     */
    @Test
    fun the_last_call_stays_out_when_it_finishes_and_folds_in_when_something_follows() {
        val a = call("a")
        val running = call("b", runId = "a", done = false)
        val finished = running.copy(done = true)
        val whileRunning = groupToolRuns(listOf(a, running))
        val afterItEnds = groupToolRuns(listOf(a, finished))
        val afterTheReply = groupToolRuns(listOf(a, finished, reply()))
        assertEquals(listOf(listOf("a"), listOf("b")), shape(whileRunning))
        assertEquals(listOf(listOf("a"), listOf("b")), shape(afterItEnds))
        assertEquals(listOf(listOf("a", "b"), listOf("reply")), shape(afterTheReply))
        // The run keeps the key it was drawn under throughout, so the list rebuilds a row rather
        // than losing its anchor.
        assertEquals(whileRunning.first().key, afterItEnds.first().key)
        assertEquals(whileRunning.first().key, afterTheReply.first().key)
    }

    @Test
    fun a_run_of_finished_calls_is_one_group_once_something_follows_it() {
        val rows =
            groupToolRuns(
                listOf(call("a"), call("b", runId = "a"), call("c", runId = "a"), reply())
            )
        assertEquals(listOf(listOf("a", "b", "c"), listOf("reply")), shape(rows))
        assertTrue(rows.first() is TranscriptRow.Tools, "$rows")
    }
}