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

package com.example.aiapp

import androidx.compose.foundation.background
import androidx.compose.foundation.clickable
import androidx.compose.foundation.gestures.awaitEachGesture
import androidx.compose.foundation.gestures.awaitFirstDown
import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.Row
import androidx.compose.foundation.layout.Spacer
import androidx.compose.foundation.layout.fillMaxWidth
import androidx.compose.foundation.layout.height
import androidx.compose.foundation.layout.padding
import androidx.compose.foundation.layout.width
import androidx.compose.foundation.shape.CornerBasedShape
import androidx.compose.foundation.shape.CornerSize
import androidx.compose.material3.Card
import androidx.compose.material3.CardDefaults
import androidx.compose.material3.CircularProgressIndicator
import androidx.compose.material3.MaterialTheme
import androidx.compose.material3.Text
import androidx.compose.runtime.Composable
import androidx.compose.runtime.remember
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.draw.clip
import androidx.compose.ui.graphics.Shape
import androidx.compose.ui.input.pointer.PointerEventPass
import androidx.compose.ui.input.pointer.pointerInput
import androidx.compose.ui.layout.LayoutCoordinates
import androidx.compose.ui.layout.onGloballyPositioned
import androidx.compose.ui.platform.LocalDensity
import androidx.compose.ui.semantics.contentDescription
import androidx.compose.ui.semantics.semantics
import androidx.compose.ui.text.font.FontFamily
import androidx.compose.ui.text.style.TextOverflow
import androidx.compose.ui.unit.Dp
import androidx.compose.ui.unit.dp

/**
 * A run of consecutive tool calls, or anything else, in the order they will be drawn.
 *
 * Grouping is decided here rather than when events are folded, because it is a display decision:
 * the transcript's own order is what paging and the event stream depend on, and one screen's idea
 * of "these belong together" must not reach back into it.
 */
sealed class TranscriptRow {
    /**
     * This row's identity in the list, which must survive everything that can happen to the row.
     *
     * The list is keyed by this so that inserting a new message at one end, or a page of history at
     * the other, moves the rows and not the reader. That makes it the load-bearing value on this
     * screen: when a key changes, the list loses its anchor and the transcript steps under whoever
     * is reading it.
     *
     * A tool row therefore keys on [TranscriptItem.ToolRun.runId] rather than on a sequence number,
     * and it is the *same* value whether the run is drawn as one card or as a group. A lone call
     * that gains a neighbour becomes a group without changing identity, which is the case a
     * seq-based key got wrong: the row the reader was looking at was replaced rather than updated.
     * Everything else keys on the seq of the event behind it, which never moves.
     */
    abstract val key: Any

    data class Single(val item: TranscriptItem) : TranscriptRow() {
        override val key: Any
            get() = (item as? TranscriptItem.ToolRun)?.runId ?: item.seq
    }

    /** Two or more calls with nothing between them; drawn as one collapsed card. */
    data class Tools(val calls: List<TranscriptItem.ToolRun>) : TranscriptRow() {
        /** The run's own name, which every call in it already carries. */
        val id: String
            get() = calls.first().runId

        override val key: Any
            get() = id
    }
}

/**
 * Runs of adjacent tool calls become one row; everything else passes through.
 *
 * A single call is left alone: "Called 1 tool" hides a card to say the same thing in more words,
 * and the run this exists for is the burst of five greps nobody wants to scroll past.
 */
fun groupToolRuns(items: List<TranscriptItem>): List<TranscriptRow> {
    val rows = mutableListOf<TranscriptRow>()
    var run = mutableListOf<TranscriptItem.ToolRun>()

    fun flush() {
        when (run.size) {
            0 -> {}
            1 -> rows += TranscriptRow.Single(run.first())
            else -> rows += TranscriptRow.Tools(run.toList())
        }
        run = mutableListOf()
    }

    items.forEach { item ->
        // Grouped by the run each call says it belongs to, not by adjacency worked out here.
        // Adjacency is the same answer most of the time and a worse one at the edges: a call
        // arriving next to an existing run, or a page of history arriving in front of one, both
        // change which call is *first*, and a group named after its first member is a different
        // group every time that happens.
        if (item is TranscriptItem.ToolRun && (run.isEmpty() || run.first().runId == item.runId)) {
            run += item
        } else {
            flush()
            if (item is TranscriptItem.ToolRun) run += item else rows += TranscriptRow.Single(item)
        }
    }
    flush()
    return rows
}

