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.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.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 setups 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 {
    /**
     * The session list, with a subagent's own transcript over it when [subagent] is set.
     *
     * A layer on this screen rather than a screen of its own, for the same reason [Session.files]
     * is: [SessionListScreen] owns which cards are expanded and what each expansion fetched, kept
     * in `remember`, and a subagent is opened from a card's expander. As a sibling `Screen` it was
     * disposed and recreated on every return, which lost that state -- an expanded card collapsed
     * itself the moment its own subagent's view was closed.
     */
    data class Main(val subagent: SubagentTarget? = null) : Screen()

    /**
     * One subagent's own transcript, read-only. See [SessionScreen]'s `subagent` parameter and
     * SUBAGENTS.md's "Phone". Closing it returns to [Main] under it, not to [Session]: a subagent
     * is opened from the session list's card rather than from inside the session it belongs to.
     */
    data class SubagentTarget(val summary: SessionSummary, val subagent: SubagentSummary)

    /**
     * One session, with the file explorer over it when [files] is set.
     *
     * The explorer is a layer on this screen rather than a screen of its own, so the session under
     * it stays composed: its event stream keeps flowing, its scroll position and draft stay put,
     * and coming back from a file costs nothing. As a sibling `Screen` it would be disposed and re-
     * created on every return, refetching the transcript over the tunnel.
     */
    data class Session(val summary: SessionSummary, val files: FilesTarget? = null) : Screen()

    data object Spawn : 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) {
        is Screen.Main ->
            Box(Modifier.imePadding()) {
                MainScreen(
                    settings = current,
                    reloadToken = reloadToken,
                    share = share,
                    onOpen = { screen = Screen.Session(it) },
                    onOpenSubagent = { summary, subagent ->
                        screen = here.copy(subagent = Screen.SubagentTarget(summary, subagent))
                    },
                    onSpawn = { screen = Screen.Spawn },
                    onImported = { imported ->
                        reloadToken++
                        screen = Screen.Session(imported)
                    },
                    onSettings = { screen = Screen.Settings },
                )
                // Its own back handler is registered after MainScreen's, so it is the one the
                // platform asks first while a subagent is open -- the same rule the files
                // explorer's handler follows over its session, below.
                here.subagent?.let { target ->
                    BackHandler { screen = here.copy(subagent = null) }
                    // Its own opaque background: this screen was always the sole content under
                    // the theme's own Surface before, so it never had to paint one -- stacked over
                    // the list here, the space between its own cards let the list underneath show
                    // through without this. The same fix FilesScreen needed over its session.
                    Box(Modifier.fillMaxSize().background(MaterialTheme.colorScheme.background)) {
                        key(target.summary.id, target.subagent.id) {
                            SessionScreen(
                                settings = current,
                                summary = target.summary,
                                onBack = { screen = here.copy(subagent = null) },
                                onFiles = {},
                                subagent = target.subagent,
                            )
                        }
                    }
                }
            }
        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 {
                    SessionScreen(
                        settings = current,
                        summary = here.summary,
                        onBack = goToMain,
                        onFiles = { screen = here.copy(files = it) },
                        share = share,
                        onShareTaken = { share = null },
                    )
                    // 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.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) } })
}