app/backend/fluksio/tui/app.py

"""The screen itself: the engine under it, and what it is doing.

See `fluksio.tui` for why the engine is a child process rather than a thread.
"""

from __future__ import annotations

import argparse
import json
import os
import signal
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Any

from textual import work
from textual.app import App, ComposeResult
from textual.containers import Horizontal, Vertical, VerticalScroll
from textual.screen import ModalScreen
from textual.widgets import (
    Button,
    DataTable,
    Footer,
    Header,
    Input,
    Label,
    RichLog,
    Static,
    TabbedContent,
    TabPane,
)

from fluksio.cli import (
    DEFAULT_PORT,
    DEFAULT_PORTAL,
    _data_dir,
    _token_for,
    probe_engine,
    read_pidfile,
)
from fluksio.sdk import stream
from fluksio.sdk.cli import WATCH_INTERVAL_S, _dur, _status_screen
from fluksio.sdk.client import Client, config_path
from fluksio.tui.compare import CompareScreen

#: How long to keep asking the child for the credential it writes on the way
#: up, before giving up and letting the panels say so.
STARTUP_TRIES = 60

#: How many runs one comparison holds. The chart draws five of them
#: (`chart.MAX_SERIES`); the table below it carries the rest, which is the
#: split the browser makes for the same reason.
MAX_PICKED = 20

#: How many runs the table holds, now that it owns a whole tab.
RUNS_SHOWN = 50

#: What the dashboard subscribes to. `message_value` is most of what the bus
#: carries and none of what this screen draws.
KINDS = frozenset(
    {"run_started", "run_finished", "flow_paused", "engine_fatal", "flow_changed"}
)

#: How long an event waits for the ones behind it. A sweep starting twenty
#: runs should cost one read of the engine, not twenty.
COALESCE_S = 0.25

#: Where the engine's own output goes, beside the data it is serving. A file
#: rather than a pipe because the screen is meant to be closed while the
#: engine keeps running: a pipe whose reader has gone breaks the next write,
#: and a node's `print` is one of those writes. It also means the log of an
#: engine this screen only *adopted* can be read here.
LOG_NAME = "serve.log"

#: How much of it to read back when the screen opens.
LOG_TAIL_BYTES = 64 * 1024

#: What it is truncated to when this screen starts an engine of its own.
#: ponytail: a size check, not rotation — add rotation when somebody wants
#: yesterday's log.
LOG_KEEP_BYTES = 5 * 1024 * 1024

#: How long the tail waits when the file has nothing new. Also how long
#: closing the screen waits for that thread.
LOG_POLL_S = 0.25

#: How long a stop waits for the engine to be gone before killing it, and how
#: long it waits for an adopted one, which it can only ask.
STOP_WAIT_S = 10.0

#: The width of the runs table's five fixed columns, padding included, which
#: is what is left over for the inputs.
FIXED_COLUMNS = 54


def child_argv(argv: list[str]) -> list[str]:
    """The same command, told to serve without a dashboard of its own.

    Every flag is passed through: whatever `serve` was asked for is what the
    engine under this screen is running with.
    """
    passed = [flag for flag in argv if flag != "--plain"]
    return [sys.executable, "-m", "fluksio.cli", *passed, "--plain"]


