package com.example.aiapp

import androidx.compose.foundation.background
import androidx.compose.foundation.clickable
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.Immutable
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.setValue
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.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

/**
 * One row as the transcript draws it: a run of consecutive tool calls, or anything else.
 *
 * 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.
 *
 * Immutable, and said so, because Compose cannot tell: a row is rebuilt from the transcript rather
 * than edited, and two rows describing the same events are equal. Compose infers stability from a
 * class's fields, and a `List` field -- which several of these carry -- makes it assume the worst,
 * so a page of history landing recomposed all 148 loaded rows including the markdown inside them,
 * measured as 701 compositions for 148 rows in one scroll.
 *
 * The promise this makes is real and has to stay true: nothing here is mutated after it is built.
 */
@Immutable
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. 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. Which value
     * that is belongs to the item ([TranscriptItem.key]), not to a `when` here.
     */
    abstract val key: Any

    /**
     * Where this row starts in the transcript: the sequence number of the oldest event behind it.
     *
     * Separate from [key], and deliberately so. [key] is the list's identity and is a display
     * decision; a seq is the server's own numbering, assigned once and meaning the same thing to
     * every device. So anything that has to point at a place in the conversation and still find it
     * later -- a saved scroll position -- points with this.
     */
    abstract val startSeq: Long

    data class Single(val item: TranscriptItem) : TranscriptRow() {
        override val key: Any
            get() = item.key

        override val startSeq: Long
            get() = 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

        override val startSeq: Long
            get() = calls.first().seq
    }
}

/**
 * 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> =
    DebugStats.timed("grouped tool runs") { groupRuns(items) }

private fun groupRuns(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*.
        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
}

/**
 * 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.
 *
 * 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; 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. The bar at the foot is the
 * same height as the heading at the top.
 */
@Composable
fun ToolGroup(
    group: TranscriptRow.Tools,
    expanded: Boolean,
    /**
     * Where it was pressed is the row's business rather than the control's -- a group has a control
     * at each end, and only the row knows where its own ends are.
     */
    onToggle: () -> Unit,
    isToolExpanded: (String) -> Boolean,
    onToolToggle: (String) -> Unit,
    onAnswer: (List<QuestionAnswer>, onSettled: () -> Unit) -> Unit,
    image: @Composable (String) -> Unit,
) {
    val heading = "Called ${group.calls.size} tools"
    if (!expanded) {
        Card(Modifier.fillMaxWidth().clickable(onClick = 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).clickable(onClick = 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 = { onToolToggle(call.id) },
                    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 the density
 * changes.
 */
@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.
 */
@Composable
private fun CollapseBar(height: Dp, onToggle: () -> Unit) {
    val colour = MaterialTheme.colorScheme.onSurfaceVariant
    Row(
        Modifier.fillMaxWidth().height(height).clickable(onClick = onToggle).semantics {
            contentDescription = "Collapse these tool calls"
        },
        horizontalArrangement = Arrangement.Center,
        verticalAlignment = Alignment.CenterVertically,
    ) {
        Chevron(Pointing.Up, 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 what a hand-written first/last pair gets wrong.
 */
@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.
 *
 * 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.
 *
 * 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: () -> Unit,
    onAnswer: (List<QuestionAnswer>, onSettled: () -> Unit) -> 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().clickable(onClick = 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.
                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, and a proportional font silently destroys the alignment that carried
                    // the meaning.
                    //
                    // Its terminal styling applied and the rest of the escapes taken out: colour is
                    // often the whole of what a diff or a test run is saying. Remembered against
                    // the text, so a card that is open through a scroll parses once.
                    val palette = remember { ansiPalette() }
                    val styled = remember(tool.output, palette) { ansiStyled(tool.output, palette) }
                    RawBlock(Modifier.padding(top = 2.dp)) {
                        Text(
                            styled,
                            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.
            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, onAnswer) }
                }
            }
        }
    }
}

/**
 * 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<QuestionAnswer>, onSettled: () -> Unit) -> Unit,
) {
    // What was pressed, before the answer has been round-tripped. Two bare words with no submit
    // step -- unlike a question card, where the answer is worth reviewing -- so the press has to be
    // its own acknowledgement or the row sits unchanged for a round trip. Cleared when the request
    // settles: by then either the answer is in `ask.answers`, or it failed and the buttons come
    // back.
    var pressed by remember(ask.id) { mutableStateOf<String?>(null) }
    Spacer(Modifier.height(8.dp))
    Text(
        ask.prompt.substringBefore('\n'),
        style = MaterialTheme.typography.bodyMedium,
        color = awaitingColor,
    )
    // Answered or not, the options stay and the one that was taken is marked -- see
    // [AskedQuestion], which is the same rule on the question card. A permission is where it
    // matters most: "Answered: Deny" alone does not say that Allow was the alternative.
    val settled = ask.answers.isNotEmpty()
    AnswerOptions(
        ask.options,
        if (settled) ask.answers else listOfNotNull(pressed),
        onPick =
            if (settled || pressed != null) null
            else
                { label ->
                    pressed = label
                    onAnswer(listOf(QuestionAnswer(ask.id, listOf(label)))) { pressed = null }
                },
    )
}

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