from collections import defaultdict import logging from rich.text import Text from textual.containers import Horizontal, Vertical from textual.widget import Widget from textual.widgets import Input, OptionList, Static, Tree from textual.widgets.option_list import Option logger = logging.getLogger(__name__) class PolicyTreeWidget(Widget): """Widget for displaying and searching a hierarchical policy tree.""" def __init__(self, policies, devices): super().__init__() self.policies = policies self.devices = devices self.last_highlighted_node = None self.leaf_counts = defaultdict(int) def compose(self): policy_tree = Tree("", id="policy_tree") # Label set in on_mount policy_tree.styles.width = "2fr" policy_tree.styles.height = "100%" label = Static("Device Search:") search_box = Input( placeholder="Search policies or devices...", id="tree_search" ) details_pane = Static("", id="details_pane") with Horizontal(): yield policy_tree with Vertical() as right_pane: right_pane.styles.width = "3fr" yield label yield search_box yield details_pane def on_mount(self) -> None: self._precompute_leaf_counts() # Update root label with total leaf count total_leaves = sum( self.leaf_counts.get(policy.groupid, 0) for policy in self.policies if policy.parent == "global-policy-settings" ) policy_tree = self.query_one("#policy_tree", Tree) policy_tree.root.set_label(f"Agents in Policies: ({total_leaves})") self._build_tree() def _precompute_leaf_counts(self): """Precompute leaf counts for each policy group.""" device_counts = defaultdict(int) for device in self.devices: device_counts[device.groupid] += 1 child_map = defaultdict(list) for policy in self.policies: child_map[policy.parent].append(policy.groupid) def count_leaves(groupid): count = device_counts[groupid] for child_id in child_map.get(groupid, []): count += count_leaves(child_id) self.leaf_counts[groupid] = count return count for policy in self.policies: if policy.parent == "global-policy-settings": count_leaves(policy.groupid) def _build_tree(self): policy_tree = self.query_one("#policy_tree", Tree) node_map = {} # Sort top-level policies top_policies = [ p for p in self.policies if p.parent == "global-policy-settings" ] top_policies.sort( key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True ) for policy in top_policies: label = f"{policy.name} ({self.leaf_counts.get(policy.groupid, 0)})" node = policy_tree.root.add(label=label, data=policy) node_map[policy.groupid] = node # Sort and add child policies children_by_parent = defaultdict(list) for policy in self.policies: if policy.parent != "global-policy-settings": children_by_parent[policy.parent].append(policy) for parent_id, children in children_by_parent.items(): children.sort( key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True ) parent_node = node_map.get(parent_id) if parent_node: for policy in children: label = f"{policy.name} ({self.leaf_counts.get(policy.groupid, 0)})" node = parent_node.add(label=label, data=policy) node_map[policy.groupid] = node # Add devices (leaf nodes) for device in self.devices: group_id = device.groupid parent_node = node_map.get(group_id) if parent_node: parent_node.add(label=device.hostname, data=device) def _collect_tree_nodes(self, node, all_nodes): all_nodes.append(node) for child in node.children: self._collect_tree_nodes(child, all_nodes) def _remove_match_selector(self): try: existing = self.query("#match_selector") for widget in existing: if widget.is_attached: widget.remove() except Exception as exc: logger.debug("Failed to remove match_selector: %s", exc) def on_tree_node_selected(self, message: Tree.NodeSelected) -> None: node = message.node data = node.data details_pane = self.query_one("#details_pane", Static) if self.last_highlighted_node is not None: original_label = str(self.last_highlighted_node.label).strip() if isinstance(self.last_highlighted_node.label, Text): original_label = self.last_highlighted_node.label.plain self.last_highlighted_node.set_label(original_label) label_text = str(node.label).strip() if isinstance(node.label, Text): label_text = node.label.plain highlighted_label = Text(label_text, style="reverse bold") node.set_label(highlighted_label) self.last_highlighted_node = node if data: details = "\n".join( f"{key}: {value}" for key, value in data.__dict__.items() ) else: details = f"Selected: {node.label}" details_pane.update(details) message.stop() def on_input_submitted(self, message: Input.Submitted) -> None: self._remove_match_selector() query = message.value.strip().lower() tree = self.query_one("#policy_tree", Tree) details_pane = self.query_one("#details_pane", Static) all_nodes = [] self._collect_tree_nodes(tree.root, all_nodes) label_to_node = {} for node in all_nodes: label_text = str(node.label).lower() label_to_node[label_text] = node if node.data: data_dict = ( node.data.__dict__ if hasattr(node.data, "__dict__") else node.data ) for key, value in data_dict.items(): if isinstance(value, str): label_to_node[value.lower()] = node matches = sorted([label for label in label_to_node if query in label]) if matches: try: option_list = self.query_one("#match_selector", OptionList) option_list.clear_options() option_list.display = True except: option_list = OptionList(id="match_selector") details_pane.parent.mount(option_list) for label in matches: option_list.add_option(Option(label, id=f"match_{label}")) details_pane.update(f"Found {len(matches)} matches. Select one below.") else: self._remove_match_selector() details_pane.update("No matches found.") def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None: selected_id = event.option.id.replace("match_", "") tree = self.query_one("#policy_tree", Tree) details_pane = self.query_one("#details_pane", Static) all_nodes = [] self._collect_tree_nodes(tree.root, all_nodes) label_to_node = {str(node.label).lower(): node for node in all_nodes} match_node = label_to_node.get(selected_id.lower()) if match_node: node = match_node path = [] while node: path.insert(0, node) node = node.parent for node in path: node.expand() tree.select_node(match_node) tree.scroll_to_node(match_node) match_node.set_label(Text(str(match_node.label), style="reverse bold")) details_pane.update(f"Selected: {match_node.label}") try: option_list = self.query_one("#match_selector", OptionList) option_list.remove() except: pass