browser-cli/browser_cli/commands/rendering/tabs_tree.py

"""Tabs tree renderer."""
from __future__ import annotations

from collections.abc import Iterable

from rich.console import Console
from rich.text import Text
from rich.tree import Tree

from browser_cli.commands.rendering.common import Row, int_value, item_value, text_value, tree_title_limit, tree_url_limit
from browser_cli.commands.rendering.labels import (
  DEFAULT_BROWSER_STYLE,
  DEFAULT_SCOPE,
  browser_label_style,
  group_tree_label,
  scoped_browser_label,
  tab_tree_label,
)

GroupId = object
GroupKey = tuple[str, str, int | None, GroupId]
TreeNodeKey = tuple[str, str]
WindowNodeKey = tuple[str, str, int]
BrowserGroupNodeKey = tuple[str, str, int, GroupId]

def browser_scope(item: Row) -> str:
  """Return the remote/local scope key used by tree renderers."""
  return text_value(item_value(item, "browser_group")) or DEFAULT_SCOPE

def browser_label_key(item: Row) -> str:
  """Return the browser/profile key used by tree renderers."""
  return text_value(item_value(item, "browser")) or DEFAULT_SCOPE

def tab_window_id(tab: Row) -> int:
  """Return a stable window id from object or dict-shaped tab responses."""
  return int_value(item_value(tab, "window_id", item_value(tab, "windowId", 0)))

def tab_group_id(tab: Row) -> GroupId | None:
  """Return a tab group id from object or dict-shaped tab responses."""
  group_id = item_value(tab, "group_id", item_value(tab, "groupId"))
  return None if group_id is None else group_id

def tab_sort_key(tab: Row) -> tuple[str, str, int, int, int, int]:
  """Stable tab ordering across multi-browser responses."""
  group_id = tab_group_id(tab)
  return (
    browser_scope(tab),
    browser_label_key(tab),
    tab_window_id(tab),
    int_value(item_value(tab, "index", 0)),
    int_value(group_id, -1) if group_id is not None else -1,
    int_value(item_value(tab, "id", 0)),
  )

class TabsTreeBuilder:
  """Stateful builder for the browser tabs tree.

  The tree has optional scope nodes (remote host/local), then browser/profile,
  then window, then browser tab-groups/tabs. Keeping this state in a helper
  keeps ``build_tabs_tree`` small while preserving stable node reuse.
  """

  tabs: list[Row]
  groups: list[Row]
  show_urls: bool
  show_browser: bool
  group_by_scope: bool
  title_limit: int
  url_limit: int
  root: Tree
  group_info: dict[GroupKey, Row]
  browser_styles: dict[str, str]
  scope_nodes: dict[str, Tree]
  browser_nodes: dict[TreeNodeKey, Tree]
  window_nodes: dict[WindowNodeKey, Tree]
  group_nodes: dict[BrowserGroupNodeKey, Tree]

  def __init__(
    self,
    tabs: Iterable[Row],
    groups: Iterable[Row],
    *,
    console: Console,
    show_urls: bool = False,
  ):
    self.tabs = sorted(tabs, key=tab_sort_key)
    self.groups = list(groups)
    self.show_urls = show_urls
    self.show_browser = any(bool(item_value(tab, "browser")) for tab in self.tabs)
    self.group_by_scope = any(bool(item_value(item, "browser_group")) for item in self.tabs + self.groups)
    self.title_limit = tree_title_limit(console=console, show_browser=self.show_browser, show_urls=show_urls)
    self.url_limit = tree_url_limit(self.title_limit, console=console)
    self.root = Tree("[bold]Tabs[/bold]")
    self.group_info = self._group_info()
    self.browser_styles = self._browser_styles()
    self.scope_nodes = {}
    self.browser_nodes = {}
    self.window_nodes = {}
    self.group_nodes = {}

  def build(self) -> Tree:
    for tab in self.tabs:
      self._add_tab(tab)
    return self.root

  def _group_info(self) -> dict[GroupKey, Row]:
    return {
      (
        browser_scope(group),
        browser_label_key(group),
        int_value(item_value(group, "window_id", item_value(group, "windowId")), 0),
        item_value(group, "id"),
      ): group
      for group in self.groups
    }

  def _browser_styles(self) -> dict[str, str]:
    styles: dict[str, str] = {}
    for item in self.tabs + self.groups:
      key = browser_label_key(item)
      styles.setdefault(key, browser_label_style(text_value(item_value(item, "browser_name")) or None))
    return styles

  def _scope_node(self, scope: str) -> Tree:
    if not self.group_by_scope:
      return self.root
    node = self.scope_nodes.get(scope)
    if node is None:
      node = self.root.add(Text(scope, style="bold"))
      self.scope_nodes[scope] = node
    return node

  def _browser_node(self, scope: str, browser: str) -> Tree:
    key = (scope, browser)
    node = self.browser_nodes.get(key)
    if node is None:
      parent = self._scope_node(scope)
      if self.show_browser:
        label = scoped_browser_label(browser, scope, grouped=self.group_by_scope)
        node = parent.add(Text(label, style=self.browser_styles.get(browser, DEFAULT_BROWSER_STYLE)))
      else:
        node = parent
      self.browser_nodes[key] = node
    return node

  def _window_node(self, scope: str, browser: str, window_id: int) -> Tree:
    key = (scope, browser, window_id)
    node = self.window_nodes.get(key)
    if node is None:
      node = self._browser_node(scope, browser).add(f"Window {window_id}")
      self.window_nodes[key] = node
    return node

  def _group_node(self, scope: str, browser: str, window_id: int, group_id: GroupId, parent: Tree) -> Tree:
    key = (scope, browser, window_id, group_id)
    node = self.group_nodes.get(key)
    if node is None:
      group = self.group_info.get(key) or self.group_info.get((scope, browser, None, group_id))
      node = parent.add(group_tree_label(group_id, group, title_limit=self.title_limit))
      self.group_nodes[key] = node
    return node

  def _add_tab(self, tab: Row) -> None:
    scope = browser_scope(tab)
    browser = browser_label_key(tab)
    window_id = tab_window_id(tab)
    window_node = self._window_node(scope, browser, window_id)
    group_id = tab_group_id(tab)
    parent = window_node if group_id is None else self._group_node(scope, browser, window_id, group_id, window_node)
    parent.add(tab_tree_label(
      tab,
      title_limit=self.title_limit,
      show_urls=self.show_urls,
      url_limit=self.url_limit,
    ))

def build_tabs_tree(
  tabs: Iterable[Row],
  groups: Iterable[Row],
  *,
  console: Console,
  show_urls: bool = False,
) -> Tree:
  """Build a remote/local → browser → window → group/tab tree from tab responses."""
  return TabsTreeBuilder(tabs, groups, console=console, show_urls=show_urls).build()