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<Machine>>>(LoadState.Loading) }

    // Machine first, then one of its providers. Choosing a machine can invalidate the provider, so
    // the
    // provider is stored by name and resolved against the current machine rather than held as an
    // object that could outlive the list it came from.
    var machineName by remember { mutableStateOf<String?>(null) }
    var providerName by remember { mutableStateOf<String?>(null) }
    var title by remember { mutableStateOf("") }
    var model by remember { mutableStateOf("") }
    var providerModels by remember { mutableStateOf<List<OfferedModel>>(emptyList()) }
    var providerModelsLoading by remember { mutableStateOf(false) }
    var providerModelsError by remember { mutableStateOf<String?>(null) }
    var cwd by remember { mutableStateOf("") }
    // Set only after the selected provider reports its own default. An empty value is not sent.
    var permissionMode by remember { mutableStateOf("") }
    // 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) }
    var contextSize by remember { mutableStateOf("") }
    var temperature by remember { mutableStateOf("") }
    // Whether a model that carries a multi-token-prediction head drafts with it. Left to the
    // server by default, which turns it on exactly where the file has one -- see `SPECULATIVE` in
    // the llama driver. Here so a machine where drafting does not pay has a way out that is not
    // an edit to config.ron.
    var speculative by remember { mutableStateOf(SPECULATIVE_AUTO) }

    LaunchedEffect(Unit) {
        // Separate from the machines 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) { fetchMachines(settings) }
                val first = fetched.firstOrNull()
                machineName = 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 machines =
            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 machine = machines.firstOrNull { it.name == machineName }
        val current = machine?.providers?.firstOrNull { it.name == providerName }
        // Coding CLIs take a working directory, model, permission mode and thinking level. Keying
        // the extra fields on the kind rather than the provider name keeps a second installation
        // from needing anything here.
        val isClaude = current?.kind == "claude_cli"
        val isCodex = current?.kind == "codex_cli"
        val isCodingCli = isClaude || isCodex
        val isLlama = current?.kind == "llama_cpp"
        // Echo is the only kind with nothing to choose between.
        val offersModels = isCodingCli || isLlama
        // Where a session's tools act, which is the only thing a working directory decides.
        val takesCwd = isCodingCli || isLlama

        // Whichever machine and provider are chosen now, asked again when either changes. The
        // previous answer is dropped first rather than left on screen: a model name from another
        // machine looks exactly like one from this one.
        LaunchedEffect(machine?.id, current?.name) {
            model = ""
            providerModels = emptyList()
            providerModelsError = null
            permissionMode = current?.defaultPermissionMode.orEmpty()
            // Every kind that offers models at all, not only the coding CLIs: a llama provider
            // answers with the GGUFs on the machine it runs on, through the same call. One
            // question with one answer is what keeps the picker free of a branch on the kind.
            if (machine == null || current == null || !offersModels) {
                providerModelsLoading = false
                return@LaunchedEffect
            }
            providerModelsLoading = true
            try {
                providerModels =
                    withContext(Dispatchers.IO) {
                        fetchProviderModels(settings, machine.id, current.name)
                    }
            } catch (e: ApiException) {
                providerModelsError = e.message
            } finally {
                providerModelsLoading = false
            }
        }

        // The machine first, because it decides what can be run at all.
        ChipGroup(
            label = "Machine",
            options = machines.map { it.name },
            selected = machineName,
            onSelect = { name ->
                machineName = 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 =
                    machines.firstOrNull { it.name == name }?.providers?.firstOrNull()?.name
            },
        )
        machine?.address?.let {
            Text(
                it,
                style = MaterialTheme.typography.bodySmall,
                color = MaterialTheme.colorScheme.onSurfaceVariant,
            )
            // The address belongs to the machine 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 machine with none says so rather than showing an
        // empty row that reads as a failure.
        if (machine != null && machine.providers.isEmpty()) {
            Text(
                "\"${machine.name}\" has no providers configured.",
                style = MaterialTheme.typography.bodyMedium,
                color = MaterialTheme.colorScheme.onSurfaceVariant,
            )
        } else {
            ChipGroup(
                label = "Provider",
                options = machine?.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 (offersModels) {
            when {
                providerModelsLoading ->
                    Text(
                        "Loading model choices…",
                        style = MaterialTheme.typography.bodySmall,
                        color = MaterialTheme.colorScheme.onSurfaceVariant,
                    )
                providerModelsError != null ->
                    Text(
                        "Model choices unavailable: $providerModelsError",
                        style = MaterialTheme.typography.bodySmall,
                        color = MaterialTheme.colorScheme.error,
                    )
                // A llama session cannot start without one, so this says what to do about it
                // rather than only that there is nothing -- the models it needs are on the
                // machine that will serve them, which is not always this backend.
                providerModels.isEmpty() && isLlama ->
                    Text(
                        "No models on ${machine?.name}. The Models screen downloads " +
                            "to the backend; another machine needs the file put there itself.",
                        style = MaterialTheme.typography.bodyMedium,
                        color = MaterialTheme.colorScheme.onSurfaceVariant,
                    )
                providerModels.isEmpty() ->
                    Text(
                        "This machine reported no selectable models.",
                        style = MaterialTheme.typography.bodySmall,
                        color = MaterialTheme.colorScheme.onSurfaceVariant,
                    )
                else -> {
                    Spacer(Modifier.height(16.dp))
                    ChipGroup(
                        label = "Model",
                        // The label, and the id is what is sent: for a llama model those differ,
                        // since it is chosen by path and named by what is inside the file.
                        options = providerModels.map { it.label },
                        selected = providerModels.firstOrNull { it.id == model }?.label,
                        onSelect = { chosen ->
                            val id = providerModels.first { it.label == chosen }.id
                            // A llama session has to have one, so choosing the same chip twice
                            // must not clear it -- there is nothing to fall back to.
                            model = if (model == id && !isLlama) "" else id
                        },
                    )
                }
            }
            Spacer(Modifier.height(16.dp))
        }

        if (isCodingCli) {
            // Free text as well as the chips above: the catalog is a shortcut, and a CLI will
            // take a name it did not list.
            OutlinedTextField(
                value = model,
                onValueChange = { model = it },
                label = { Text("Model (blank = the CLI's default)") },
                singleLine = true,
                modifier = Modifier.fillMaxWidth(),
            )
            Spacer(Modifier.height(16.dp))
        }

        if (isLlama) {
            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))

            // Said as what it is rather than as "MTP": the reader is choosing whether the session
            // goes faster, and most models have nothing to turn on here at all.
            ChipGroup(
                label = "Speculative decoding (models that carry a draft head)",
                options = listOf(SPECULATIVE_AUTO, SPECULATIVE_OFF),
                selected = speculative,
                onSelect = { speculative = it },
            )
            Spacer(Modifier.height(16.dp))
        }

        // Every session whose tools act on files needs one, which is both kinds that have
        // tools -- a llama session's built-in tools run in it exactly as a CLI's do.
        if (takesCwd) {
            OutlinedTextField(
                value = cwd,
                onValueChange = { cwd = it },
                label = { Text("Working directory") },
                placeholder = { Text("/home/…") },
                singleLine = true,
                modifier = Modifier.fillMaxWidth(),
            )
            Spacer(Modifier.height(16.dp))
        }

        // Offered wherever the provider has modes, rather than where this screen believes it
        // does: the server is what knows, and llama.cpp grew them without this line changing.
        if (current != null && current.permissionModes.isNotEmpty()) {
            ChipGroup(
                label = "Permissions",
                options = current.permissionModes,
                selected = permissionMode,
                onSelect = { permissionMode = it },
            )
            Spacer(Modifier.height(16.dp))
        }

        if (isCodingCli) {
            // 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 (isCodingCli) {
                                    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
                                    // `machine`'s own provider list.
                                    machine = machine.id,
                                    provider = chosen.name,
                                    title = title.trim(),
                                    model = model.trim().takeIf { offersModels },
                                    cwd = cwd.trim().takeIf { takesCwd },
                                    permissionMode = permissionMode.takeIf { it.isNotEmpty() },
                                    effort = effort.takeIf { isCodingCli },
                                    // 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) }
                                                // Only the choice that changes anything: "auto"
                                                // is the absence of the setting, not a value of
                                                // it, so a session spawned without an opinion
                                                // carries none.
                                                if (speculative == SPECULATIVE_OFF) {
                                                    put("speculative", "off")
                                                }
                                            }
                                        },
                                )
                            }
                        onSpawned(spawned)
                    } catch (e: ApiException) {
                        spawnError = e.message
                        busy = false
                    }
                }
            },
            // A llama session names the file to load, so there is nothing to spawn without one.
            enabled = !busy && current != null && !(isLlama && model.isEmpty()),
        ) {
            Text(if (busy) "Spawning..." else "Spawn")
        }
    }
}

/**
 * Leave the draft head to the server, which uses one wherever the model file has one. Spelled the
 * same as the absence of the `speculative` parameter, because that is what it means.
 */
private const val SPECULATIVE_AUTO = "auto"

/** The `speculative` parameter's only other value; see the llama driver's `SPECULATIVE`. */
private const val SPECULATIVE_OFF = "off"

/**
 * 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) },
            )
        }
    }
}