/** Where a click went down, and what it went down on. See [clickableAt]. */
private class TapPoint {
    var coords: LayoutCoordinates? = null
    var y = 0f
}

/**
 * Clickable, and tells the click where on the screen the finger went down.
 *
 * A row holds still the end nearest the tap when it changes height, so the toggle has to say where
 * it was touched. It cannot say which *end* it was: a group is one row with a control at each end
 * and a call in the middle, and only the row knows where its own ends are. So this reports a
 * position in root coordinates and leaves the meaning to whoever owns the row.
 *
 * Built on `clickable` rather than replacing it, because `clickable` is what draws the ripple and
 * what puts a click action in front of assistive technology. The position is read on the initial
 * pass and nothing is consumed, so the click still happens exactly as it would have.
 */
@Composable
fun Modifier.clickableAt(onClick: (Float) -> Unit): Modifier {
    val tap = remember { TapPoint() }
    return onGloballyPositioned { tap.coords = it }
        .pointerInput(Unit) {
            awaitEachGesture {
                val down =
                    awaitFirstDown(requireUnconsumed = false, pass = PointerEventPass.Initial)
                tap.y = tap.coords?.localToRoot(down.position)?.y ?: 0f
            }
        }
        .clickable { onClick(tap.y) }
}

/**
 * Several calls under one heading, closed until somebody asks.
 *
 * What says the calls belong together is the surface behind them, which is the one cue rather than
 * two half-cues -- rounded to the same corner every other card in the app has, so a group reads as
 * one object rather than as a square patch behind round things. The calls sit on it inset by
 * [GROUP_INSET], which is the container's own padding rather than an indent: they are the same rows
 * they would be on their own, and a rounded corner drawn hard against a rounded corner reads as a
 * notch.
 *
 * Inside, the calls are a connected stack. Facing corners are square and the outer ones are not, so
 * the run reads as one thing broken into its parts; [GROUP_GAP] keeps the parts legible without
 * separating them. See [connectedShape].
 *
 * It closes from either end. A long group's header scrolls off while its last call is still on
 * screen, and the reader who wants it shut is looking at the bottom, not hunting for the top. The
 * bar at the foot is the same height as the heading at the top, so the surface the calls sit on is
 * as thick below them as above.
 */
@Composable
fun ToolGroup(
    group: TranscriptRow.Tools,
    expanded: Boolean,
    /**
     * Told where it was pressed, because this row has a control at each end -- see [clickableAt].
     */
    onToggle: (Float) -> Unit,
    isToolExpanded: (String) -> Boolean,
    onToolToggle: (String, Float) -> Unit,
    onAnswer: (questionId: String, answers: List<String>) -> Unit,
    image: @Composable (String) -> Unit,
) {
    val heading = "Called ${group.calls.size} tools"
    if (!expanded) {
        Card(Modifier.fillMaxWidth().clickableAt(onToggle)) {
            Text(
                heading,
                style = MaterialTheme.typography.titleSmall,
                modifier = Modifier.padding(GROUP_INSET_LARGE),
            )
        }
        return
    }
    Column(
        Modifier.fillMaxWidth()
            .clip(MaterialTheme.shapes.medium)
            .background(MaterialTheme.colorScheme.surfaceContainerLow)
    ) {
        val barHeight = groupBarHeight()
        Row(
            Modifier.fillMaxWidth().height(barHeight).clickableAt(onToggle),
            verticalAlignment = Alignment.CenterVertically,
        ) {
            Text(
                heading,
                style = MaterialTheme.typography.titleSmall,
                modifier = Modifier.padding(horizontal = GROUP_INSET_LARGE),
            )
        }
        Column(
            Modifier.padding(horizontal = GROUP_INSET),
            verticalArrangement = Arrangement.spacedBy(GROUP_GAP),
        ) {
            group.calls.forEachIndexed { index, call ->
                ToolCard(
                    tool = call,
                    expanded = isToolExpanded(call.id),
                    onToggle = { at -> onToolToggle(call.id, at) },
                    onAnswer = onAnswer,
                    image = image,
                    shape = connectedShape(index, group.calls.size),
                )
            }
        }
        // Shutting it from here anchors the other end: the reader is at the bottom of a long
        // group, and what they are looking at is what follows it.
        CollapseBar(barHeight, onToggle)
    }
}

