vortex/tui/app/vortex_tui/main.py

"""Vortex TUI application: room list, timeline, bridges, drafts."""

from __future__ import annotations

from nio import (
    AsyncClient,
    MatrixRoom,
    RoomCreateResponse,
    RoomMessageMedia,
    RoomMessageNotice,
    RoomMessageText,
    SyncResponse,
)
from textual.app import App, ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal, Vertical
from textual.screen import Screen
from textual.widgets import Footer, Input, Label, ListItem, ListView, RichLog, Static

from . import drafts as drafts_db
from .provisioning import BridgeProvisioning, ProvisioningError, load_bridge_registry
from .qr import looks_like_qr_payload, render_qr
from .session import build_client, load_credentials

NETWORKS = ["gmessages", "telegram", "whatsapp", "twitter", "linkedin", "discord", "meta"]
LABELS = {
    "gmessages": "gm",
    "telegram": "tg",
    "whatsapp": "wa",
    "twitter": "tw",
    "linkedin": "li",
    "discord": "dc",
    "meta": "ig",
    "matrix": "mx",
}


class LoginWizardScreen(Screen[bool]):
    """Guided bridge login driven by the bridgev2 provisioning API.

    Step loop: pick a flow -> user_input/cookies steps become sequential
    prompts -> display_and_wait steps render the QR/code and long-poll until
    scanned -> complete.
    """

    BINDINGS = [Binding("escape", "cancel", "Cancel login")]

    def __init__(self, network: str, prov: BridgeProvisioning):
        super().__init__()
        self.network = network
        self.prov = prov
        self.login_id: str | None = None
        self._fields: list[dict] = []
        self._values: dict[str, str] = {}
        self._step: dict | None = None

    def compose(self) -> ComposeResult:
        yield Static(f"Login to {self.network} - esc cancels", id="wiz-title")
        yield Static("Loading login flows...", id="wiz-status")
        yield ListView(id="wiz-flows")
        yield RichLog(id="wiz-display", wrap=True, markup=False, highlight=False)
        yield Input(id="wiz-input")
        yield Footer()

    def on_mount(self) -> None:
        self.query_one("#wiz-input", Input).display = False
        self.run_worker(self._load_flows(), exclusive=True)

    def _status(self, text: str) -> None:
        self.query_one("#wiz-status", Static).update(text)

    async def _load_flows(self) -> None:
        try:
            flows = await self.prov.flows()
        except ProvisioningError as exc:
            self._status(f"Cannot reach bridge: {exc}")
            return
        lv = self.query_one("#wiz-flows", ListView)
        for flow in flows:
            item = ListItem(Label(f"{flow['name']} - {flow.get('description', '')}"))
            item.flow_id = flow["id"]
            await lv.append(item)
        self._status("Pick a login method (Enter):")
        lv.focus()

    async def on_list_view_selected(self, event: ListView.Selected) -> None:
        event.stop()
        flow_id = getattr(event.item, "flow_id", None)
        if not flow_id or self.login_id:
            return
        self.query_one("#wiz-flows", ListView).display = False
        self._status("Starting login...")
        self.run_worker(self._start(flow_id), exclusive=True)

    async def _start(self, flow_id: str) -> None:
        try:
            resp = await self.prov.start(flow_id)
        except ProvisioningError as exc:
            self._status(f"Login start failed: {exc}")
            return
        self.login_id = resp["login_id"]
        await self._handle_step(resp)

    async def _handle_step(self, step: dict) -> None:
        self._step = step
        kind = step["type"]
        instructions = step.get("instructions") or ""
        display = self.query_one("#wiz-display", RichLog)
        if kind == "complete":
            self.notify(f"{self.network}: logged in.")
            self.dismiss(True)
        elif kind in ("user_input", "cookies"):
            params = step.get(kind) or {}
            self._fields = list(params.get("fields") or [])
            self._values = {}
            if url := params.get("url"):
                display.write(f"Open {url} in a browser to obtain the values.")
            self._status(instructions or "Fill in the requested values:")
            self._next_field()
        elif kind == "display_and_wait":
            params = step.get("display_and_wait") or {}
            display.clear()
            if instructions:
                display.write(instructions)
            if params.get("type") == "qr":
                display.write(render_qr(params.get("data", "")))
                display.write(params.get("data", ""))
            else:
                display.write(f"Code: {params.get('data', '')}")
            self._status("Waiting for you to scan/confirm (refreshes automatically)...")
            self.run_worker(self._wait(step), exclusive=True)
        else:  # webauthn etc.
            self._status(f"Unsupported step type {kind!r} - use the bot DM instead.")

    async def _wait(self, step: dict) -> None:
        try:
            resp = await self.prov.step(self.login_id, step["step_id"], step["type"])
        except ProvisioningError as exc:
            self._status(f"Login failed: {exc}")
            return
        await self._handle_step(resp)

    def _next_field(self) -> None:
        box = self.query_one("#wiz-input", Input)
        if not self._fields:
            box.display = False
            self.run_worker(self._submit_fields(), exclusive=True)
            return
        field = self._fields[0]
        name = field.get("name") or field.get("id")
        desc = field.get("description") or ""
        box.placeholder = f"{name}{' - ' + desc if desc else ''}"
        box.password = field.get("type") == "password"
        box.value = field.get("default_value") or ""
        box.display = True
        box.focus()

    def on_input_submitted(self, event: Input.Submitted) -> None:
        event.stop()
        if event.input.id != "wiz-input" or not self._fields:
            return
        field = self._fields.pop(0)
        self._values[field["id"]] = event.value.strip()
        event.input.value = ""
        self._next_field()

    async def _submit_fields(self) -> None:
        self._status("Submitting...")
        try:
            resp = await self.prov.step(
                self.login_id, self._step["step_id"], self._step["type"], self._values
            )
        except ProvisioningError as exc:
            self._status(f"Login failed: {exc}")
            return
        await self._handle_step(resp)

    def action_cancel(self) -> None:
        if self.login_id:
            self.run_worker(self.prov.cancel(self.login_id))
        self.dismiss(False)


