"""
Tree node classes for the hierarchical inspector view.
Provides extensible tree structure for entity inspection.
"""

from abc import ABC, abstractmethod
from typing import List, Optional, Dict, Any
import pygame


class TreeNode(ABC):
    """Base class for tree nodes in the inspector."""

    def __init__(self, label: str, parent: Optional['TreeNode'] = None):
        self.label = label
        self.parent = parent
        self.children: List['TreeNode'] = []
        self.is_expanded = False
        self.is_selected = False
        self.depth = 0 if parent is None else parent.depth + 1
        self.rect = pygame.Rect(0, 0, 0, 20)  # Will be updated during layout

    def add_child(self, child: 'TreeNode') -> None:
        """Add a child node to this node."""
        child.parent = self
        child.depth = self.depth + 1
        self.children.append(child)

    def remove_child(self, child: 'TreeNode') -> None:
        """Remove a child node from this node."""
        if child in self.children:
            self.children.remove(child)
            child.parent = None

    def toggle_expand(self) -> None:
        """Toggle the expanded state of this node."""
        if self.children:
            self.is_expanded = not self.is_expanded

    def expand(self) -> None:
        """Expand this node."""
        if self.children:
            self.is_expanded = True

    def collapse(self) -> None:
        """Collapse this node."""
        self.is_expanded = False

    def get_visible_children(self) -> List['TreeNode']:
        """Get children that are currently visible."""
        if not self.is_expanded:
            return []
        return self.children

    def get_all_visible_descendants(self) -> List['TreeNode']:
        """Get all visible descendants of this node."""
        visible = []
        if self.is_expanded:
            for child in self.children:
                visible.append(child)
                visible.extend(child.get_all_visible_descendants())
        return visible

    def is_leaf(self) -> bool:
        """Check if this node has no children."""
        return len(self.children) == 0

    @abstractmethod
    def get_display_text(self) -> str:
        """Get the display text for this node."""
        pass

    @abstractmethod
    def can_expand(self) -> bool:
        """Check if this node can be expanded."""
        pass

    def get_indent(self) -> int:
        """Get the indentation width for this node."""
        return self.depth * 20


class SimulationNode(TreeNode):
    """Root node representing the entire simulation."""

    def __init__(self, world):
        super().__init__("Simulation")
        self.world = world
        self.is_expanded = True  # Always expanded

    def get_display_text(self) -> str:
        return f"Simulation"

    def can_expand(self) -> bool:
        return True

    def update_entity_counts(self) -> None:
        """Update entity type nodes with current world state."""
        entity_types = {}

        # Count entities by type
        for entity in self.world.get_objects():
            entity_type = type(entity).__name__
            if entity_type not in entity_types:
                entity_types[entity_type] = []
            entity_types[entity_type].append(entity)

        # Update children to match current entity types
        existing_types = {child.label: child for child in self.children}

        # Remove types that no longer exist
        for type_name in existing_types:
            if type_name not in entity_types:
                self.remove_child(existing_types[type_name])

        # Add or update type nodes
        for type_name, entities in entity_types.items():
            if type_name in existing_types:
                # Update existing node
                type_node = existing_types[type_name]
                type_node.update_entities(entities)
            else:
                # Create new type node
                type_node = EntityTypeNode(type_name, entities)
                self.add_child(type_node)


class EntityTypeNode(TreeNode):
    """Node representing a category of entities."""

    def __init__(self, entity_type: str, entities: List[Any]):
        super().__init__(entity_type)
        self.entities = entities
        self._update_children()

    def _update_children(self) -> None:
        """Update child entity nodes to match current entities."""
        existing_ids = {child.entity_id: child for child in self.children}
        current_ids = {id(entity): entity for entity in self.entities}

        # Remove entities that no longer exist
        for entity_id in existing_ids:
            if entity_id not in current_ids:
                self.remove_child(existing_ids[entity_id])

        # Add or update entity nodes
        for entity_id, entity in current_ids.items():
            if entity_id in existing_ids:
                # Update existing node
                entity_node = existing_ids[entity_id]
                entity_node.entity = entity
            else:
                # Create new entity node
                entity_node = EntityNode(entity)
                self.add_child(entity_node)

    def update_entities(self, entities: List[Any]) -> None:
        """Update the entities for this type."""
        self.entities = entities
        self._update_children()

    def get_display_text(self) -> str:
        count = len(self.entities)
        return f"{self.label} ({count})"

    def can_expand(self) -> bool:
        return len(self.entities) > 0


class EntityNode(TreeNode):
    """Node representing an individual entity."""

    def __init__(self, entity: Any):
        # Use entity ID as the initial label, will be updated in get_display_text
        super().__init__("", None)
        self.entity = entity
        self.entity_id = id(entity)
        self.is_leaf = True  # Entity nodes can't have children

    def get_display_text(self) -> str:
        # Start with just ID, designed for extensibility
        return f"Entity {self.entity_id}"

    def can_expand(self) -> bool:
        return False

    def is_leaf(self) -> bool:
        return True

    def get_entity_info(self) -> Dict[str, Any]:
        """Get entity information for display purposes."""
        # Base implementation - can be extended later
        info = {
            'id': self.entity_id,
            'type': type(self.entity).__name__
        }

        # Add position if available
        if hasattr(self.entity, 'position'):
            info['position'] = {
                'x': getattr(self.entity.position, 'x', 0),
                'y': getattr(self.entity.position, 'y', 0)
            }

        return info


class TreeSelectionManager:
    """Manages selection state for tree nodes."""

    def __init__(self):
        self.selected_nodes: List[TreeNode] = []
        self.last_selected_node: Optional[TreeNode] = None

    def select_node(self, node: TreeNode, multi_select: bool = False) -> None:
        """Select a node, optionally with multi-select."""
        if not multi_select:
            # Clear existing selection
            for selected_node in self.selected_nodes:
                selected_node.is_selected = False
            self.selected_nodes.clear()

        if node.is_selected:
            # Deselect if already selected
            node.is_selected = False
            if node in self.selected_nodes:
                self.selected_nodes.remove(node)
            if node == self.last_selected_node:
                self.last_selected_node = None
        else:
            # Select the node
            node.is_selected = True
            self.selected_nodes.append(node)
            self.last_selected_node = node

    def select_range(self, from_node: TreeNode, to_node: TreeNode, all_nodes: List[TreeNode]) -> None:
        """Select a range of nodes between from_node and to_node."""
        if from_node not in all_nodes or to_node not in all_nodes:
            return

        # Clear existing selection
        for selected_node in self.selected_nodes:
            selected_node.is_selected = False
        self.selected_nodes.clear()

        # Find indices
        from_index = all_nodes.index(from_node)
        to_index = all_nodes.index(to_node)

        # Select range
        start = min(from_index, to_index)
        end = max(from_index, to_index)

        for i in range(start, end + 1):
            node = all_nodes[i]
            node.is_selected = True
            self.selected_nodes.append(node)

        self.last_selected_node = to_node

    def clear_selection(self) -> None:
        """Clear all selections."""
        for node in self.selected_nodes:
            node.is_selected = False
        self.selected_nodes.clear()
        self.last_selected_node = None

    def get_selected_entities(self) -> List[Any]:
        """Get the actual entities for selected entity nodes."""
        entities = []
        for node in self.selected_nodes:
            if isinstance(node, EntityNode):
                entities.append(node.entity)
        return entities