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

package com.example.aiapp

import androidx.activity.compose.BackHandler
import androidx.compose.foundation.background
import androidx.compose.foundation.layout.Box
import androidx.compose.foundation.layout.fillMaxSize
import androidx.compose.foundation.layout.imePadding
import androidx.compose.foundation.layout.padding
import androidx.compose.material3.AlertDialog
import androidx.compose.material3.MaterialTheme
import androidx.compose.material3.Text
import androidx.compose.material3.TextButton
import androidx.compose.runtime.Composable
import androidx.compose.runtime.CompositionLocalProvider
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.getValue
import androidx.compose.runtime.key
import androidx.compose.runtime.mutableIntStateOf
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.rememberCoroutineScope
import androidx.compose.runtime.setValue
import androidx.compose.ui.Modifier
import androidx.compose.ui.platform.LocalContext
import androidx.compose.ui.semantics.clearAndSetSemantics
import androidx.compose.ui.unit.dp
import com.example.wgapplink.localNetworkAllowed
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext

/**
 * One `when` rather than a navigation library: a handful of screens, with [Screen.Main] as the root
 * and the back button the only other way between them.
 *
 * Import, models and machines are tabs inside [MainScreen] -- four views of the same backend, none
 * of them a step down from another -- and what is left here is only what genuinely is a step down:
 * one session, spawning one, and settings.
 */
private sealed class Screen {
    data object Main : Screen()

    /**
     * One session, with the file explorer or a subagent transcript over it when set.
     *
     * Both are layers on this screen rather than screens of their own, so the session under them
     * stays composed: its event stream keeps flowing, its scroll position and draft stay put, and
     * coming back costs nothing. As sibling `Screen`s they would dispose and recreate it on every
     * return, refetching the transcript over the tunnel.
     */
    data class Session(
        val summary: SessionSummary,
        val files: FilesTarget? = null,
        val subagent: SubagentSummary? = null,
    ) : Screen()

    data object Spawn : Screen()

    /**
     * One provider on one machine: its settings, and what its shared server is holding.
     *
     * A step down from the machines tab rather than a tab of its own, because it is about one
     * machine rather than about the backend. Addressed by ids and names rather than by the
     * [Provider] it was tapped from: what it shows is fetched, and a stale copy of a card would be
     * a second version of the same truth.
     */
    data class ProviderSettings(val machineId: String, val provider: String) : Screen()

    data object Settings : Screen()
}

/**
 * A session a notification tap asked to open, before it is a screen.
 *
 * The notification names an id and nothing else, so opening it means fetching the session first.
 * [serial] tells two taps on the same session's notification apart, since they are two requests and
 * would otherwise compare equal.
 */
data class SessionOpenRequest(val sessionId: String, val serial: Int)

/** A tap that could not be turned into a screen, kept with its request so Try again knows what. */
private data class FailedOpen(val request: SessionOpenRequest, val message: String)

/**
 * [settingsVersion] bumps when enrollment lands via an `aiapp://` intent (see MainActivity), re-
 * reading the stored settings -- a plain `remember` would keep serving the pre-enrollment null.
 *
 * [openRequest] is the session a notification tap asked for, likewise from MainActivity.
 *
 * [shareRequest] is what another app shared in, likewise. It is held here until a session takes it,
 * because the share arrives before anyone has said which session it is for.
 */
