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

package com.example.aiapp

import androidx.compose.foundation.layout.Row
import androidx.compose.foundation.layout.fillMaxWidth
import androidx.compose.foundation.layout.padding
import androidx.compose.material3.LinearProgressIndicator
import androidx.compose.material3.MaterialTheme
import androidx.compose.material3.Text
import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableIntStateOf
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.setValue
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.graphics.Color
import androidx.compose.ui.unit.dp
import java.time.Duration
import java.time.OffsetDateTime
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.delay
import kotlinx.coroutines.withContext

/** What one machine's rate limits came back as, or why they didn't. */
sealed class SessionUsage {
    /** Nothing has come back yet. Distinct from every answer, including an empty one. */
    data object Waiting : SessionUsage()

    /** Every window the machine reported, in the order it reported them. */
    data class Known(val windows: List<UsageWindow>) : SessionUsage()

    /**
     * This machine meters nothing, so there is no window to show.
     *
     * Separate from [Unavailable], and the distinction is the point: a session on `echo` or on a
     * local llama.cpp has no paid quota at all, which is a fact about how it was set up and not a
     * failure to find something out. The backend never asks such a machine, and reading that
     * silence as "couldn't find out" is answering with the nearest available word.
     */
    data object NotMetered : SessionUsage()

    /**
     * The question could not be answered, and why.
     *
     * Its own state because "we couldn't find out" and "none of it is used" must never share an
     * appearance: a bar sitting at zero because a machine is unreachable reads as plenty of
     * headroom.
     */
    data class Unavailable(val why: String) : SessionUsage()
}

/** How often to ask again. The backend caches, so this re-reads its cache rather than the API. */
private const val REFRESH_MS = 60_000L

/**
 * One poll of every machine's limits, and the handle to ask again.
 *
 * A screen shows this answer in more than one place -- the bar under the session header, the colour
 * of the button beside it, and the dialog that button opens -- and each used to fetch for itself.
 * Two fetches say one thing twice and then disagree: the bar's copy can be a whole refresh interval
 * old when the dialog opens with a fresh one, so the header read 42% while the screen over it read
 * 47%.
 */
class UsageFeed(
    val snapshots: LoadState<List<UsageSnapshot>>,
    /** A fetch is outstanding. Only ever true over an answer already shown. */
    val refreshing: Boolean,
    /** Ask the backend again now. The dialog's refresh button; the poll does it on its own. */
    val refresh: () -> Unit,
) {
    /** What [setup]'s own limits came back as. See [usageFor] for why the states are these. */
    fun forSetup(setup: String): SessionUsage =
        when (val state = snapshots) {
            is LoadState.Loading -> SessionUsage.Waiting
            is LoadState.Error -> SessionUsage.Unavailable(state.message)
            is LoadState.Loaded -> usageFor(state.value, setup)
        }
}

/**
 * The one poll of the machines' rate limits, polled and refreshable.
 *
 * Hoisted out of [SessionUsageBar] because everything on a session's screen that reports on usage
 * has to be reporting the same measurement; see [UsageFeed].
 */
@Composable
fun rememberUsageFeed(settings: ServerSettings): UsageFeed {
    var snapshots by remember { mutableStateOf<LoadState<List<UsageSnapshot>>>(LoadState.Loading) }
    var refreshing by remember { mutableStateOf(true) }
    // Bumped to ask again now. The poll below restarts from the new value, so a manual refresh also
    // resets the countdown rather than leaving one due immediately after.
    var asked by remember { mutableIntStateOf(0) }
    LaunchedEffect(asked) {
        while (true) {
            refreshing = true
            // Replaces the answer only once the next one is in hand: dropping back to Loading would
            // blank a bar somebody is reading for the length of a round trip, and what was on
            // screen is still the last thing the machine actually said.
            snapshots =
                try {
                    LoadState.Loaded(withContext(Dispatchers.IO) { fetchUsage(settings) })
                } catch (e: ApiException) {
                    LoadState.failed(e)
                }
            refreshing = false
            delay(REFRESH_MS)
        }
    }
    return remember(snapshots, refreshing) { UsageFeed(snapshots, refreshing) { asked++ } }
}

/**
 * The colour for a control that reports on [usage] as a whole: the worst window's.
 *
 * Worst rather than the five-hour one, because the button it colours opens *all* of them, and a
 * blue icon over a weekly quota at 97% would be the interface answering a question nobody asked.
 * Taken over however many windows came back rather than the three Claude sends today -- the backend
 * passes windows it does not recognise straight through.
 *
 * Every state that is not a measurement takes the ordinary control colour instead. That is the
 * point where colour stops being able to help: blue is the low end of a scale here, so colouring an
 * unknown blue would say "measured, and fine" about a machine nobody could reach.
 */
@Composable
fun usageGlyphColour(usage: SessionUsage): Color =
    when (usage) {
        is SessionUsage.Known ->
            usage.windows.maxOfOrNull { it.percent }?.let { quotaColor(it) }
                ?: MaterialTheme.colorScheme.primary
        else -> MaterialTheme.colorScheme.primary
    }