class Enroll(ModalScreen[None]):
    """The claim code a portal minted, and which portal minted it.

    The work happens here rather than back on the dashboard: enrolment is a
    round trip to the portal, and the person who pressed the button is looking
    at this modal while it happens.
    """

    BINDINGS = [("escape", "dismiss(None)", "cancel")]

    def __init__(self, data_dir: Path) -> None:
        super().__init__()
        self.data_dir = data_dir

    def compose(self) -> ComposeResult:
        with Vertical(id="enroll"):
            yield Label("Pair this instance with a portal")
            yield Input(placeholder="claim code", id="code")
            yield Input(value=DEFAULT_PORTAL, id="portal")
            yield Label("", id="note")
            with Horizontal():
                yield Button("Enroll", variant="primary", id="go")
                yield Button("Cancel", id="cancel")

    def on_mount(self) -> None:
        self.query_one("#code", Input).focus()

    def on_button_pressed(self, event: Button.Pressed) -> None:
        if event.button.id == "cancel":
            self.dismiss(None)
            return
        self.submit()

    def on_input_submitted(self, event: Input.Submitted) -> None:
        self.submit()

    def submit(self) -> None:
        code = self.query_one("#code", Input).value.strip()
        portal = self.query_one("#portal", Input).value.strip() or DEFAULT_PORTAL
        if not code:
            self.note("A claim code is what pairs this instance.")
            return
        self.busy(True)
        self.note(f"Enrolling with {portal}…")
        self.enroll(code, portal)

    def note(self, message: str) -> None:
        self.query_one("#note", Label).update(message)

    def busy(self, working: bool) -> None:
        for one in self.query(Input):
            one.disabled = working
        for button in self.query(Button):
            button.disabled = working

    @work(thread=True)
    def enroll(self, code: str, portal: str) -> None:
        """`fluksio enroll`, as its own process for the same reason serve is.

        A running engine picks the configuration up on its own; this screen
        only has to report what the command said.
        """
        done = subprocess.run(  # noqa: S603
            [
                sys.executable,
                "-m",
                "fluksio.cli",
                "enroll",
                code,
                "--portal",
                portal,
                "--data-dir",
                str(self.data_dir),
            ],
            capture_output=True,
            text=True,
        )
        self.app.call_from_thread(
            self.finished, done.returncode, done.stdout + done.stderr
        )

    def finished(self, code: int, output: str) -> None:
        lines = [line for line in output.splitlines() if line.strip()]
        for line in lines:
            self.app.note(line)  # type: ignore[attr-defined]
        said = lines[-1] if lines else ""
        if code == 0:
            self.app.notify(said or "Connected to the portal.")
            self.dismiss(None)
            return
        self.busy(False)
        self.note(said or "Enrolment failed; the log has what it said.")


class Artifacts(ModalScreen[None]):
    """What a run left behind, and one key to fetch it.

    The rows are already on `GET /runs/{id}` — the same list the browser draws
    on a run's page — so this is a view of them rather than a second reading.
    """

    BINDINGS = [
        ("escape", "dismiss(None)", "close"),
        ("enter", "download", "download"),
    ]

    def __init__(self, client: Client, run_id: str) -> None:
        super().__init__()
        self.client = client
        self.run_id = run_id
        self.rows: list[dict[str, Any]] = []

    def compose(self) -> ComposeResult:
        with Vertical(id="artifacts"):
            yield Label(f"Artifacts of {self.run_id[-8:]}")
            table: DataTable[str] = DataTable(id="files", cursor_type="row")
            yield table
            yield Label("", id="note")

    def on_mount(self) -> None:
        table = self.query_one("#files", DataTable)
        table.add_columns("name", "node", "type", "bytes")
        table.focus()
        self.read()

    @work(thread=True)
    def read(self) -> None:
        try:
            detail = self.client.run(self.run_id)
        except Exception as exc:  # noqa: BLE001 — the modal reports it
            self.app.call_from_thread(self.note, f"Could not read the run: {exc}")
            return
        self.app.call_from_thread(self.show, detail.get("artifacts") or [])

    def show(self, rows: list[dict[str, Any]]) -> None:
        self.rows = rows
        table = self.query_one("#files", DataTable)
        for row in rows:
            table.add_row(
                str(row.get("name", "")),
                str(row.get("node", "")),
                str(row.get("media_type", "")),
                f"{int(row.get('size') or 0):,}",
                key=str(row.get("digest", "")),
            )
        self.note("This run produced no artifacts." if not rows else "Enter downloads.")

    def note(self, message: str) -> None:
        self.query_one("#note", Label).update(message)

    def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
        self.action_download()

    def action_download(self) -> None:
        table = self.query_one("#files", DataTable)
        if not table.row_count:
            return
        digest = table.coordinate_to_cell_key(table.cursor_coordinate).row_key.value
        row = next((one for one in self.rows if one.get("digest") == digest), None)
        if row is not None:
            self.fetch(row)

    @work(thread=True)
    def fetch(self, row: dict[str, Any]) -> None:
        # The name it was written under reads better than the message's, which
        # is chosen for the graph — the same choice `fluksio artifacts` makes.
        out = Path(str(row.get("filename") or row.get("name")))
        try:
            out.write_bytes(self.client.download(str(row["digest"])))
        except Exception as exc:  # noqa: BLE001 — the modal reports it
            self.app.call_from_thread(self.note, f"Could not fetch it: {exc}")
            return
        self.app.call_from_thread(self.note, f"Wrote {out.resolve()}")