@Composable
fun AppRoot(
    settingsVersion: Int,
    openRequest: SessionOpenRequest?,
    shareRequest: ShareRequest? = null,
) {
    val context = LocalContext.current
    val scope = rememberCoroutineScope()
    var settings by remember(settingsVersion) { mutableStateOf(loadServerSettings(context)) }
    var screen by remember { mutableStateOf<Screen>(Screen.Main) }
    // A notification tap this could not follow, and why. Null both before one is asked for and
    // after one succeeds, since success is a screen rather than a message.
    var failedOpen by remember { mutableStateOf<FailedOpen?>(null) }
    // Bumped whenever another screen changes something the list shows, so returning to it
    // refetches.
    var reloadToken by remember { mutableIntStateOf(0) }
    // Cleared by the session screen that attached it, not when a newer request arrives: a share
    // must be attached exactly once, and only the screen that did it knows that it has.
    var share by remember { mutableStateOf<ShareRequest?>(null) }
    LaunchedEffect(shareRequest) {
        if (shareRequest != null) {
            share = shareRequest
            // A session already open takes it. Otherwise the list is where the choice is made,
            // whatever screen was showing: Spawn and Settings have nowhere to put a file.
            if (screen !is Screen.Session) screen = Screen.Main
        }
    }

    // A standing condition rather than a per-request failure, so it is stated once here instead of
    // appended to every error it might cause. Without this the app is simply unreachable and every
    // screen blames the server or the tunnel for it.
    if (!localNetworkAllowed(context)) {
        Text(
            "This app is not allowed to reach local network addresses, so it cannot " +
                "connect to the backend at all. Grant \"local network\" in Android's app " +
                "settings; until then every screen here will look like the server is down.",
            color = MaterialTheme.colorScheme.error,
            style = MaterialTheme.typography.bodySmall,
            modifier = Modifier.padding(16.dp),
        )
    }

    val current = settings
    if (current == null) {
        // Not enrolled yet: settings is the only usable screen. The QR path lands in MainActivity
        // and recomposes from the top.
        Box(Modifier.imePadding()) {
            SettingsScreen(
                existing = null,
                onSaved = { saved ->
                    settings = saved
                    screen = Screen.Main
                },
                onBack = null,
            )
        }
        return
    }

    // The one way back, whichever screen is showing and whether it was reached by the system back
    // gesture or a screen's own Back button. Every leaf screen can have changed something the list
    // shows, so it always refetches.
    val goToMain = {
        reloadToken++
        screen = Screen.Main
    }
    if (screen !is Screen.Main) {
        BackHandler(onBack = goToMain)
    }

    // Turning a notification into the screen it points at. The id has to be resolved to a session
    // first, because that is what SessionScreen is given -- and unlike a list row, there is nothing
    // here to seed it from.
    //
    // A failure is reported rather than swallowed: somebody deliberately tapped a notification, so
    // an app that opens to the session list with no explanation looks like the tap missed.
    val open: suspend (SessionOpenRequest) -> Unit = { request ->
        failedOpen = null
        try {
            val session = withContext(Dispatchers.IO) { fetchSession(current, request.sessionId) }
            screen = Screen.Session(session)
        } catch (e: ApiException) {
            failedOpen = FailedOpen(request, e.message ?: "Unknown error")
        }
    }
    LaunchedEffect(openRequest) { openRequest?.let { open(it) } }

    val failed = failedOpen
    if (failed != null) {
        AlertDialog(
            onDismissRequest = { failedOpen = null },
            title = { Text("Couldn't open that session") },
            text = { Text(failed.message) },
            confirmButton = {
                TextButton(onClick = { scope.launch { open(failed.request) } }) {
                    Text("Try again")
                }
            },
            dismissButton = { TextButton(onClick = { failedOpen = null }) { Text("Cancel") } },
        )
    }

    // Every screen but the session takes the keyboard as bottom padding here. The session screen
    // deliberately does not: resizing a whole screen on every frame of the keyboard animation is
    // the cost that made it lag, so it moves only its composer and transcript.
    when (val here = screen) {
        Screen.Main ->
            Box(Modifier.imePadding()) {
                MainScreen(
                    settings = current,
                    reloadToken = reloadToken,
                    share = share,
                    onOpen = { screen = Screen.Session(it) },
                    onSpawn = { screen = Screen.Spawn },
                    onImported = { imported ->
                        reloadToken++
                        screen = Screen.Session(imported)
                    },
                    onSettings = { screen = Screen.Settings },
                    onProvider = { machineId, provider ->
                        screen = Screen.ProviderSettings(machineId, provider)
                    },
                )
            }
        is Screen.Session ->
            // Keyed on the id, because a different session is a different screen rather than this
            // one showing other rows. SessionScreen remembers a transcript, an open stream, a draft
            // and a scroll position, and without the key Compose keeps all of it across the change
            // and merges two conversations -- which crashes the list on the first duplicate row
            // key. Only reachable since a notification can move straight from one session to
            // another.
            key(here.summary.id) {
                // A Box so the explorer can be drawn *over* the session rather than instead of it.
                // No imePadding here, for the reason above -- the explorer adds its own.
                Box {
                    val fileLinkHandler = rememberFileLinkHandler { path ->
                        screen = here.copy(files = here.summary.filesTarget(path))
                    }
                    Box(
                        Modifier.then(
                            if (here.subagent != null || here.files != null)
                                Modifier.clearAndSetSemantics {}
                            else Modifier
                        )
                    ) {
                        // How much background work the session has, from the one subscription
                        // to its events the screen below holds. Here because the panel and that
                        // screen both draw it, and must draw the same number.
                        var backgroundTasks by
                            remember(here.summary.id) {
                                mutableIntStateOf(here.summary.backgroundTasks)
                            }
                        // The two panels this session can be pulled aside for: its subagents
                        // from the right, and the whole main screen from the left. Both are here
                        // rather than screens of their own for the same reason the explorer is --
                        // the session under them stays composed. The main panel exists only
                        // inside a session, which is what makes it unswipeable until one has been
                        // opened.
                        SidePanels(
                            left = { active, close ->
                                MainPanel(
                                    settings = current,
                                    sessionId = here.summary.id,
                                    active = active,
                                    onOpen = { screen = Screen.Session(it) },
                                    onSpawn = { screen = Screen.Spawn },
                                    onImported = { imported ->
                                        reloadToken++
                                        screen = Screen.Session(imported)
                                    },
                                    onSettings = { screen = Screen.Settings },
                                    onProvider = { machineId, provider ->
                                        screen = Screen.ProviderSettings(machineId, provider)
                                    },
                                    onClose = close,
                                    onGone = goToMain,
                                )
                            },
                            // The whole width: it stands in for the screen Back would have shown,
                            // rather than sitting over the session the way the subagents do.
                            leftFraction = 1f,
                            right = { active, close ->
                                SubagentPanel(
                                    settings = current,
                                    summary = here.summary,
                                    active = active,
                                    backgroundTasks = backgroundTasks,
                                    onClose = close,
                                    onOpenSubagent = { screen = here.copy(subagent = it) },
                                )
                            },
                        ) {
                            CompositionLocalProvider(
                                LocalFileLinkHandler provides fileLinkHandler
                            ) {
                                SessionScreen(
                                    settings = current,
                                    summary = here.summary,
                                    onBack = goToMain,
                                    onFiles = { screen = here.copy(files = it) },
                                    share = share,
                                    onShareTaken = { share = null },
                                    onBackgroundTasks = { backgroundTasks = it },
                                )
                            }
                        }
                    }
                    here.subagent?.let { subagent ->
                        BackHandler { screen = here.copy(subagent = null) }
                        Box(
                            Modifier.fillMaxSize().background(MaterialTheme.colorScheme.background)
                        ) {
                            key(subagent.id) {
                                SessionScreen(
                                    settings = current,
                                    summary = here.summary,
                                    onBack = { screen = here.copy(subagent = null) },
                                    onFiles = {},
                                    subagent = subagent,
                                )
                            }
                        }
                    }
                    // Its own back handler is registered after this screen's, so it is the one the
                    // platform asks first, and it steps back inside itself before closing.
                    here.files?.let { target ->
                        FilesScreen(
                            settings = current,
                            target = target,
                            onClose = { screen = here.copy(files = null) },
                        )
                    }
                }
            }
        is Screen.ProviderSettings ->
            Box(Modifier.imePadding()) {
                ProviderScreen(
                    settings = current,
                    machineId = here.machineId,
                    provider = here.provider,
                    onBack = goToMain,
                )
            }
        is Screen.Spawn ->
            Box(Modifier.imePadding()) {
                SpawnScreen(
                    settings = current,
                    onSpawned = { spawned ->
                        reloadToken++
                        screen = Screen.Session(spawned)
                    },
                    onBack = goToMain,
                )
            }
        is Screen.Settings ->
            Box(Modifier.imePadding()) {
                SettingsScreen(
                    existing = current,
                    onSaved = { saved ->
                        settings = saved
                        goToMain()
                    },
                    onBack = goToMain,
                )
            }
    }

    // Last, so it draws over the screen above rather than under it: these are stacked in the Box
    // the activity puts around this, and that Box paints in the order it was given. A session
    // wanting attention is not a fact about the page somebody happens to be on. Tapping one is the
    // same act as tapping a notification, so it goes through the same `open`.
    SessionAlerts(onOpen = { request -> scope.launch { open(request) } })
}