/**
 * The height of a group's heading, and so of the bar at its foot.
 *
 * Derived from the type the heading is set in rather than written down, because the two have to
 * match and a pair of numbers chosen to look equal stops being equal the moment either the style or
 * the density changes. Taking the line height also means the heading cannot be clipped by it.
 */
@Composable
private fun groupBarHeight(): Dp {
    val line = MaterialTheme.typography.titleSmall.lineHeight
    return with(LocalDensity.current) { line.toDp() } + GROUP_INSET_LARGE * 2
}

/**
 * The bottom half of a group's toggle: an arrow back up to its heading.
 *
 * Given the heading's height rather than padded to something that looks close, so the surface the
 * calls sit on is the same thickness at both ends. See [groupBarHeight].
 */
@Composable
private fun CollapseBar(height: Dp, onToggle: (Float) -> Unit) {
    val colour = MaterialTheme.colorScheme.onSurfaceVariant
    Row(
        Modifier.fillMaxWidth().height(height).clickableAt(onToggle).semantics {
            contentDescription = "Collapse these tool calls"
        },
        horizontalArrangement = Arrangement.Center,
        verticalAlignment = Alignment.CenterVertically,
    ) {
        Chevron(pointingUp = true, colour = colour)
    }
}

/**
 * The shape of one card in a stack of [count]: square where it faces a neighbour, rounded where it
 * does not.
 *
 * Written once and given an index rather than branched at each end, because a stack has three cases
 * that are one rule -- and the middle one is the case a hand-written first/last pair gets wrong
 * when a run turns out to have three calls in it.
 */
@Composable
private fun connectedShape(index: Int, count: Int): CornerBasedShape {
    val shape = MaterialTheme.shapes.medium
    val square = CornerSize(0.dp)
    return shape.copy(
        topStart = if (index == 0) shape.topStart else square,
        topEnd = if (index == 0) shape.topEnd else square,
        bottomStart = if (index == count - 1) shape.bottomStart else square,
        bottomEnd = if (index == count - 1) shape.bottomEnd else square,
    )
}

/** The padding inside a card, and so the height a bar of one line of text comes to. */
private val GROUP_INSET_LARGE = 12.dp

/** How far the stack of calls is held off the edge of the surface it sits on. */
private val GROUP_INSET = 4.dp

/** Enough to read the join as a join rather than as one tall card. */
private val GROUP_GAP = 2.dp

/**
 * One tool call.
 *
 * Closed, it is a single line: the tool's name and what the call is for. The command itself is not
 * on it, because a wrapped command turns one row into four and a run of them into a wall -- and the
 * name plus the intent is what somebody scanning the transcript is reading for.
 *
 * Open, it shows the command, whatever else the input carried, and the output. The timeout sits at
 * the top right: it is a limit on the call rather than part of what the call does, and it is worth
 * seeing beside the command it constrains rather than buried in the fields below it.
 *
 * A call waiting on permission is shown open whatever the reader last chose, since the command is
 * the thing being decided and a row saying only "Bash" cannot be decided on.
 */