class BridgeScreen(Screen[str | None]):
    """Bridge accounts: Enter starts a guided login, m opens the bot DM."""

    BINDINGS = [
        Binding("escape", "back", "Back"),
        Binding("m", "bot_dm", "Bot DM"),
    ]

    def compose(self) -> ComposeResult:
        yield Static(
            "Bridges - Enter starts a guided login; m opens the bot DM for "
            "manual commands.",
            id="bridge-help",
        )
        yield ListView(id="bridge-list")
        yield Footer()

    def on_mount(self) -> None:
        app: VortexApp = self.app
        self.registry = load_bridge_registry()
        lv = self.query_one("#bridge-list", ListView)
        for net in NETWORKS:
            bot = f"@{net}bot:{app.domain}"
            item = ListItem(Label(self._row(net, "...")))
            item.net, item.bot = net, bot
            lv.append(item)
        lv.focus()
        self.run_worker(self._load_statuses(), exclusive=True)

    def _row(self, net: str, status: str) -> str:
        return f"[{LABELS[net]}] {net:<10} {status}"

    def _prov(self, net: str) -> BridgeProvisioning | None:
        info = self.registry.get(net)
        if not info or info.get("kind") != "v2":
            return None
        app: VortexApp = self.app
        return BridgeProvisioning(
            info["base_url"], app.client.user_id, app.client.access_token
        )

    async def _load_statuses(self) -> None:
        for item in list(self.query_one("#bridge-list", ListView).children):
            net = item.net
            if net not in self.registry:
                status = "not running (check BRIDGES_ENABLED / TELEGRAM_API_* in .env)"
            elif (prov := self._prov(net)) is None:
                status = "manual login via bot DM (press m)"
            else:
                try:
                    logins = (await prov.whoami()).get("logins") or []
                    status = (
                        f"logged in: {', '.join(l.get('name') or l['id'] for l in logins)}"
                        if logins
                        else "not logged in (Enter to log in)"
                    )
                except ProvisioningError:
                    status = "bridge unreachable"
            item.query_one(Label).update(self._row(net, status))

    async def on_list_view_selected(self, event: ListView.Selected) -> None:
        event.stop()
        net = event.item.net
        if net not in self.registry:
            self.notify(f"{net} is not in BRIDGES_ENABLED", severity="warning")
            return
        prov = self._prov(net)
        if prov is None:
            await self._open_bot_dm(event.item.bot)
            return

        def done(_ok: bool | None) -> None:
            self.run_worker(self._load_statuses(), exclusive=True)

        self.app.push_screen(LoginWizardScreen(net, prov), done)

    async def action_bot_dm(self) -> None:
        item = self.query_one("#bridge-list", ListView).highlighted_child
        if item is not None:
            await self._open_bot_dm(item.bot)

    async def _open_bot_dm(self, bot: str) -> None:
        app: VortexApp = self.app
        dm = app.find_bot_dm(bot)
        if dm:
            self.dismiss(dm)
            return
        resp = await app.client.room_create(invite=[bot], is_direct=True)
        if isinstance(resp, RoomCreateResponse):
            self.dismiss(resp.room_id)
        else:
            self.notify(f"room_create failed: {resp}", severity="error")

    def action_back(self) -> None:
        self.dismiss(None)


