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

package com.example.aiapp

import android.graphics.BitmapFactory
import androidx.activity.compose.rememberLauncherForActivityResult
import androidx.activity.result.PickVisualMediaRequest
import androidx.activity.result.contract.ActivityResultContracts
import androidx.compose.foundation.Image
import androidx.compose.foundation.clickable
import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.layout.Box
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.PaddingValues
import androidx.compose.foundation.layout.Row
import androidx.compose.foundation.layout.Spacer
import androidx.compose.foundation.layout.fillMaxSize
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.lazy.LazyColumn
import androidx.compose.foundation.lazy.items
import androidx.compose.foundation.lazy.rememberLazyListState
import androidx.compose.material3.Button
import androidx.compose.material3.Card
import androidx.compose.material3.CardDefaults
import androidx.compose.material3.CircularProgressIndicator
import androidx.compose.material3.DropdownMenu
import androidx.compose.material3.DropdownMenuItem
import androidx.compose.material3.MaterialTheme
import androidx.compose.material3.OutlinedButton
import androidx.compose.material3.OutlinedTextField
import androidx.compose.material3.Text
import androidx.compose.material3.TextButton
import androidx.compose.runtime.Composable
import androidx.compose.runtime.DisposableEffect
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableLongStateOf
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.rememberCoroutineScope
import androidx.compose.runtime.setValue
import androidx.compose.runtime.snapshotFlow
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.ImageBitmap
import androidx.compose.ui.graphics.asImageBitmap
import androidx.compose.ui.platform.LocalContext
import androidx.compose.ui.unit.dp
import java.util.concurrent.atomic.AtomicLong
import java.util.concurrent.atomic.AtomicReference
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext

private const val RECONNECT_DELAY_MS = 1500L

/**
 * What the transcript renders: the event stream folded into displayable rows (see [foldEvent]). The
 * stream is the only data source -- opening this screen replays from seq 0, and a reconnect resumes
 * from the last seq seen, so there is no separate history fetch to drift from it.
 */
sealed class TranscriptItem {
    data class UserMsg(val text: String) : TranscriptItem()

    data class AssistantMsg(val text: String) : TranscriptItem()

    data class ToolRun(
        val id: String,
        val tool: String,
        val input: String,
        val output: String,
        val done: Boolean,
    ) : TranscriptItem()

    data class QuestionCard(
        val id: String,
        val prompt: String,
        val options: List<String>,
        val answer: String?,
    ) : TranscriptItem()

    data class ErrorMsg(val message: String) : TranscriptItem()

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

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

fun foldEvent(items: List<TranscriptItem>, event: SessionEvent): List<TranscriptItem> =
    when (event) {
        is SessionEvent.UserMessage -> items + TranscriptItem.UserMsg(event.text)
        is SessionEvent.AssistantText -> {
            // Deltas accumulate into the message they're streaming.
            val last = items.lastOrNull()
            if (last is TranscriptItem.AssistantMsg) {
                items.dropLast(1) + last.copy(text = last.text + event.delta)
            } else {
                items + TranscriptItem.AssistantMsg(event.delta)
            }
        }
        is SessionEvent.ToolStart ->
            items + TranscriptItem.ToolRun(event.id, event.tool, event.input, "", done = false)
        is SessionEvent.ToolUpdate -> updateTool(items, event.id) { it.copy(output = event.output) }
        is SessionEvent.ToolEnd ->
            updateTool(items, event.id) { it.copy(output = event.output, done = true) }
        is SessionEvent.Question ->
            items +
                TranscriptItem.QuestionCard(event.id, event.prompt, event.options, answer = null)
        is SessionEvent.Answered ->
            items.map {
                if (it is TranscriptItem.QuestionCard && it.id == event.id)
                    it.copy(answer = event.answer)
                else it
            }
        is SessionEvent.Error -> items + TranscriptItem.ErrorMsg(event.message)
        is SessionEvent.Image -> items + TranscriptItem.ImageItem(event.ref)
        is SessionEvent.Unknown -> items + TranscriptItem.Note("[${event.type}]")
        // Screen-level state, not transcript rows -- see SessionScreen.
        is SessionEvent.Status,
        is SessionEvent.UsageDelta -> items
    }

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
}