@Composable
fun ToolCard(
    tool: TranscriptItem.ToolRun,
    expanded: Boolean,
    onToggle: (Float) -> Unit,
    onAnswer: (questionId: String, answers: List<String>) -> Unit,
    image: @Composable (String) -> Unit = {},
    /** Square where this card faces another in a group; see [connectedShape]. */
    shape: Shape = CardDefaults.shape,
) {
    val parsed = remember(tool.tool, tool.input) { parseToolInput(tool.tool, tool.input) }
    val deciding = tool.asks.any { it.answers.isEmpty() }
    val open = expanded || deciding
    Card(Modifier.fillMaxWidth().clickableAt(onToggle), shape = shape) {
        Column(Modifier.padding(GROUP_INSET_LARGE)) {
            Row(verticalAlignment = Alignment.CenterVertically) {
                Text(tool.tool, style = MaterialTheme.typography.titleSmall)
                if (open) {
                    Spacer(Modifier.weight(1f))
                    parsed.timeout?.let {
                        Text(
                            "timeout $it",
                            style = MaterialTheme.typography.labelSmall,
                            color = MaterialTheme.colorScheme.onSurfaceVariant,
                        )
                    }
                } else {
                    parsed.title?.let {
                        Text(
                            it,
                            style = MaterialTheme.typography.bodySmall,
                            color = MaterialTheme.colorScheme.onSurfaceVariant,
                            maxLines = 1,
                            overflow = TextOverflow.Ellipsis,
                            modifier = Modifier.weight(1f).padding(start = 8.dp),
                        )
                    } ?: Spacer(Modifier.weight(1f))
                }
                // A spinner says the machine is working. While this call is waiting on an
                // answer the machine is doing nothing at all -- the turn is stopped on the
                // person reading it -- so it says whose move it is instead, in the colour this
                // app uses everywhere for that.
                if (deciding) {
                    Spacer(Modifier.width(8.dp))
                    Text(
                        "your turn",
                        style = MaterialTheme.typography.labelLarge,
                        color = awaitingColor,
                    )
                } else if (!tool.done) {
                    Spacer(Modifier.width(8.dp))
                    CircularProgressIndicator(
                        modifier = Modifier.width(16.dp).height(16.dp),
                        strokeWidth = 2.dp,
                    )
                }
            }
            if (open) {
                parsed.description?.let {
                    Text(
                        it,
                        style = MaterialTheme.typography.bodySmall,
                        color = MaterialTheme.colorScheme.onSurfaceVariant,
                        modifier = Modifier.padding(top = 4.dp),
                    )
                }
                // Everything AskUserQuestion carries is the questions, and those are drawn
                // below as something answerable; dumping the same JSON above them would be the
                // decision stated twice, once unreadably.
                if (tool.tool != ASK_USER_QUESTION) {
                    ToolInputView(tool.tool, tool.input, Modifier.padding(top = 4.dp))
                }
                if (tool.output.isNotEmpty()) {
                    Spacer(Modifier.height(8.dp))
                    Text("Output", style = MaterialTheme.typography.labelSmall)
                    // What the tool printed, on the surface everything verbatim gets and in the
                    // face it was written for: this is column-aligned far more often than it is
                    // prose -- a directory listing, a diff, a table of numbers -- and a
                    // proportional font silently destroys the alignment that carried the meaning.
                    RawBlock(Modifier.padding(top = 2.dp)) {
                        Text(
                            tool.output,
                            style = MaterialTheme.typography.bodySmall,
                            fontFamily = FontFamily.Monospace,
                        )
                    }
                }
            }
            // Shown open or closed. A call that produced a picture is one
            // whose result *is* the picture, and a row that hides it says
            // less than the one line it replaced -- unlike a command, which
            // is what the closed line already summarises.
            tool.images.forEach { ref -> image(ref) }
            if (tool.asks.isNotEmpty()) {
                if (tool.tool == ASK_USER_QUESTION) {
                    AskUserQuestionBody(tool.asks, onAnswer)
                } else {
                    tool.asks.forEach { ask ->
                        PermissionAsk(ask) { answers -> onAnswer(ask.id, answers) }
                    }
                }
            }
        }
    }
}

/**
 * The permission ask on the call it is about.
 *
 * Only the question, not the prompt's second half: the backend sends the tool's input with it so
 * the ask can stand alone, and here it does not have to -- the card above is showing exactly that.
 */
@Composable
private fun PermissionAsk(ask: TranscriptItem.QuestionCard, onAnswer: (List<String>) -> Unit) {
    Spacer(Modifier.height(8.dp))
    Text(
        ask.prompt.substringBefore('\n'),
        style = MaterialTheme.typography.bodyMedium,
        color = awaitingColor,
    )
    if (ask.answers.isNotEmpty()) {
        Text(
            "Answered: ${ask.answers.joinToString(", ")}",
            style = MaterialTheme.typography.labelLarge,
            color = MaterialTheme.colorScheme.onSurfaceVariant,
        )
    } else {
        AnswerOptions(ask.options, onAnswer)
    }
}

/**
 * The tool whose input is a question rather than a command; see [AskUserQuestionBody].
 *
 * Also what [runIdFor] breaks a run of calls on, so the row a reader answered is never folded
 * inside a collapsed group.
 */
const val ASK_USER_QUESTION = "AskUserQuestion"