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

package com.example.aiapp

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.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.itemsIndexed
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.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.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.rememberCoroutineScope
import androidx.compose.runtime.setValue
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.asImageBitmap
import androidx.compose.ui.unit.dp
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import java.util.concurrent.atomic.AtomicLong
import java.util.concurrent.atomic.AtomicReference

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) {
    val scope = rememberCoroutineScope()
    var items by remember { mutableStateOf(listOf<TranscriptItem>()) }
    var status by remember { mutableStateOf(summary.status) }
    var totalTokens by remember { mutableStateOf(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>()) }
    val context = androidx.compose.ui.platform.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() }
    }

    LaunchedEffect(items.size) {
        if (items.isNotEmpty()) listState.animateScrollToItem(items.size - 1)
    }

    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 = androidx.activity.compose.rememberLauncherForActivityResult(
        androidx.activity.result.contract.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.kind,
                        summary.model,
                        if (totalTokens > 0) "$totalTokens tok" else null,
                    ).joinToString(" · "),
                    style = MaterialTheme.typography.bodySmall,
                    color = MaterialTheme.colorScheme.onSurfaceVariant,
                )
            }
            StatusText(status)
        }

        (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 = androidx.compose.foundation.layout.PaddingValues(16.dp),
            verticalArrangement = Arrangement.spacedBy(8.dp),
        ) {
            itemsIndexed(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.
        Row(
            verticalAlignment = Alignment.CenterVertically,
            modifier = Modifier.fillMaxWidth().padding(8.dp),
        ) {
            TextButton(onClick = {
                pickImage.launch(
                    androidx.activity.result.PickVisualMediaRequest(
                        androidx.activity.result.contract.ActivityResultContracts
                            .PickVisualMedia.ImageOnly,
                    ),
                )
            }) {
                Text(if (pendingAttachments.isEmpty()) "+" else "+${pendingAttachments.size}")
            }
            OutlinedTextField(
                value = input,
                onValueChange = { input = it },
                modifier = Modifier.weight(1f),
                placeholder = { Text(if (pendingAttachments.isEmpty()) "Message" else "Message (+image)") },
                maxLines = 4,
            )
            Spacer(Modifier.width(8.dp))
            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<androidx.compose.ui.graphics.ImageBitmap?>(null)
    }
    var failed by remember(ref) { mutableStateOf(false) }
    LaunchedEffect(ref) {
        try {
            val bytes = withContext(Dispatchers.IO) { fetchSessionFile(settings, sessionId, ref) }
            bitmap = android.graphics.BitmapFactory.decodeByteArray(bytes, 0, bytes.size)
                ?.asImageBitmap()
            failed = bitmap == null
        } catch (e: 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 -> androidx.compose.foundation.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) }
                    }
                }
            }
        }
    }
}