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") }
    // Null until the server has been asked, and null again if it answers "no level chosen" -- the
    // two are told apart by [defaultsAsked], because a picker that shows a level before the answer
    // arrives is one you can spawn at without having chosen it.
    var effort by remember { mutableStateOf<String?>(null) }
    var defaultsAsked by remember { mutableStateOf(false) }
    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) }
    // The models on the *chosen machine*, 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. Refetched when the machine changes, because a model is a file on one
    // machine -- see [fetchSetupModels].
    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) {
        // Separate from the setups fetch below and deliberately not fatal: failing to learn the
        // default must leave a screen you can still spawn from, so the picker stays on "default"
        // and says so rather than the whole form refusing to draw.
        runCatching { withContext(Dispatchers.IO) { fetchDefaultEffort(settings) } }
            .onSuccess { effort = it }
        defaultsAsked = true
        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)
            }
    }

    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 }
        // Whichever machine is chosen now, asked again when that changes. The old machine's list
        // is dropped first rather than left on screen: a file name from another machine looks
        // exactly like one from this one.
        LaunchedEffect(setup?.id) {
            models = emptyList()
            modelKey = null
            val id = setup?.id ?: return@LaunchedEffect
            models =
                runCatching { withContext(Dispatchers.IO) { fetchSetupModels(settings, id) } }
                    .getOrDefault(emptyList())
        }
        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 on the machine it will run on, so the
            // choice is that list rather than free text -- a name that is not on that machine's
            // disk is a session that cannot start.
            if (models.isEmpty()) {
                Text(
                    "No models on ${setup?.name ?: "this machine"}. The Models screen downloads " +
                        "to the backend; another machine needs the file put there itself.",
                    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(16.dp))

            // Says what it does to *later* spawns as well, because it does: the level chosen here
            // is stored as the default, which is the whole way that default is set. A picker that
            // quietly changed a global would be the same control with the fact left out.
            ChipGroup(
                label = "Thinking (kept as the default for new sessions)",
                options = listOf(DEFAULT_EFFORT) + EFFORT_LEVELS,
                // The CLI's own default is a level in the list, so this cannot be a one-way trip.
                // Disabled-looking until the server has answered, for the reason above.
                selected = if (defaultsAsked) effort ?: DEFAULT_EFFORT else null,
                onSelect = { chosen -> effort = chosen.takeIf { it != DEFAULT_EFFORT } },
            )
        }
        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) {
                                // Stored before the spawn and not after it: choosing a level is
                                // an intent about new sessions in general, so a spawn that then
                                // fails must not also lose the choice. Non-fatal for the same
                                // reason the fetch above is -- the session is what was asked for.
                                if (isClaude) {
                                    runCatching { setDefaultEffort(settings, effort) }
                                }
                                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 },
                                    effort = effort.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) },
            )
        }
    }
}