"""
Tree widget for displaying hierarchical inspector view.
Handles rendering, interaction, and navigation of tree nodes.
"""

import pygame
import pygame_gui
from pygame_gui.core import UIElement
from typing import List, Optional, Tuple, Any
from ui.inspector_tree import TreeNode, SimulationNode, TreeSelectionManager
from config.constants import (
    PANEL_BACKGROUND_COLOR, PANEL_SELECTED_COLOR, PANEL_HOVER_COLOR,
    PANEL_TEXT_COLOR, PANEL_ICON_COLOR, PANEL_NODE_HEIGHT, PANEL_INDENTATION
)


class TreeWidget(UIElement):
    """Interactive tree widget for hierarchical entity inspection."""

    def __init__(self, relative_rect: pygame.Rect, manager: pygame_gui.UIManager, world: Any):
        super().__init__(relative_rect, manager, container=None,
                        starting_height=1, layer_thickness=1)

        self.world = world
        self.font = pygame.font.Font(None, 16)
        self.small_font = pygame.font.Font(None, 12)

        # Tree structure
        self.root_node = SimulationNode(world)
        self.visible_nodes: List[TreeNode] = []
        self.all_nodes: List[TreeNode] = []

        # Selection management
        self.selection_manager = TreeSelectionManager()

        # Visual properties (using standardized panel styling)
        self.node_height = PANEL_NODE_HEIGHT
        self.expand_collapse_width = PANEL_INDENTATION
        self.icon_size = 8
        self.text_color = PANEL_TEXT_COLOR
        self.selected_color = PANEL_SELECTED_COLOR
        self.hover_color = PANEL_HOVER_COLOR
        self.expand_icon_color = PANEL_ICON_COLOR

        # Interaction state
        self.hovered_node: Optional[TreeNode] = None
        self.drag_start_node: Optional[TreeNode] = None
        self.is_dragging = False
        self.last_mouse_pos: Optional[Tuple[int, int]] = None

        # Scrolling
        self.scroll_offset = 0
        self.total_height = 0

        # Initialize tree structure
        self._update_tree_structure()
        self._update_visible_nodes()

    def _update_tree_structure(self) -> None:
        """Update the tree structure based on current world state."""
        self.root_node.update_entity_counts()
        self.all_nodes = self._build_all_nodes_list(self.root_node)

    def _build_all_nodes_list(self, node: TreeNode) -> List[TreeNode]:
        """Build a flat list of all nodes in the tree."""
        nodes = [node]
        for child in node.children:
            nodes.extend(self._build_all_nodes_list(child))
        return nodes

    def _update_visible_nodes(self) -> None:
        """Update the list of visible nodes based on expanded state."""
        self.visible_nodes = [self.root_node]
        self.visible_nodes.extend(self.root_node.get_all_visible_descendants())
        self._update_node_layout()
        self.total_height = len(self.visible_nodes) * self.node_height

    def _update_node_layout(self) -> None:
        """Update the layout rectangles for all visible nodes."""
        y_offset = 0
        for node in self.visible_nodes:
            node.rect = pygame.Rect(
                0,
                y_offset - self.scroll_offset,
                self.rect.width,
                self.node_height
            )
            y_offset += self.node_height

    def _get_node_at_position_local(self, local_pos: Tuple[int, int]) -> Optional[TreeNode]:
        """Get the node at the given local position (already converted)."""
        for node in self.visible_nodes:
            if node.rect.collidepoint(local_pos):
                return node
        return None

    def _get_node_at_position(self, pos: Tuple[int, int]) -> Optional[TreeNode]:
        """Get the node at the given screen position."""
        # Convert screen coordinates to local tree widget coordinates
        local_pos = (pos[0] - self.rect.x, pos[1] - self.rect.y)
        return self._get_node_at_position_local(local_pos)

    def _get_expand_collapse_rect(self, node: TreeNode) -> pygame.Rect:
        """Get the rectangle for the expand/collapse icon."""
        indent = node.get_indent()
        return pygame.Rect(
            indent + 5,
            node.rect.y + (self.node_height - self.icon_size) // 2,
            self.icon_size,
            self.icon_size
        )

    def _is_click_on_expand_collapse_local(self, node: TreeNode, local_pos: Tuple[int, int]) -> bool:
        """Check if a click is on the expand/collapse icon (local coordinates)."""
        if not node.can_expand():
            return False

        expand_rect = self._get_expand_collapse_rect(node)
        return expand_rect.collidepoint(local_pos)

    def _is_click_on_expand_collapse(self, node: TreeNode, pos: Tuple[int, int]) -> bool:
        """Check if a click is on the expand/collapse icon."""
        if not node.can_expand():
            return False

        local_pos = (pos[0] - self.rect.x, pos[1] - self.rect.y)
        expand_rect = self._get_expand_collapse_rect(node)
        return expand_rect.collidepoint(local_pos)

    def _render_node(self, surface: pygame.Surface, node: TreeNode) -> None:
        """Render a single tree node."""
        if not node.rect.colliderect(surface.get_rect()):
            return

        # Background
        if node.is_selected:
            pygame.draw.rect(surface, self.selected_color, node.rect)
        elif node == self.hovered_node:
            pygame.draw.rect(surface, self.hover_color, node.rect)

        
        # Expand/collapse icon
        if node.can_expand():
            icon_rect = self._get_expand_collapse_rect(node)
            icon_color = self.expand_icon_color

            # Draw + or - icon
            center_x = icon_rect.centerx
            center_y = icon_rect.centery
            size = 3

            # Horizontal line
            pygame.draw.line(surface, icon_color,
                           (center_x - size, center_y),
                           (center_x + size, center_y), 1)

            # Vertical line (only for collapsed nodes)
            if not node.is_expanded:
                pygame.draw.line(surface, icon_color,
                               (center_x, center_y - size),
                               (center_x, center_y + size), 1)

        # Node text
        indent = node.get_indent()
        text_x = indent + (self.expand_collapse_width if node.can_expand() else 10)

        text_surface = self.font.render(node.get_display_text(), True, self.text_color)
        text_rect = pygame.Rect(
            text_x,
            node.rect.y + (self.node_height - text_surface.get_height()) // 2,
            text_surface.get_width(),
            text_surface.get_height()
        )

        # Clip text if it's too wide
        max_text_width = self.rect.width - text_x - 10
        if text_rect.width > max_text_width:
            # Truncate text
            text = node.get_display_text()
            while text and text_surface.get_width() > max_text_width - 20:
                text = text[:-1]
                text_surface = self.font.render(text + "...", True, self.text_color)

        surface.blit(text_surface, text_rect)

    def handle_event(self, event: pygame.event.Event) -> bool:
        """Handle input events for the tree widget."""
        if event.type == pygame.MOUSEBUTTONDOWN:
            if event.button == 1:  # Left click
                # Event position is already converted to local coordinates by HUD
                self.last_mouse_pos = event.pos
                node = self._get_node_at_position_local(event.pos)
                if node:
                    # Check for expand/collapse click
                    if self._is_click_on_expand_collapse_local(node, event.pos):
                        node.toggle_expand()
                        self._update_visible_nodes()
                    else:
                        # Selection click
                        multi_select = pygame.key.get_pressed()[pygame.K_LSHIFT]
                        self.selection_manager.select_node(node, multi_select)
                return True

        elif event.type == pygame.MOUSEMOTION:
            # Event position is already converted to local coordinates by HUD
            self.last_mouse_pos = event.pos
            self.hovered_node = self._get_node_at_position_local(event.pos)
            return True

        elif event.type == pygame.MOUSEWHEEL:
            # Handle scrolling
            if event.y > 0:
                self.scroll_offset = max(0, self.scroll_offset - self.node_height)
            else:
                max_scroll = max(0, self.total_height - self.rect.height)
                self.scroll_offset = min(max_scroll, self.scroll_offset + self.node_height)
            self._update_node_layout()
            return True

        return False

    def update(self, time_delta: float) -> None:
        """Update the tree widget with smart performance optimizations."""
        # Smart update system: only update tree structure when necessary
        if self._should_update_tree():
            self._update_tree_structure()
            self._update_visible_nodes()
            # Recalculate hover state based on last mouse position after structure changes
            if self.last_mouse_pos:
                self.hovered_node = self._get_node_at_position_local(self.last_mouse_pos)

    def _should_update_tree(self) -> bool:
        """Determine if the tree structure needs updating."""
        if not hasattr(self, '_last_entity_count'):
            self._last_entity_count = len(self.world.get_objects())
            return True

        current_count = len(self.world.get_objects())
        if current_count != self._last_entity_count:
            self._last_entity_count = current_count
            return True

        # Also check if any previously selected entities have been removed
        if self.selection_manager.selected_nodes:
            for node in self.selection_manager.selected_nodes[:]:
                if hasattr(node, 'entity'):
                    try:
                        # Check if entity still exists in world
                        if node.entity not in self.world.get_objects():
                            return True
                    except:
                        return True

        return False

    def draw(self, surface: pygame.Surface) -> None:
        """Draw the tree widget."""
        # Create a clipping surface for the tree area
        tree_surface = pygame.Surface((self.rect.width, self.rect.height))
        tree_surface.fill(PANEL_BACKGROUND_COLOR)  # Using standardized background color

        # Set up clipping rect to prevent rendering outside bounds
        clip_rect = pygame.Rect(0, 0, self.rect.width, self.rect.height)
        tree_surface.set_clip(clip_rect)

        # Render visible nodes (only those that should be visible after scrolling)
        for node in self.visible_nodes:
            # Only render if node is within the visible tree area
            if node.rect.bottom > 0 and node.rect.top < self.rect.height:
                self._render_node(tree_surface, node)

        # Blit to main surface at position (0,0) since this is already positioned
        surface.blit(tree_surface, (0, 0))

    def expand_all(self) -> None:
        """Expand all nodes in the tree."""
        for node in self.all_nodes:
            if node.can_expand():
                node.expand()
        self._update_visible_nodes()

    def collapse_all(self) -> None:
        """Collapse all nodes except the root."""
        for node in self.all_nodes[1:]:  # Skip root node
            if node.can_expand():
                node.collapse()
        self._update_visible_nodes()

    def get_selected_entities(self) -> List[Any]:
        """Get the entities currently selected in the tree."""
        return self.selection_manager.get_selected_entities()

    def select_entities(self, entities: List[Any]) -> None:
        """Select specific entities in the tree."""
        self.selection_manager.clear_selection()

        for entity in entities:
            # Find the entity node for this entity
            for node in self.all_nodes:
                if hasattr(node, 'entity') and node.entity == entity:
                    self.selection_manager.select_node(node, multi_select=True)
                    break

        # Expand parent nodes to show selected entities
        for node in self.selection_manager.selected_nodes:
            parent = node.parent
            while parent:
                if parent.can_expand():
                    parent.expand()
                parent = parent.parent

        self._update_visible_nodes()

    def scroll_to_node(self, node: TreeNode) -> None:
        """Scroll the tree to show the given node."""
        if node not in self.visible_nodes:
            return

        node_top = node.rect.y
        node_bottom = node.rect.y + node.rect.height

        if node_top < 0:
            self.scroll_offset = max(0, self.scroll_offset + node_top)
        elif node_bottom > self.rect.height:
            self.scroll_offset = min(
                self.total_height - self.rect.height,
                self.scroll_offset + (node_bottom - self.rect.height)
            )

        self._update_node_layout()

    def clear_hover(self) -> None:
        """Clear the hover state (called when mouse leaves tree widget)."""
        self.hovered_node = None
        self.last_mouse_pos = None