class DraftsScreen(Screen):
    """Pending drafts from the shared SQLite DB, polled every 2s."""

    BINDINGS = [
        Binding("escape", "back", "Back"),
        Binding("e", "edit", "Edit"),
        Binding("s", "send", "Send"),
        Binding("x", "discard", "Discard"),
    ]

    def compose(self) -> ComposeResult:
        yield Static("Drafts - e edit, s send, x discard, esc back", id="drafts-help")
        yield ListView(id="draft-list")
        yield Input(placeholder="Edit draft body, Enter saves", id="draft-edit")
        yield Footer()

    def on_mount(self) -> None:
        self.con = drafts_db.connect()
        self._editing: str | None = None
        self.query_one("#draft-edit", Input).display = False
        self.call_later(self._refresh)
        # ponytail: 2s poll instead of watchfiles; watch the db file if this chafes
        self.set_interval(2, self._refresh)
        self.query_one("#draft-list", ListView).focus()

    async def _refresh(self) -> None:
        app: VortexApp = self.app
        lv = self.query_one("#draft-list", ListView)
        idx = lv.index
        items = []
        for row in drafts_db.pending(self.con):
            room = app.client.rooms.get(row["room_id"]) if app.client else None
            name = (room.display_name if room else row["room_id"])[:24]
            body = row["body"].replace("\n", " ")[:48]
            item = ListItem(
                Label(f"{name:<24} | {body:<48} | {row['created_by']} @ {row['created_at']}")
            )
            item.draft = dict(row)
            items.append(item)
        await lv.clear()
        await lv.extend(items)
        if idx is not None and items:
            lv.index = min(idx, len(items) - 1)

    def _current(self) -> dict | None:
        item = self.query_one("#draft-list", ListView).highlighted_child
        return getattr(item, "draft", None)

    def action_edit(self) -> None:
        draft = self._current()
        if not draft:
            return
        self._editing = draft["draft_id"]
        box = self.query_one("#draft-edit", Input)
        box.value = draft["body"]
        box.display = True
        box.focus()

    def on_input_submitted(self, event: Input.Submitted) -> None:
        event.stop()
        if event.input.id != "draft-edit" or not self._editing:
            return
        drafts_db.update_body(self.con, self._editing, event.value)
        self._editing = None
        event.input.display = False
        self.query_one("#draft-list", ListView).focus()
        self.call_later(self._refresh)

    async def action_send(self) -> None:
        draft = self._current()
        app: VortexApp = self.app
        if not draft or not app.client:
            return
        resp = await app.client.room_send(
            draft["room_id"],
            "m.room.message",
            {"msgtype": "m.text", "body": draft["body"]},
            ignore_unverified_devices=True,
        )
        drafts_db.mark(self.con, draft["draft_id"], "sent", getattr(resp, "event_id", None))
        await self._refresh()

    def action_discard(self) -> None:
        draft = self._current()
        if draft:
            drafts_db.mark(self.con, draft["draft_id"], "discarded")
            self.call_later(self._refresh)

    def action_back(self) -> None:
        self.con.close()
        self.dismiss()


