#    Copyright (C) 2025 James Brotosky, Brandon Wickline
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.

from collections import defaultdict
import logging

from rich.text import Text
from textual.containers import Horizontal, Vertical
from textual.message import Message
from textual.widget import Widget
from textual.widgets import Button, 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."""

    class ViewExecutionHistory(Message):
        """Message sent when user wants to view execution history for a device."""

        def __init__(self, device):
            super().__init__()
            self.device = device

    class GenerateOTP(Message):
        """Message sent when user wants to generate OTP for a device."""

        def __init__(self, device):
            super().__init__()
            self.device = device

    class ToggleEnforcement(Message):
        """Message sent when user wants to toggle audit/enforcement for a device."""

        def __init__(self, device):
            super().__init__()
            self.device = device

    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
        self.selected_device = None  # Track currently selected device

    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"
        )
        exec_history_button = Button(
            "📊 Execution History", id="view_exec_history_button", disabled=True
        )
        exec_history_button.styles.margin = (0, 1, 0, 0)  # Right margin

        otp_button = Button("🎫 Generate OTP", id="generate_otp_button", disabled=True)
        otp_button.styles.margin = (0, 1, 0, 0)  # Right margin

        toggle_enforcement_button = Button(
            "🔄 Toggle Enforcement/Audit", id="toggle_enforcement_button", disabled=True
        )
        # No right margin on last button
        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
                # Action buttons in a horizontal row
                with Horizontal() as button_row:
                    button_row.styles.height = "auto"
                    yield exec_history_button
                    yield otp_button
                    yield toggle_enforcement_button
                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 refresh_data(self, policies, devices):
        """Refresh the widget with new data and rebuild the tree."""
        self.policies = policies
        self.devices = devices
        self.selected_device = None

        # Disable all buttons since selection is lost
        try:
            self.query_one("#view_exec_history_button", Button).disabled = True
            self.query_one("#generate_otp_button", Button).disabled = True
            self.query_one("#toggle_enforcement_button", Button).disabled = True
        except:
            pass

        # Rebuild tree with new data
        self._precompute_leaf_counts()
        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()
        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)
        exec_history_button = self.query_one("#view_exec_history_button", Button)
        otp_button = self.query_one("#generate_otp_button", Button)
        toggle_enforcement_button = self.query_one("#toggle_enforcement_button", Button)

        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

        # Check if selected node is a device (has Agent data)
        from models.agent import Agent

        if data and isinstance(data, Agent):
            self.selected_device = data
            exec_history_button.disabled = False
            otp_button.disabled = False
            toggle_enforcement_button.disabled = False
        else:
            self.selected_device = None
            exec_history_button.disabled = True
            otp_button.disabled = True
            toggle_enforcement_button.disabled = True

        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

    def on_button_pressed(self, event: Button.Pressed) -> None:
        """Handle button presses."""
        if event.button.id == "view_exec_history_button":
            if self.selected_device:
                self.post_message(self.ViewExecutionHistory(self.selected_device))
                event.stop()
        elif event.button.id == "generate_otp_button":
            if self.selected_device:
                self.post_message(self.GenerateOTP(self.selected_device))
                event.stop()
        elif event.button.id == "toggle_enforcement_button":
            if self.selected_device:
                self.post_message(self.ToggleEnforcement(self.selected_device))
                event.stop()