@Composable
fun SessionScreen(
    settings: ServerSettings,
    summary: SessionSummary,
    onBack: () -> Unit,
    onUsage: () -> Unit,
) {
    val scope = rememberCoroutineScope()
    var items by remember { mutableStateOf(listOf<TranscriptItem>()) }
    var status by remember { mutableStateOf(summary.status) }
    var totalTokens by remember { mutableLongStateOf(0L) }
    var streamError by remember { mutableStateOf<String?>(null) }
    var actionError by remember { mutableStateOf<String?>(null) }
    var input by remember { mutableStateOf("") }
    var expandedTools by remember { mutableStateOf(setOf<String>()) }
    // Uploaded-but-not-yet-sent attachment ids; sent with the next message.
    var pendingAttachments by remember { mutableStateOf(listOf<String>()) }
    // What this session is set to now, seeded from the row that opened it and
    // then owned here, because changing either is something this screen does.
    var model by remember { mutableStateOf(summary.model) }
    var permissionMode by remember { mutableStateOf(summary.permissionMode ?: "auto") }
    // The models this provider actually offers, asked of the server rather
    // than listed here: a hardcoded list is a claim about a machine.
    var offeredModels by remember { mutableStateOf<List<String>>(emptyList()) }
    val context = LocalContext.current
    // The resume cursor, written from the stream's IO thread.
    val lastSeq = remember { AtomicLong(0) }
    val activeStream = remember { AtomicReference<EventStream?>(null) }
    val listState = rememberLazyListState()

    fun apply(entry: SeqEvent) {
        lastSeq.set(entry.seq)
        when (val event = entry.event) {
            is SessionEvent.Status -> status = event.state
            is SessionEvent.UsageDelta -> totalTokens += event.tokens
            else -> items = foldEvent(items, event)
        }
    }

    // The stream lifecycle: connect, follow, and on any drop reconnect
    // from the cursor -- so a flaky link (or a backend restart) costs
    // nothing but the gap's latency.
    LaunchedEffect(summary.id) {
        while (true) {
            val stream = EventStream(settings, summary.id)
            activeStream.set(stream)
            try {
                withContext(Dispatchers.IO) {
                    stream.run(lastSeq.get()) { entry ->
                        apply(entry)
                        streamError = null
                    }
                }
            } catch (e: ApiException) {
                streamError = e.message
            } finally {
                stream.close()
            }
            delay(RECONNECT_DELAY_MS)
        }
    }
    // Coroutine cancellation can't interrupt a blocking socket read;
    // closing the stream is what unblocks it when this screen goes away.
    DisposableEffect(summary.id) { onDispose { activeStream.get()?.close() } }

    // Whether the view is pinned to the newest item. It is the reader's
    // scroll that decides: settling anywhere above the bottom releases it,
    // settling back at the bottom re-arms it. Written only when a scroll
    // *ends* so that the pin's state survives the moments when new content
    // has just pushed the bottom away but the reader never moved.
    var followTail by remember { mutableStateOf(true) }
    LaunchedEffect(listState) {
        snapshotFlow { listState.isScrollInProgress }
            .collect { scrolling -> if (!scrolling) followTail = !listState.canScrollForward }
    }
    // Two things move the bottom out from under the reader: a new item,
    // and the viewport shrinking when the keyboard opens. Watching only
    // item count handled the first and left the input box typing into a
    // view whose tail had slid under the IME. `scrollToItem` rather than
    // animated: on an imported session hundreds of items arrive at once,
    // and animating through them is a light show, not scrolling.
    LaunchedEffect(listState) {
        snapshotFlow { items.size to listState.layoutInfo.viewportSize.height }
            .collect { (count, _) ->
                if (followTail && count > 0) listState.scrollToItem(count - 1)
            }
    }

    LaunchedEffect(summary.setupName, summary.provider) {
        offeredModels =
            try {
                withContext(Dispatchers.IO) {
                    fetchSetups(settings)
                        .firstOrNull { it.name == summary.setupName }
                        ?.providers
                        ?.firstOrNull { it.name == summary.provider }
                        ?.models
                        .orEmpty()
                }
            } catch (_: Exception) {
                // Not worth reporting: the picker simply has nothing to
                // offer, which is visible, and the session is unaffected.
                emptyList()
            }
    }

    fun act(action: () -> Unit) {
        scope.launch {
            try {
                withContext(Dispatchers.IO) { action() }
                actionError = null
            } catch (e: ApiException) {
                actionError = e.message
            }
        }
    }

    fun send() {
        val text = input.trim()
        val attachments = pendingAttachments
        if (text.isEmpty() && attachments.isEmpty()) return
        input = ""
        pendingAttachments = emptyList()
        act { sendMessage(settings, summary.id, text, attachments) }
    }

    // The system photo picker; the image uploads as soon as it's chosen,
    // so Send only has ids to reference.
    val pickImage =
        rememberLauncherForActivityResult(ActivityResultContracts.PickVisualMedia()) { uri ->
            if (uri != null) {
                scope.launch {
                    try {
                        val id =
                            withContext(Dispatchers.IO) {
                                val bytes =
                                    context.contentResolver.openInputStream(uri)?.use {
                                        it.readBytes()
                                    } ?: throw ApiException("couldn't read the picked image")
                                val mime = context.contentResolver.getType(uri) ?: "image/jpeg"
                                uploadAttachment(settings, summary.id, bytes, mime)
                            }
                        pendingAttachments = pendingAttachments + id
                        actionError = null
                    } catch (e: ApiException) {
                        actionError = e.message
                    }
                }
            }
        }

    Column(Modifier.fillMaxSize()) {
        Row(
            verticalAlignment = Alignment.CenterVertically,
            modifier = Modifier.fillMaxWidth().padding(horizontal = 8.dp, vertical = 4.dp),
        ) {
            TextButton(onClick = onBack) { Text("Back") }
            Column(Modifier.weight(1f)) {
                Text(summary.title, style = MaterialTheme.typography.titleMedium)
                Text(
                    listOfNotNull(
                            summary.provider,
                            "on ${summary.setupName}",
                            summary.model,
                            if (totalTokens > 0) "$totalTokens tok" else null,
                        )
                        .joinToString(" · "),
                    style = MaterialTheme.typography.bodySmall,
                    color = MaterialTheme.colorScheme.onSurfaceVariant,
                )
            }
            StatusText(status)
            // Beside the provider it reports on, which is the line directly to its left.
            //
            // Its real home is this provider's settings, which do not exist yet; until they do,
            // the session is the only place the provider is already named, so it is the only
            // place the button can sit without inventing a scope for itself. What it shows is
            // the paid service's own numbers, so a session on a provider with no such service
            // gets an honest "unavailable" rather than a hidden button -- a control that comes
            // and goes makes its absence the signal, and absence cannot say why.
            TextButton(onClick = onUsage) { Text("Usage") }
        }

        (streamError ?: actionError)?.let { message ->
            Text(
                message,
                color = MaterialTheme.colorScheme.error,
                style = MaterialTheme.typography.bodySmall,
                modifier = Modifier.padding(horizontal = 16.dp, vertical = 4.dp),
            )
        }

        LazyColumn(
            state = listState,
            modifier = Modifier.weight(1f).fillMaxWidth(),
            contentPadding = PaddingValues(16.dp),
            verticalArrangement = Arrangement.spacedBy(8.dp),
        ) {
            items(items) { item ->
                when (item) {
                    is TranscriptItem.UserMsg -> UserBubble(item.text)
                    is TranscriptItem.AssistantMsg ->
                        Text(item.text, style = MaterialTheme.typography.bodyLarge)
                    is TranscriptItem.ToolRun ->
                        ToolCard(
                            tool = item,
                            expanded = item.id in expandedTools,
                            onToggle = {
                                expandedTools =
                                    if (item.id in expandedTools) expandedTools - item.id
                                    else expandedTools + item.id
                            },
                        )
                    is TranscriptItem.QuestionCard ->
                        QuestionRow(item) { answer ->
                            act { answerQuestion(settings, summary.id, item.id, answer) }
                        }
                    is TranscriptItem.ErrorMsg ->
                        Text(
                            item.message,
                            color = MaterialTheme.colorScheme.error,
                            style = MaterialTheme.typography.bodyMedium,
                        )
                    is TranscriptItem.ImageItem -> SessionImage(settings, summary.id, item.ref)
                    is TranscriptItem.Note ->
                        Text(
                            item.text,
                            style = MaterialTheme.typography.bodySmall,
                            color = MaterialTheme.colorScheme.onSurfaceVariant,
                        )
                }
            }
        }

        // Always enabled -- a send while the session is running becomes a
        // steering message injected at the next tool boundary, which is
        // the point of the whole app.
        //
        // The field gets a row of its own, above the buttons: sharing one
        // put the full width behind three controls, so the thing being
        // typed into was the narrowest thing on the row.
        Column(Modifier.fillMaxWidth().padding(8.dp)) {
            OutlinedTextField(
                value = input,
                onValueChange = { input = it },
                modifier = Modifier.fillMaxWidth(),
                placeholder = {
                    Text(if (pendingAttachments.isEmpty()) "Message" else "Message (+image)")
                },
                maxLines = 4,
            )
            Row(
                verticalAlignment = Alignment.CenterVertically,
                modifier = Modifier.fillMaxWidth(),
            ) {
                TextButton(
                    onClick = {
                        pickImage.launch(
                            PickVisualMediaRequest(
                                ActivityResultContracts.PickVisualMedia.ImageOnly
                            )
                        )
                    }
                ) {
                    Text(if (pendingAttachments.isEmpty()) "+" else "+${pendingAttachments.size}")
                }
                // Beside the field they govern, and showing their current
                // value rather than a label: what this session is set to is
                // the thing worth reading at a glance, and the control for
                // changing it is the same object.
                if (offeredModels.isNotEmpty()) {
                    PickerButton(
                        current = model ?: "default",
                        options = offeredModels,
                        onPick = { chosen ->
                            model = chosen
                            act { setSessionModel(settings, summary.id, chosen) }
                        },
                    )
                }
                PickerButton(
                    current = permissionMode,
                    options = PERMISSION_MODES,
                    onPick = { chosen ->
                        permissionMode = chosen
                        act { setSessionPermissionMode(settings, summary.id, chosen) }
                    },
                )
                Spacer(Modifier.weight(1f))
                if (status == "running" || status == "compacting") {
                    OutlinedButton(onClick = { act { interruptSession(settings, summary.id) } }) {
                        Text("Stop")
                    }
                    Spacer(Modifier.width(8.dp))
                }
                Button(onClick = { send() }) { Text("Send") }
            }
        }
    }
}