class RunsTable(DataTable[str]):
    """The runs, and the keys that only mean anything while they are in front.

    Textual resolves a binding from the focused widget outwards, so hanging
    them here is what makes the footer read as run tools on this tab and as
    engine tools on the others.
    """

    BINDINGS = [
        ("space", "app.pick", "pick for comparison"),
        ("c", "app.cancel_run", "cancel run"),
        ("a", "app.artifacts", "artifacts"),
    ]


class ServeApp(App[int]):
    """Three tabs: how the engine is, what has run, and what it is saying."""

    CSS = """
    TabbedContent { height: 1fr; }
    TabPane { height: 1fr; padding: 0; }
    #overview { padding: 0 1; }
    #enroll { width: 60; height: auto; padding: 1 2; background: $surface; }
    #enroll #note { color: $text-muted; height: auto; }
    #artifacts { width: 80; height: auto; padding: 1 2; background: $surface; }
    #artifacts DataTable { height: auto; max-height: 14; }
    """

    BINDINGS = [
        ("q", "quit", "quit (engine keeps running)"),
        ("1", "tab('overview-tab')", "overview"),
        ("2", "tab('runs-tab')", "runs"),
        ("3", "tab('logs-tab')", "logs"),
        ("s", "stop_start", "stop/start"),
        ("r", "restart", "restart"),
        ("e", "enroll", "enroll"),
    ]

    def __init__(self, args: argparse.Namespace) -> None:
        super().__init__()
        self.args = args
        self.data_dir: Path = _data_dir(args.data_dir, args.shared)
        self.child: subprocess.Popen[bytes] | None = None
        #: The pid of an engine this dashboard did not start. Only ever one
        #: whose data directory is this one — the probe is what proves it.
        self.adopted: int | None = None
        self.client: Client | None = None
        self.url = ""
        #: The runs ticked for a comparison, in the order they were ticked.
        self.picked: list[str] = []
        #: Set by the socket thread, drained by the coalescing timer — so a
        #: cascade of events costs one read of the engine rather than one each.
        self.stirred = False
        self.stop_stream = threading.Event()
        self.stop_log = threading.Event()
        #: The credential's mtime when the engine under this screen was
        #: started, so its own is told from the one before it.
        self.config_stamp = 0

    # -- layout ---------------------------------------------------------------

    def compose(self) -> ComposeResult:
        yield Header()
        with TabbedContent(initial="overview-tab"):
            with TabPane("Overview", id="overview-tab"):
                with VerticalScroll():
                    yield Static(id="overview")
            with TabPane("Runs", id="runs-tab"):
                yield RunsTable(id="runs", cursor_type="row")
            with TabPane("Logs", id="logs-tab"):
                yield RichLog(id="log", markup=False, highlight=False, max_lines=5000)
        yield Footer()

    def on_mount(self) -> None:
        self.title = f"fluksio — {self.data_dir}"
        table = self.query_one("#runs", RunsTable)
        table.add_column(" ", width=1)
        table.add_column("run", width=8)
        table.add_column("status", width=9)
        table.add_column("flow", width=16)
        table.add_column("took", width=8)
        table.add_column("params")
        self.tail_log()
        self.start_engine()
        # The socket is what makes a run that starts and finishes between two
        # ticks visible; this stays as the heartbeat for what is not on the bus.
        self.set_interval(WATCH_INTERVAL_S, self.refresh_panels)
        self.set_interval(COALESCE_S, self.drain)

    def action_tab(self, tab: str) -> None:
        self.query_one(TabbedContent).active = tab

    def on_tabbed_content_tab_activated(
        self, event: TabbedContent.TabActivated
    ) -> None:
        """Focus what the tab is about, so its keys are the ones in the footer."""
        focusable = event.pane.query("RunsTable, RichLog, VerticalScroll")
        if focusable:
            focusable.first().focus()

    async def action_quit(self) -> None:
        """Close the screen and leave the engine running.

        The threads are told first: both are worker threads on the loop's own
        executor, which is joined while the loop closes — so a reader still
        blocked on the log or the socket would hold the process after the
        screen is gone, which is what `q` used to do.
        """
        self.stop_log.set()
        self.stop_stream.set()
        self.exit(0)

    def on_unmount(self) -> None:
        self.stop_log.set()
        self.stop_stream.set()

    # -- the engine under the screen ------------------------------------------

    def note(self, message: str) -> None:
        self.query_one("#log", RichLog).write(message)

    def say(self, message: str) -> None:
        """`note`, from a worker thread."""
        self.call_from_thread(self.note, message)

    def log_path(self) -> Path:
        return self.data_dir / LOG_NAME

    def start_engine(self) -> None:
        """Adopt whatever is already serving this directory, or start one."""
        host = self.args.host
        reachable = "127.0.0.1" if host in ("0.0.0.0", "::", "") else host
        # Where this directory's engine last said it was, which is not always
        # where it was asked to be: a taken port moves.
        running = read_pidfile(self.data_dir)
        wanted = running["port"] if running else (self.args.port or DEFAULT_PORT)
        url = f"http://{reachable}:{wanted}"
        who = probe_engine(url, _token_for(self.data_dir))

        if who == "ours":
            self.adopted = running["pid"] if running else None
            self.url = url
            named = f" (pid {self.adopted})" if self.adopted else ""
            self.note(f"Adopted the engine already serving this directory{named}.")
            if self.adopted is None:
                self.note("It wrote no pidfile, so this screen cannot stop it.")
            self.connect()
            return
        if who == "foreign":
            self.note(f"Port {wanted} holds another instance's Fluksio.")

        # What the credential looked like before the child wrote its own, so
        # the wait below cannot mistake the last engine's url for this one's.
        self.config_stamp = self._config_stamp()
        path = self.log_path()
        if path.exists() and path.stat().st_size > LOG_KEEP_BYTES:
            path.write_text("")
        handle = path.open("ab", buffering=0)
        try:
            self.child = subprocess.Popen(  # noqa: S603
                child_argv(sys.argv[1:]),
                stdout=handle,
                stderr=subprocess.STDOUT,
                env={**os.environ, "PYTHONUNBUFFERED": "1"},
            )
        finally:
            # The child holds a descriptor of its own; this one is what would
            # otherwise keep the file open for as long as the screen lives.
            handle.close()
        self.await_engine()

    @work(thread=True, exclusive=True, group="log")
    def tail_log(self) -> None:
        """The engine's own output, which is why a second terminal was needed."""
        path = self.log_path()
        while not self.stop_log.is_set() and self.is_running:
            try:
                with path.open("r", errors="replace") as handle:
                    handle.seek(max(0, path.stat().st_size - LOG_TAIL_BYTES))
                    if handle.tell():
                        handle.readline()  # whatever line the seek landed in
                    while not self.stop_log.is_set() and self.is_running:
                        line = handle.readline()
                        if line:
                            self.call_from_thread(self.note, line.rstrip())
                            continue
                        if path.stat().st_size < handle.tell():
                            break  # truncated under us; read it from the top
                        time.sleep(LOG_POLL_S)
            except OSError:
                # Not written yet, or gone. Either way it may appear.
                time.sleep(LOG_POLL_S)

    @work(thread=True, exclusive=True, group="startup")
    def await_engine(self) -> None:
        """Wait for the child to say where it landed, then talk to it.

        The port is the child's to choose — it moves off a taken one — and
        `client.json` is where it says which it took.
        """
        for _ in range(STARTUP_TRIES):
            if self.child is not None and self.child.poll() is not None:
                return
            try:
                stored = json.loads(config_path(self.data_dir).read_text())
            except (OSError, ValueError):
                stored = {}
            if self._config_stamp() == self.config_stamp:
                # Still the one the last engine left, which names a port this
                # one may not have taken.
                time.sleep(0.5)
                continue
            if stored.get("url") and stored.get("token"):
                self.url = str(stored["url"])
                self.call_from_thread(self.connect)
                return
            time.sleep(0.5)
        self.call_from_thread(self.note, "The engine did not come up.")

    def connect(self) -> None:
        stored = config_path(self.data_dir)
        try:
            config = json.loads(stored.read_text())
            self.client = Client(
                url=str(config["url"]), token=str(config["token"]), retries=0
            )
        except (OSError, ValueError, KeyError) as exc:
            self.note(f"No credential to talk to the engine with: {exc}")
            return
        self.url = self.client.url
        self.stop_stream = threading.Event()
        self.watch_engine(self.client.url, self.client.token)
        self.refresh_panels()

    def _config_stamp(self) -> int:
        """When the credential was last written, or 0 if it never was."""
        try:
            return config_path(self.data_dir).stat().st_mtime_ns
        except OSError:
            return 0

    def engine_pid(self) -> int | None:
        return self.child.pid if self.child is not None else self.adopted

    def stop_engine(self) -> None:
        """Stop the engine and wait for it to be gone. **From a thread.**

        Waiting is what makes `r` work: a new engine started while the old one
        still holds the port would move off it, and the screen would end up
        watching one engine while the client talks to another. An engine
        draining its flows takes as long as it takes, which is why this is not
        something to do on the loop that draws.
        """
        pid = self.engine_pid()
        if self.child is not None:
            self.say(f"Stopping the engine (pid {self.child.pid}).")
            self.child.terminate()
            try:
                self.child.wait(timeout=STOP_WAIT_S)
            except subprocess.TimeoutExpired:
                self.child.kill()
            self.child = None
        elif self.adopted is not None:
            self.say(f"Stopping the adopted engine (pid {self.adopted}).")
            try:
                os.kill(self.adopted, signal.SIGTERM)
            except OSError as exc:
                self.say(f"Could not stop it: {exc}")
            self._await_exit(self.adopted)
            self.adopted = None
        self.stop_stream.set()
        self.client = None
        if pid is not None:
            self.say("The engine stopped.")

    def _await_exit(self, pid: int) -> None:
        deadline = time.time() + STOP_WAIT_S
        while time.time() < deadline:
            try:
                os.kill(pid, 0)
            except OSError:
                return
            time.sleep(0.2)
        self.say(f"pid {pid} is still running.")

    # -- the engine's own events ----------------------------------------------

    @work(thread=True, exclusive=True, group="stream")
    def watch_engine(self, url: str, token: str) -> None:
        """The bus, over the websocket the browser subscribes to."""
        stream.subscribe(
            url,
            token,
            lambda event: setattr(self, "stirred", True),
            self.stop_stream,
            kinds=KINDS,
            on_error=lambda exc: self.call_from_thread(
                self.note, f"Not receiving events: {exc}"
            ),
        )

    def drain(self) -> None:
        """Whatever the bus said since the last tick, as one read."""
        if not self.stirred:
            return
        self.stirred = False
        self.refresh_panels()
        # A run finishing is exactly when a comparison's numbers become final,
        # and the screen showing them cannot hear the bus itself.
        if isinstance(self.screen, CompareScreen):
            self.screen.action_reload()

    # -- what the panels show -------------------------------------------------

    def refresh_panels(self) -> None:
        if self.client is not None:
            self.read_engine()

    @work(thread=True, exclusive=True, group="poll")
    def read_engine(self) -> None:
        client = self.client
        if client is None:
            return
        try:
            screen = _status_screen(client)
            rows = client.runs(limit=RUNS_SHOWN)
        except Exception as exc:
            self.call_from_thread(self.show_offline, exc)
            return
        self.call_from_thread(self.show, screen, rows)

    def show(self, screen: Any, rows: list[dict[str, Any]]) -> None:
        self.query_one("#overview", Static).update(screen)
        table = self.query_one("#runs", RunsTable)
        cursor = table.cursor_row
        # Whatever the fixed columns do not use, so a wide terminal shows the
        # inputs rather than a stripe of empty table.
        room = max(20, table.size.width - FIXED_COLUMNS)
        table.clear()
        for row in rows:
            run_id = str(row.get("id", ""))
            params = json.dumps(row.get("params") or {})
            table.add_row(
                "·" if run_id in self.picked else " ",
                run_id[-8:],
                str(row.get("status", "")),
                str(row.get("flow", "")),
                _dur(row.get("duration_ms")),
                params.ljust(room) if len(params) <= room else params[: room - 1] + "…",
                key=run_id,
            )
        # The rows are rewritten wholesale every refresh, and a cursor that
        # jumped back to the top on each one would make picking three runs a
        # race against the clock.
        if 0 <= cursor < table.row_count:
            table.move_cursor(row=cursor)

    def show_offline(self, exc: Exception) -> None:
        which = "starting" if self.engine_pid() is not None else "not running"
        self.query_one("#overview", Static).update(
            f"[yellow]The engine is {which}.[/] ({type(exc).__name__})"
        )

    # -- keys -----------------------------------------------------------------

    @work(thread=True, group="engine")
    def action_stop_start(self) -> None:
        if self.engine_pid() is not None:
            self.stop_engine()
        else:
            self.call_from_thread(self.start_engine)

    @work(thread=True, group="engine")
    def action_restart(self) -> None:
        """Stop, wait for the port to come back, start."""
        self.stop_engine()
        self.call_from_thread(self.start_engine)

    def cursor_run(self) -> str:
        """The whole id of the run the cursor is on.

        The cell holds its tail, which is what fits in a column; the row's key
        is the whole of it, which is what the engine is asked about.
        """
        table = self.query_one("#runs", RunsTable)
        if self.client is None or not table.row_count:
            return ""
        key = table.coordinate_to_cell_key(table.cursor_coordinate).row_key
        return str(key.value or "")

    def action_cancel_run(self) -> None:
        run_id = self.cursor_run()
        if run_id:
            self.cancel_run(run_id)

    def action_pick(self) -> None:
        """Tick the run under the cursor into the comparison, or untick it."""
        run_id = self.cursor_run()
        if not run_id:
            return
        if run_id in self.picked:
            self.picked.remove(run_id)
        elif len(self.picked) >= MAX_PICKED:
            self.note(f"A comparison carries {MAX_PICKED} runs; untick one first.")
            return
        else:
            self.picked.append(run_id)
        self.refresh_panels()

    def action_compare(self) -> None:
        """The ticked runs, or the one under the cursor if none are ticked."""
        if self.client is None:
            return
        ids = self.picked or [self.cursor_run()]
        if any(ids):
            self.push_screen(CompareScreen(self.client, [one for one in ids if one]))

    def on_data_table_row_selected(self, event: DataTable.RowSelected) -> None:
        """Enter on the focused runs table, which never reaches the binding."""
        if event.data_table.id == "runs":
            self.action_compare()

    def action_artifacts(self) -> None:
        run_id = self.cursor_run()
        if self.client is not None and run_id:
            self.push_screen(Artifacts(self.client, run_id))

    @work(thread=True)
    def cancel_run(self, run_id: str) -> None:
        client = self.client
        if client is None:
            return
        try:
            client.cancel(run_id)
        except Exception as exc:
            self.call_from_thread(self.note, f"Could not cancel {run_id}: {exc}")
            return
        self.call_from_thread(self.note, f"Cancelled {run_id}.")
        self.call_from_thread(self.refresh_panels)

    def action_enroll(self) -> None:
        self.push_screen(Enroll(self.data_dir))


def run_tui(args: argparse.Namespace) -> int:
    """Open the dashboard, and say what is still running when it closes."""
    app = ServeApp(args)
    app.run()
    pid, url = app.engine_pid(), app.url
    if pid is not None:
        print(f"The engine is still running: pid {pid} at {url or 'its port'}.")
        print(f"  fluksio serve reattaches to it; kill {pid} stops it.")
        print(f"  Its log: {app.log_path()}")
    return 0