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, Switch, 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)
        self.match_type = "Count"  # Default to sorting by count

    def compose(self):
        # Create the switch and its label
        switch = Switch(value=False, id="match_switch")
        switch.styles.margin = (0, 0, 0, 0)  # top, right, bottom, left
        switch.styles.padding = (0, 0, 0, 0)

        switch_label = Static("Sort: Count", id="match_switch_label")
        switch_label.styles.margin = (1, 0, 0, 0)
        switch_label.styles.padding = (0, 0, 0, 0)

        # Create the tree
        policy_tree = Tree("", id="policy_tree")  # Label set in on_mount
        policy_tree.styles.width = "2fr"
        policy_tree.styles.height = "100%"

        # Create the search box and details pane
        label = Static("Device Search:")
        search_box = Input(
            placeholder="Search policies or devices...", id="tree_search"
        )
        details_pane = Static("", id="details_pane")

        # Layout the UI
        with Horizontal():
            yield policy_tree
            with Vertical() as right_pane:
                right_pane.styles.width = "3fr"
                # Use a Horizontal container for the switch and label
                with Horizontal() as switch_container:
                    switch_container.styles.height = 3
                    switch_container.styles.margin = (0, 0, 0, 1)
                    switch_container.styles.padding = (0, 0, 0, 0)
                    yield switch
                    yield switch_label
                # Add the search box and details pane
                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()
        # Expand the root node
        policy_tree.root.expand()

    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)
        policy_tree.clear()  # Clear existing nodes
        node_map = {}

        # Sort top-level policies
        top_policies = [
            p for p in self.policies if p.parent == "global-policy-settings"
        ]

        # Sort by count (default) or alphabetically
        if getattr(self, "match_type", "Count") == "Count":
            top_policies.sort(
                key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True
            )
        else:  # Alphabetical
            top_policies.sort(key=lambda p: p.name.lower())

        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():
            if getattr(self, "match_type", "Count") == "Count":
                children.sort(
                    key=lambda p: self.leaf_counts.get(p.groupid, 0), reverse=True
                )
            else:  # Alphabetical
                children.sort(key=lambda p: p.name.lower())

            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) - always sort alphabetically
        devices_by_group = defaultdict(list)
        for device in self.devices:
            devices_by_group[device.groupid].append(device)

        for group_id, devices in devices_by_group.items():
            devices.sort(key=lambda d: d.hostname.lower())  # Always sort alphabetically
            parent_node = node_map.get(group_id)
            if parent_node:
                for device in devices:
                    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_switch_changed(self, event: Switch.Changed):
        self.match_type = "Alpha" if event.value else "Count"
        self.query_one("#match_switch_label", Static).update(
            f"Sort: {self.match_type.capitalize()}"
        )
        self._build_tree()

    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
