path: root/packages/meshbay-cli/src/meshbay_cli/tui/screens.py

"""
The sign-in screen, and the main screen: two panes and the transfers.

F5 copies the selection to the other pane's folder — a download or an upload,
with the rules `get` and `put` follow: an upload is decided whole before it
starts (names the node refuses stop it, names already there are asked about
once for the batch), nothing on either side is replaced without being asked,
and a name from a group is made writable here before it is written.
"""

import asyncio
import contextlib
import threading
from pathlib import Path

from rich.text import Text
from textual.app import ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal, Vertical
from textual.screen import Screen
from textual.widgets import Button, Footer, Input, Static

from ..client.account import Account
from ..client.names import UnsafeName, local_path
from ..client.paths import handle
from ..client.pool import Pool
from ..client.store import FILE_KEY_NOTICE, KEY_FILE
from ..client.transfers import PART_SUFFIX, TransferError, create_dir, download, upload
from ..client.tree import Tree
from ..client.upload_plan import plan_upload, refused_names, taken_names
from ..errors import CliError
from ..output import safe
from .panes import LocalPane, RemotePane
from .widgets import Choice, TransferStrip

PER_GROUP = 2


class SignInScreen(Screen):
    DEFAULT_CSS = """
    SignInScreen { align: center middle; }
    SignInScreen > Vertical { width: 64; height: auto; border: thick $accent;
                              padding: 1 2; background: $surface; }
    SignInScreen Input { margin-bottom: 1; }
    SignInScreen #error { color: $error; height: auto; }
    """

    def compose(self) -> ComposeResult:
        ctx = self.app.ctx
        with Vertical():
            yield Static(Text("Sign in to MeshBay"))
            yield Input(value=ctx.hub_url, placeholder="hub", id="hub")
            yield Input(value=ctx.username or ctx.stored_username or "",
                        placeholder="username", id="user")
            yield Input(placeholder="passphrase", password=True, id="pass")
            yield Button("Sign in", id="go", variant="primary")
            yield Static("", id="error", markup=False)

    def on_mount(self) -> None:
        empty = next((i for i in self.query(Input) if not i.value), None)
        (empty or self.query_one("#pass", Input)).focus()

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

    def on_button_pressed(self, event: Button.Pressed) -> None:
        self._submit()

    def _submit(self) -> None:
        hub = self.query_one("#hub", Input).value.strip().rstrip("/")
        user = self.query_one("#user", Input).value.strip()
        passphrase = self.query_one("#pass", Input).value
        if not (hub and user and passphrase):
            self.query_one("#error", Static).update("Hub, username and passphrase.")
            return
        self.query_one("#error", Static).update("Signing in…")
        self.run_worker(self._sign_in(hub, user, passphrase), exclusive=True)

    async def _sign_in(self, hub_url: str, user: str, passphrase: str) -> None:
        app = self.app
        try:
            await app.use_hub(hub_url)
            account = Account(app.hub, user, lambda: passphrase)
            await account.hub.check_version()
            await account.sign_in(remember=True)
            kind = await account.remember(app.ctx.store)
        except CliError as e:
            self.query_one("#error", Static).update(e.message)
            return
        if kind == KEY_FILE:
            app.notify(FILE_KEY_NOTICE, timeout=10, markup=False)
        app.start(account)