class VortexApp(App):
    """Three panes: room list | timeline over input, plus footer bindings."""

    CSS = """
    #rooms { width: 36; border-right: solid $accent; }
    #timeline { height: 1fr; }
    #bridge-help, #drafts-help { padding: 0 1; color: $text-muted; }
    #wiz-title, #wiz-status { padding: 0 1; }
    #wiz-flows { max-height: 10; }
    #wiz-display { height: 1fr; }
    """
    TITLE = "Vortex"
    SUB_TITLE = "/reply <text> replies to the last message in the room"
    BINDINGS = [
        Binding("q", "quit", "Quit"),
        Binding("b", "bridges", "Bridges"),
        Binding("d", "drafts", "Drafts"),
        Binding("tab", "focus_next", "Cycle focus"),
    ]

    def __init__(self, creds: dict):
        super().__init__()
        self.creds = creds
        self.client: AsyncClient | None = None
        self.domain = creds["user_id"].split(":", 1)[1]
        self.selected_room: str | None = None
        self.pending_select: str | None = None
        self.timelines: dict[str, list[str]] = {}
        self.last_event: dict[str, str] = {}
        self.unread: dict[str, int] = {}
        self.first_sync = False

    def compose(self) -> ComposeResult:
        with Horizontal():
            yield ListView(id="rooms")
            with Vertical():
                yield RichLog(id="timeline", wrap=True, markup=False, highlight=False)
                yield Input(
                    placeholder="Message (Enter sends, /reply <text> replies)", id="msg"
                )
        yield Footer()

    def on_mount(self) -> None:
        self.run_worker(self._connect_and_sync(), exclusive=True)

    # ---- matrix -----------------------------------------------------------
    async def _connect_and_sync(self) -> None:
        log = self.query_one("#timeline", RichLog)
        log.write("Connecting...")
        try:
            self.client = await build_client(self.creds)
        except Exception as exc:  # surface login problems in the UI
            log.write(f"Login failed: {exc}")
            return
        self.client.add_event_callback(
            self._on_event, (RoomMessageText, RoomMessageNotice, RoomMessageMedia)
        )
        self.client.add_response_callback(self._on_sync, SyncResponse)
        log.write(
            f"Logged in as {self.client.user_id} ({self.client.device_id}). Syncing..."
        )
        await self.client.sync_forever(timeout=30000, full_state=True)

    async def _on_sync(self, _response) -> None:
        self.first_sync = True
        await self._refresh_rooms()
        if self.pending_select and self.pending_select in self.client.rooms:
            rid, self.pending_select = self.pending_select, None
            self._select_room(rid)

    async def _on_event(self, room: MatrixRoom, event) -> None:
        line = self._format_event(room, event)
        self.timelines.setdefault(room.room_id, []).append(line)
        self.last_event[room.room_id] = event.event_id
        if room.room_id == self.selected_room:
            self.query_one("#timeline", RichLog).write(line)
        elif self.first_sync:
            self.unread[room.room_id] = self.unread.get(room.room_id, 0) + 1
            await self._refresh_rooms()

    def _format_event(self, room: MatrixRoom, event) -> str:
        sender = room.user_name(event.sender) or event.sender
        if isinstance(event, RoomMessageMedia):
            path = (event.url or "")[len("mxc://"):]
            url = f"{self.client.homeserver}/_matrix/media/v3/download/{path}"
            return f"<{sender}> [file] {event.body}: {url}"
        body = event.body or ""
        is_bot = any(event.sender == f"@{n}bot:{self.domain}" for n in NETWORKS)
        if looks_like_qr_payload(body, is_bot):
            return f"<{sender}> QR for {body}\n{render_qr(body)}"
        return f"<{sender}> {body}"

    # ---- rooms ------------------------------------------------------------
    def room_network(self, room: MatrixRoom) -> str:
        for uid in room.users:
            for net in NETWORKS:
                if uid.startswith(f"@{net}_") or uid == f"@{net}bot:{self.domain}":
                    return net
        return "matrix"

    def find_bot_dm(self, bot: str) -> str | None:
        # ponytail: "DM" = bot is a member and room has <=2 members
        for room in self.client.rooms.values():
            if bot in room.users and len(room.users) <= 2:
                return room.room_id
        return None

    def _room_label(self, room: MatrixRoom) -> str:
        n = self.unread.get(room.room_id, 0)
        badge = f" ({n})" if n else ""
        return f"[{LABELS[self.room_network(room)]}] {room.display_name or room.room_id}{badge}"

    async def _refresh_rooms(self) -> None:
        lv = self.query_one("#rooms", ListView)
        rooms = sorted(
            self.client.rooms.values(),
            key=lambda r: (self.room_network(r), (r.display_name or r.room_id).lower()),
        )
        items = []
        for room in rooms:
            item = ListItem(Label(self._room_label(room)))
            item.room_id = room.room_id
            items.append(item)
        idx = next(
            (i for i, r in enumerate(rooms) if r.room_id == self.selected_room), None
        )
        await lv.clear()
        await lv.extend(items)
        if idx is not None:
            lv.index = idx

    def on_list_view_selected(self, event: ListView.Selected) -> None:
        rid = getattr(event.item, "room_id", None)
        if rid:
            self._select_room(rid)

    def _select_room(self, rid: str) -> None:
        self.selected_room = rid
        self.unread.pop(rid, None)
        log = self.query_one("#timeline", RichLog)
        log.clear()
        for line in self.timelines.get(rid, []):
            log.write(line)
        room = self.client.rooms.get(rid) if self.client else None
        if room:  # update badge in place, no full rebuild
            for item in self.query_one("#rooms", ListView).children:
                if getattr(item, "room_id", None) == rid:
                    item.query_one(Label).update(self._room_label(room))
        self.query_one("#msg", Input).focus()

    # ---- sending ----------------------------------------------------------
    async def on_input_submitted(self, event: Input.Submitted) -> None:
        if event.input.id != "msg":
            return
        text = event.value.strip()
        event.input.value = ""
        if not text or not self.client or not self.selected_room:
            return
        content: dict = {"msgtype": "m.text", "body": text}
        if text.startswith("/reply "):
            content["body"] = text[len("/reply "):]
            last = self.last_event.get(self.selected_room)
            if last:
                content["m.relates_to"] = {"m.in_reply_to": {"event_id": last}}
        await self.client.room_send(
            self.selected_room,
            "m.room.message",
            content,
            ignore_unverified_devices=True,
        )

    # ---- actions ----------------------------------------------------------
    def action_bridges(self) -> None:
        if self.client:
            self.push_screen(BridgeScreen(), self._after_bridge)

    def _after_bridge(self, room_id: str | None) -> None:
        if not room_id:
            return
        if self.client and room_id in self.client.rooms:
            self._select_room(room_id)
        else:
            self.pending_select = room_id
            self.notify("DM created - it will open once it syncs.")

    def action_drafts(self) -> None:
        self.push_screen(DraftsScreen())

    async def action_quit(self) -> None:
        if self.client:
            await self.client.close()
        self.exit()


def run() -> None:
    creds = load_credentials()
    VortexApp(creds).run()


if __name__ == "__main__":
    run()