/**
 * An inline transcript image, fetched (authenticated, pinned) from the session's files route. The
 * bitmap is remembered per ref, so scrolling doesn't refetch.
 */
@Composable
private fun SessionImage(settings: ServerSettings, sessionId: String, ref: String) {
    var bitmap by remember(ref) { mutableStateOf<ImageBitmap?>(null) }
    var failed by remember(ref) { mutableStateOf(false) }
    LaunchedEffect(ref) {
        try {
            val bytes = withContext(Dispatchers.IO) { fetchSessionFile(settings, sessionId, ref) }
            bitmap = BitmapFactory.decodeByteArray(bytes, 0, bytes.size)?.asImageBitmap()
            failed = bitmap == null
        } catch (_: ApiException) {
            failed = true
        }
    }
    when (val image = bitmap) {
        null ->
            Text(
                if (failed) "[image $ref unavailable]" else "[loading image…]",
                style = MaterialTheme.typography.bodySmall,
                color = MaterialTheme.colorScheme.onSurfaceVariant,
            )
        else ->
            Image(
                bitmap = image,
                contentDescription = "session image",
                modifier = Modifier.fillMaxWidth(),
            )
    }
}

@Composable
private fun UserBubble(text: String) {
    Box(Modifier.fillMaxWidth()) {
        Card(
            colors =
                CardDefaults.cardColors(
                    containerColor = MaterialTheme.colorScheme.primaryContainer
                ),
            modifier = Modifier.align(Alignment.CenterEnd).padding(start = 48.dp),
        ) {
            Text(text, modifier = Modifier.padding(12.dp))
        }
    }
}