/**
 * The five-hour window for the machine this session runs on, under the session's own header.
 *
 * Here rather than only in the usage dialog because it is the number that decides whether to keep
 * going, and it was a screen away from the place that decision gets made. It reports on this
 * session's machine alone -- the dialog is still where every machine is compared.
 *
 * What it shows is the paid service's own metering, never derived from what this app has watched go
 * past: the transcript's token counts are a different quantity, measured differently, and a bar
 * built out of them would be a guess wearing a measurement's clothes.
 */
@Composable
fun SessionUsageBar(usage: SessionUsage, modifier: Modifier = Modifier) {
    DebugStats.count("usage bar recomposed")
    // The countdown moves even when the numbers do not, so it is driven by a clock of its own
    // rather than recomputed at draw time: a percentage that comes back unchanged is an equal
    // value, Compose skips the recomposition, and a "left" that only ticked when the quota moved
    // would sit at a stale figure for hours.
    var now by remember { mutableStateOf(OffsetDateTime.now()) }
    LaunchedEffect(Unit) {
        while (true) {
            delay(REFRESH_MS)
            now = OffsetDateTime.now()
        }
    }

    // Nothing at all for a machine that meters nothing: a row saying "unknown" there would report a
    // problem about a setup somebody chose, on every screen, forever.
    if (usage is SessionUsage.NotMetered) {
        return
    }

    Row(
        verticalAlignment = Alignment.CenterVertically,
        modifier = modifier.fillMaxWidth().padding(horizontal = 16.dp, vertical = 2.dp),
    ) {
        // Words, not a colour and not an empty bar: every one of these is a different kind of
        // answer from "this much is used", and only words carry a difference in kind.
        when (val state = usage) {
            SessionUsage.NotMetered -> Unit
            is SessionUsage.Unavailable -> UsageNote("5-hour usage unknown -- ${state.why}")
            SessionUsage.Waiting -> UsageNote("5-hour usage: checking")
            is SessionUsage.Known -> {
                val window = state.windows.firstOrNull { it.kind == "session" }
                if (window == null) {
                    UsageNote("5-hour usage unknown -- no five-hour window reported")
                } else {
                    LinearProgressIndicator(
                        progress = { (window.percent / 100.0).toFloat().coerceIn(0f, 1f) },
                        // The same step at the same percentages as the dialog's bars: this is the
                        // same measurement, and a reader who learned the colour there has to be
                        // able to read it here without checking which screen they are on.
                        color = quotaColor(window.percent),
                        modifier = Modifier.weight(1f),
                    )
                    Text(
                        fiveHourLabel(window, now),
                        style = MaterialTheme.typography.labelSmall,
                        color = MaterialTheme.colorScheme.onSurfaceVariant,
                        modifier = Modifier.padding(start = 8.dp),
                    )
                }
            }
        }
    }
}

/** Anything this row says instead of drawing a bar, so all of them look the same. */
@Composable
private fun UsageNote(text: String) {
    Text(
        text,
        style = MaterialTheme.typography.labelSmall,
        color = MaterialTheme.colorScheme.onSurfaceVariant,
    )
}

/**
 * "42% -- 2h 15m left": how much is gone, then how long what is left has to last.
 *
 * The percentage on its own does not answer the question it gets asked, which is whether to start
 * something now; 80% with twenty minutes to go and 80% with four hours to go are opposite answers.
 *
 * The window's end has two missing cases, worded differently on purpose; see [WindowEnd]. A window
 * that is not running gets the percentage and nothing else.
 */
private fun fiveHourLabel(window: UsageWindow, now: OffsetDateTime): String {
    val percent = "${window.percent.toInt()}%"
    return when (val end = windowEnd(window.resetsAt, now)) {
        // Between blocks the five-hour window has no reset time, and saying so is a fact about
        // nothing: there is no window to run out. The percentage is the whole answer.
        WindowEnd.NotRunning -> percent
        WindowEnd.Unreadable -> "$percent · reset time unreadable"
        is WindowEnd.Ends ->
            // Under a minute, including past the end: the number would round to "0m left", which
            // reads as a measurement rather than as the window having run out.
            if (end.until < Duration.ofMinutes(1)) "$percent · refresh soon"
            else "$percent · ${formatSpan(end.until)} left"
    }
}

/**
 * One machine's snapshot, out of every machine's.
 *
 * Every way of having *failed* to get numbers is [SessionUsage.Unavailable] with the reason in it.
 * None of them may look like zero, and none may look like [SessionUsage.NotMetered], which is the
 * machine having no quota rather than the question going unanswered.
 */
fun usageFor(snapshots: List<UsageSnapshot>, setup: String): SessionUsage {
    // No snapshot at all means the backend never asked, which it only does for a machine with
    // nothing metered on it. That is a different answer from having asked and failed.
    val mine = snapshots.firstOrNull { it.setup == setup } ?: return SessionUsage.NotMetered
    if (mine.state != "ok") {
        return SessionUsage.Unavailable(mine.detail ?: mine.state)
    }
    return SessionUsage.Known(mine.windows)
}