class MainScreen(Screen):
    BINDINGS = [
        Binding("f5", "copy", "Copy"),
        Binding("q", "app.quit", "Quit"),
        Binding("tab", "app.focus_next", "Switch", show=False),
    ]

    DEFAULT_CSS = """
    MainScreen > Horizontal { height: 1fr; }
    """

    def __init__(self, account: Account):
        super().__init__()
        self.account = account
        self.pool = Pool(account)
        self._trees: dict[str, Tree] = {}
        self._gates: dict[str, asyncio.Semaphore] = {}
        self._watching: set[str] = set()

    def compose(self) -> ComposeResult:
        with Horizontal():
            yield LocalPane(Path.cwd(), id="local")
            yield RemotePane(self, id="remote")
        yield TransferStrip(id="transfers")
        yield Footer()

    def on_mount(self) -> None:
        self.query_one(LocalPane).table.focus()

    @property
    def local(self) -> LocalPane:
        return self.query_one(LocalPane)

    @property
    def remote(self) -> RemotePane:
        return self.query_one(RemotePane)

    @property
    def strip(self) -> TransferStrip:
        return self.query_one(TransferStrip)

    # ── Groups ───────────────────────────────────────────────────────────────

    async def tree(self, group: dict) -> Tree:
        """The group's index, asked once and then kept current by the node's pushes."""
        if group["id"] not in self._trees:
            session = await self.pool.session(group)
            self._trees[group["id"]] = Tree.from_index(await session.index())
            if group["id"] not in self._watching:
                self._watching.add(group["id"])
                session.ch.on_push("index_delta", lambda _msg, g=group: self._changed(g))
        return self._trees[group["id"]]

    def forget_tree(self, group: dict) -> None:
        self._trees.pop(group["id"], None)

    def _changed(self, group: dict) -> None:
        """Something changed in the group: the index is asked for again, and the
        pane redrawn if it is showing that group."""
        self.forget_tree(group)
        remote = self.remote
        if remote.group is not None and remote.group["id"] == group["id"]:
            remote.load()

    # ── F5 ───────────────────────────────────────────────────────────────────

    def action_copy(self) -> None:
        focused = self.app.focused
        if focused is not None and self.remote in focused.ancestors_with_self:
            self.run_worker(self._download(), group="copy")
        else:
            self.run_worker(self._upload(), group="copy")

    def _gate(self, group: dict) -> asyncio.Semaphore:
        return self._gates.setdefault(group["id"], asyncio.Semaphore(PER_GROUP))

    async def _ask(self, message: str, options: list[tuple[str, str]]) -> str:
        return await self.app.push_screen(Choice(message, options), wait_for_dismiss=True)

    async def _download(self) -> None:
        remote, local = self.remote, self.local
        rows = remote.selection()
        if remote.group is None or not rows:
            self.app.notify("Open a group and choose what to download.", timeout=4)
            return
        group, dest = remote.group, local.path
        try:
            tree = await self.tree(group)
        except CliError as e:
            self.app.notify(e.message, severity="error", markup=False)
            return
        plan: list[tuple[dict, Path]] = []
        dirs: list[Path] = []
        problems: list[str] = []
        for row in rows:
            path = row.data["path"]
            top = path.rsplit("/", 1)[-1]
            if not tree.root_available(path):
                problems.append(f"{top}: its folder on the node is unavailable")
                continue
            if not row.is_dir:
                try:
                    plan.append((row.data["entry"], local_path(dest, [top])[0]))
                except UnsafeName as e:
                    problems.append(str(e))
                continue
            for folder, items in tree.walk(path):
                rel = [top] + [p for p in folder[len(path):].strip("/").split("/") if p]
                try:
                    dirs.append(local_path(dest, rel)[0])
                    plan += [(i.entry, local_path(dest, rel + [i.name])[0])
                             for i in items if not i.is_dir]
                except UnsafeName as e:
                    problems.append(str(e))
        for p in problems:
            self.app.notify(p, severity="warning", markup=False)
        existing = [t for _, t in plan if t.exists()
                    and not t.with_name(t.name + PART_SUFFIX).exists()]
        overwrite = False
        if existing:
            names = ", ".join(safe(t.name) for t in existing[:5]) + (
                f" and {len(existing) - 5} more" if len(existing) > 5 else "")
            answer = await self._ask(
                f"{len(existing)} file(s) already in {dest}: {names}.",
                [("skip", "Skip them"), ("overwrite", "Overwrite"), ("cancel", "Cancel")])
            if answer == "cancel":
                return
            overwrite = answer == "overwrite"
            if not overwrite:
                plan = [(e, t) for e, t in plan if t not in existing]
        for d in dirs:
            await asyncio.to_thread(d.mkdir, parents=True, exist_ok=True)
        remote.marks.clear()
        remote.redraw()
        session = await self.pool.session(group)
        await asyncio.gather(*(self._one_download(group, session, e, t, overwrite)
                               for e, t in plan))
        local.action_refresh()

    async def _one_download(self, group, session, entry: dict, target: Path,
                            overwrite: bool) -> None:
        job = self.strip.add("↓", entry["name"], int(entry["size"]))
        async with self._gate(group):
            self.strip.update_job(job, state="running")
            try:
                await download(
                    session, entry, target, overwrite=overwrite,
                    progress=lambda n: self.strip.update_job(job, done=job.done + n),
                    on_queued=lambda n: self.strip.update_job(job, state="queued",
                                                              note=f"({n} ahead)"))
                self.strip.update_job(job, state="done",
                                      note="" if target.name == entry["name"]
                                      else f"saved as {target.name}")
            except TransferError as e:
                self.strip.update_job(job, state="failed", note=e.message)

    async def _upload(self) -> None:
        remote, local = self.remote, self.local
        rows = local.selection()
        if not rows:
            return
        if remote.group is None or not remote.path:
            self.app.notify("Open a folder of a group on the right to upload into.",
                            timeout=4)
            return
        group, folder = remote.group, remote.path
        try:
            tree = await self.tree(group)
        except CliError as e:
            self.app.notify(e.message, severity="error", markup=False)
            return
        root = tree.roots.get(folder.split("/", 1)[0], {})
        if not root.get("writable") or not root.get("available", True):
            self.app.notify(f"{safe(handle(group))}/{safe(folder)} does not accept "
                            "uploads (read-only or unavailable).", severity="error",
                            markup=False)
            return
        try:
            dirs, files, problems = plan_upload([str(local.path / r.name) for r in rows],
                                                recursive=True)
        except CliError as e:
            self.app.notify(e.message, severity="error", markup=False)
            return
        problems += refused_names([f.rel[-1] for f in files], dirs)
        if problems:
            await self._ask("Nothing was sent:\n" + "\n".join(safe(p) for p in problems),
                            [("ok", "OK")])
            return
        tops = sorted({d[0] for d in dirs} | {f.rel[0] for f in files})
        taken = taken_names(tree, folder, tops)
        if taken:
            answer = await self._ask(
                f"Already in {safe(folder)}: {', '.join(safe(t) for t in sorted(taken))}. "
                "The node never replaces a file.",
                [("skip", "Skip them"), ("cancel", "Cancel")])
            if answer != "skip":
                return
            dirs = [d for d in dirs if d[0] not in taken]
            files = [f for f in files if f.rel[0] not in taken]
        local.marks.clear()
        local.redraw()
        session = await self.pool.session(group)
        remote_dir = {(): folder}
        try:
            for d in dirs:
                remote_dir[tuple(d)] = await create_dir(session, remote_dir[tuple(d[:-1])],
                                                        d[-1])
        except TransferError as e:
            self.app.notify(e.message, severity="error", markup=False)
            return
        await asyncio.gather(*(self._one_upload(group, session, f,
                                                remote_dir[tuple(f.rel[:-1])])
                               for f in files))
        self._changed(group)

    async def _one_upload(self, group, session, item, folder: str) -> None:
        name = item.rel[-1]
        job = self.strip.add("↑", name, item.size)
        async with self._gate(group):
            self.strip.update_job(job, state="running")
            lock = threading.Lock()
            try:
                with open(item.local, "rb") as fh:
                    def read(offset: int, length: int) -> bytes:
                        with lock:
                            fh.seek(offset)
                            return fh.read(length)
                    stored = await upload(
                        session, read, item.size, name, folder,
                        progress=lambda n: self.strip.update_job(job, done=job.done + n),
                        on_queued=lambda n: self.strip.update_job(job, state="queued",
                                                                  note=f"({n} ahead)"))
                self.strip.update_job(job, state="done", note="" if stored == name
                                      else f"stored as {stored}")
            except OSError as e:
                self.strip.update_job(job, state="failed", note=e.strerror or str(e))
            except TransferError as e:
                self.strip.update_job(job, state="failed", note=e.message)

    async def on_unmount(self) -> None:
        with contextlib.suppress(Exception):
            await self.pool.close()


class DeviceCode:
    """`account.on_device_code` for the interactive mode: the code in a
    notification that stays until the approval lands; the node is retried meanwhile."""

    def __init__(self, app):
        self.app = app

    def __call__(self, group: dict, pending, *, first: bool) -> bool:
        if first:
            self.app.notify(
                f"This machine is not yet a device of yours on the node hosting "
                f"{safe(handle(group))}. Approve it from the MeshBay application or a "
                f"browser: the group's Settings → \"Your devices on this node\" → "
                f"enter {pending.link_code}. Waiting…",
                title="Approve this machine", timeout=600, markup=False)
        return True