/**
 * Collapsed by default: name plus a spinner while running, expandable to the input and output. The
 * spinner-while-unfinished is exactly "ToolStart with no matching ToolEnd yet".
 */
@Composable
private fun ToolCard(tool: TranscriptItem.ToolRun, expanded: Boolean, onToggle: () -> Unit) {
    Card(Modifier.fillMaxWidth().clickable(onClick = onToggle)) {
        Column(Modifier.padding(12.dp)) {
            Row(verticalAlignment = Alignment.CenterVertically) {
                Text(
                    tool.tool,
                    style = MaterialTheme.typography.titleSmall,
                    modifier = Modifier.weight(1f),
                )
                if (!tool.done) {
                    CircularProgressIndicator(
                        modifier = Modifier.width(16.dp).height(16.dp),
                        strokeWidth = 2.dp,
                    )
                }
            }
            if (expanded) {
                Spacer(Modifier.height(8.dp))
                Text("Input", style = MaterialTheme.typography.labelSmall)
                Text(tool.input, style = MaterialTheme.typography.bodySmall)
                if (tool.output.isNotEmpty()) {
                    Spacer(Modifier.height(8.dp))
                    Text("Output", style = MaterialTheme.typography.labelSmall)
                    Text(tool.output, style = MaterialTheme.typography.bodySmall)
                }
            }
        }
    }
}

