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

package com.example.aiapp

import androidx.compose.foundation.layout.Arrangement
import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.ExperimentalLayoutApi
import androidx.compose.foundation.layout.FlowRow
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.rememberScrollState
import androidx.compose.foundation.verticalScroll
import androidx.compose.material3.Button
import androidx.compose.material3.CircularProgressIndicator
import androidx.compose.material3.FilterChip
import androidx.compose.material3.MaterialTheme
import androidx.compose.material3.OutlinedTextField
import androidx.compose.material3.Text
import androidx.compose.material3.TextButton
import androidx.compose.runtime.Composable
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.unit.dp
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext

/**
 * The spawn screen: what to run, where to run it, and the per-kind fields.
 *
 * Providers and hosts both come from the server, so adding either to its config.ron shows up here
 * with no app rebuild.
 */
@Composable
fun SpawnScreen(
    settings: ServerSettings,
    onSpawned: (SessionSummary) -> Unit,
    onBack: () -> Unit,
) {
    val scope = rememberCoroutineScope()
    // What the form is made of, and whether we have it yet. A failure here is not the same as a
    // server with nothing to offer, so it must not reach the pickers as empty lists.
    var options by remember { mutableStateOf<LoadState<List<Setup>>>(LoadState.Loading) }

    // Setup first, then one of its providers. Choosing a setup can invalidate the provider, so the
    // provider is stored by name and resolved against the current setup rather than held as an
    // object that could outlive the list it came from.
    var setupName by remember { mutableStateOf<String?>(null) }
    var providerName by remember { mutableStateOf<String?>(null) }
    var title by remember { mutableStateOf("") }
    var model by remember { mutableStateOf("") }
    var cwd by remember { mutableStateOf("") }
    // "auto" rather than "manual": on a phone every ask is a round trip to a question card, and
    // answering "allow Bash?" dozens of times per task is what this app exists to avoid.
    var permissionMode by remember { mutableStateOf("auto") }
    var busy by remember { mutableStateOf(false) }
    // Only the spawn's own failure. The fetch's lives in `options`: this one leaves a filled-in
    // form worth keeping, and that one leaves nothing to fill in.
    var spawnError by remember { mutableStateOf<String?>(null) }
    // Downloaded models, for a llama provider to choose between. Kept separate from the setups: a
    // Claude session needs none, so failing to list them must not stop the screen rendering.
    var models by remember { mutableStateOf<List<LocalModel>>(emptyList()) }
    var modelKey by remember { mutableStateOf<String?>(null) }
    var contextSize by remember { mutableStateOf("") }
    var temperature by remember { mutableStateOf("") }

    LaunchedEffect(Unit) {
        options =
            try {
                val fetched = withContext(Dispatchers.IO) { fetchSetups(settings) }
                val first = fetched.firstOrNull()
                setupName = first?.name
                providerName = first?.providers?.firstOrNull()?.name
                LoadState.Loaded(fetched)
            } catch (e: ApiException) {
                LoadState.failed(e)
            }
        models =
            runCatching { withContext(Dispatchers.IO) { fetchModels(settings).local } }
                .getOrDefault(emptyList())
    }

    Column(Modifier.fillMaxSize().verticalScroll(rememberScrollState()).padding(16.dp)) {
        Row(verticalAlignment = Alignment.CenterVertically, modifier = Modifier.fillMaxWidth()) {
            Text(
                "New session",
                style = MaterialTheme.typography.headlineSmall,
                modifier = Modifier.weight(1f),
            )
            TextButton(onClick = onBack) { Text("Cancel") }
        }
        Spacer(Modifier.height(16.dp))

        // Nothing below is fillable until the options are here, and a failure to fetch them leaves
        // no form worth showing -- so this reports and stops, rather than offering empty pickers
        // under an error message.
        val setups =
            when (val state = options) {
                is LoadState.Loading -> {
                    CircularProgressIndicator()
                    return@Column
                }
                is LoadState.Error -> {
                    Text(state.message, color = MaterialTheme.colorScheme.error)
                    return@Column
                }
                is LoadState.Loaded -> state.value
            }
        val setup = setups.firstOrNull { it.name == setupName }
        val current = setup?.providers?.firstOrNull { it.name == providerName }
        // Only the Claude CLI has models, a working directory and permission modes; keying the
        // extra fields on the kind rather than the provider name keeps a second Claude provider
        // from needing anything here.
        val isClaude = current?.kind == "claude_cli"
        val isLlama = current?.kind == "llama_cpp"

        // The machine first, because it decides what can be run at all.
        ChipGroup(
            label = "Setup",
            options = setups.map { it.name },
            selected = setupName,
            onSelect = { name ->
                setupName = name
                // The provider list changes with the machine, so a name carried over from the
                // previous one would be a selection that isn't in the picker. Take that machine's
                // first.
                providerName =
                    setups.firstOrNull { it.name == name }?.providers?.firstOrNull()?.name
            },
        )
        setup?.address?.let {
            Text(
                it,
                style = MaterialTheme.typography.bodySmall,
                color = MaterialTheme.colorScheme.onSurfaceVariant,
            )
            // The address belongs to the setup above it, not to the provider label below; without
            // this they read as one block.
            Spacer(Modifier.height(8.dp))
        }

        // Only what this machine actually has. A setup with none says so rather than showing an
        // empty row that reads as a failure.
        if (setup != null && setup.providers.isEmpty()) {
            Text(
                "\"${setup.name}\" has no providers configured.",
                style = MaterialTheme.typography.bodyMedium,
                color = MaterialTheme.colorScheme.onSurfaceVariant,
            )
        } else {
            ChipGroup(
                label = "Provider",
                options = setup?.providers?.map { it.name }.orEmpty(),
                selected = providerName,
                onSelect = { providerName = it },
            )
        }

        Spacer(Modifier.height(16.dp))

        OutlinedTextField(
            value = title,
            onValueChange = { title = it },
            label = { Text("Title") },
            singleLine = true,
            modifier = Modifier.fillMaxWidth(),
        )

        if (isLlama) {
            // A llama session names one of the models this backend has downloaded, so the choice is
            // that list rather than free text -- a name that is not on disk is a session that
            // cannot start.
            if (models.isEmpty()) {
                Text(
                    "No models downloaded yet. Get one from the Models screen first.",
                    style = MaterialTheme.typography.bodyMedium,
                    color = MaterialTheme.colorScheme.onSurfaceVariant,
                )
            } else {
                ChipGroup(
                    label = "Model",
                    // The file, not the whole key: the repository is the same for every
                    // quantisation of a model, so the file name is what tells two of them apart.
                    options = models.map { it.file },
                    selected = models.firstOrNull { it.key == modelKey }?.file,
                    onSelect = { file -> modelKey = models.first { it.file == file }.key },
                )
            }
            Spacer(Modifier.height(16.dp))

            OutlinedTextField(
                value = contextSize,
                onValueChange = { contextSize = it },
                label = { Text("Context size (blank = the model's default)") },
                singleLine = true,
                modifier = Modifier.fillMaxWidth(),
            )
            Spacer(Modifier.height(16.dp))

            OutlinedTextField(
                value = temperature,
                onValueChange = { temperature = it },
                label = { Text("Temperature (blank = llama.cpp's default)") },
                singleLine = true,
                modifier = Modifier.fillMaxWidth(),
            )
            Spacer(Modifier.height(16.dp))
        }

        if (isClaude) {
            if (current.models.isNotEmpty()) {
                Spacer(Modifier.height(16.dp))
                ChipGroup(
                    label = "Model",
                    options = current.models,
                    selected = model.ifEmpty { null },
                    onSelect = { chosen -> model = if (model == chosen) "" else chosen },
                )
            }
            Spacer(Modifier.height(8.dp))
            OutlinedTextField(
                value = model,
                onValueChange = { model = it },
                label = { Text("Model (blank = the CLI's default)") },
                singleLine = true,
                modifier = Modifier.fillMaxWidth(),
            )
            Spacer(Modifier.height(16.dp))

            OutlinedTextField(
                value = cwd,
                onValueChange = { cwd = it },
                label = { Text("Working directory") },
                placeholder = { Text("/home/…") },
                singleLine = true,
                modifier = Modifier.fillMaxWidth(),
            )
            Spacer(Modifier.height(16.dp))

            ChipGroup(
                label = "Permissions",
                options = PERMISSION_MODES,
                selected = permissionMode,
                onSelect = { permissionMode = it },
            )
        }
        Spacer(Modifier.height(24.dp))

        // Beside the button that produced it.
        spawnError?.let {
            Text(it, color = MaterialTheme.colorScheme.error)
            Spacer(Modifier.height(8.dp))
        }

        Button(
            onClick = {
                val chosen = current ?: return@Button
                busy = true
                scope.launch {
                    try {
                        val spawned =
                            withContext(Dispatchers.IO) {
                                spawnSession(
                                    settings,
                                    // The id, not the label: labels are editable and the server
                                    // resolves by id. Non-null here, since `chosen` came from
                                    // `setup`'s own provider list.
                                    setup = setup.id,
                                    provider = chosen.name,
                                    title = title.trim(),
                                    model =
                                        if (isLlama) modelKey else model.trim().takeIf { isClaude },
                                    cwd = cwd.trim().takeIf { isClaude },
                                    permissionMode = permissionMode.takeIf { isClaude },
                                    // Sent only when set, so blank means "whatever llama.cpp does
                                    // by default" rather than a zero.
                                    params =
                                        buildMap {
                                            if (isLlama) {
                                                contextSize
                                                    .trim()
                                                    .takeIf { it.isNotEmpty() }
                                                    ?.let { put("contextSize", it) }
                                                temperature
                                                    .trim()
                                                    .takeIf { it.isNotEmpty() }
                                                    ?.let { put("temperature", it) }
                                            }
                                        },
                                )
                            }
                        onSpawned(spawned)
                    } catch (e: ApiException) {
                        spawnError = e.message
                        busy = false
                    }
                }
            },
            enabled = !busy && current != null && !(isLlama && modelKey == null),
        ) {
            Text(if (busy) "Spawning..." else "Spawn")
        }
    }
}

/**
 * A labeled row of choices that wraps onto as many lines as it needs.
 *
 * FlowRow rather than Row: a plain Row gives every chip an equal share of a single line, so once
 * the options don't fit, the text inside each one wraps to one character per line instead of the
 * row wrapping.
 */
@OptIn(ExperimentalLayoutApi::class)
@Composable
fun ChipGroup(
    label: String,
    options: List<String>,
    selected: String?,
    onSelect: (String) -> Unit,
) {
    Text(label, style = MaterialTheme.typography.labelLarge)
    FlowRow(
        horizontalArrangement = Arrangement.spacedBy(8.dp),
        verticalArrangement = Arrangement.spacedBy(4.dp),
        modifier = Modifier.fillMaxWidth(),
    ) {
        options.forEach { option ->
            FilterChip(
                selected = selected == option,
                onClick = { onSelect(option) },
                label = { Text(option) },
            )
        }
    }
}