/**
 * A question (or permission request -- same shape) inline in the transcript. Option buttons until
 * answered; then the chosen answer, which the `answered` event also resolves on every other
 * connected device.
 */
@Composable
private fun QuestionRow(question: TranscriptItem.QuestionCard, onAnswer: (String) -> Unit) {
    Card(Modifier.fillMaxWidth()) {
        Column(Modifier.padding(12.dp)) {
            Text(question.prompt, style = MaterialTheme.typography.bodyLarge)
            Spacer(Modifier.height(8.dp))
            if (question.answer != null) {
                Text(
                    "Answered: ${question.answer}",
                    style = MaterialTheme.typography.labelLarge,
                    color = MaterialTheme.colorScheme.onSurfaceVariant,
                )
            } else {
                Row(horizontalArrangement = Arrangement.spacedBy(8.dp)) {
                    question.options.forEach { option ->
                        OutlinedButton(onClick = { onAnswer(option) }) { Text(option) }
                    }
                }
            }
        }
    }
}

/** The modes the CLI accepts, in the order they give up asking. */
private val PERMISSION_MODES = listOf("manual", "acceptEdits", "auto", "bypassPermissions", "plan")

/**
 * A control that reads as its own value.
 *
 * The button *is* the current setting rather than a label beside one, so the row says what the
 * session is set to without spending a second line on saying it.
 */
@Composable
private fun PickerButton(current: String, options: List<String>, onPick: (String) -> Unit) {
    var open by remember { mutableStateOf(false) }
    Box {
        TextButton(onClick = { open = true }) {
            Text(current, style = MaterialTheme.typography.bodySmall)
        }
        DropdownMenu(expanded = open, onDismissRequest = { open = false }) {
            options.forEach { option ->
                DropdownMenuItem(
                    text = { Text(option) },
                    onClick = {
                        open = false
                        if (option != current) onPick(option)
                    },
                )
            }
        }
    }
}