"""Kruskal's Algorithm for maze generation."""

import random
from typing import Dict, List, Tuple

from ..core.cell import Cell
from ..core.maze import Maze
from .base import BaseGenerator


class KruskalGenerator(BaseGenerator):
    """Generates mazes using Kruskal's Algorithm.

    This algorithm treats the maze as a graph and uses a union-find structure
    to create a minimum spanning tree. Creates mazes with many short paths.

    Time Complexity: O(E log E) where E is number of edges
    Space Complexity: O(V) where V is number of vertices
    """

    def __init__(self):
        """Initialize the Kruskal generator."""
        super().__init__("Kruskal's Algorithm")

    def _generate_maze(self, maze: Maze) -> None:
        """Generate maze using Kruskal's algorithm.

        Args:
            maze: Maze instance to generate
        """
        # Initialize union-find structure
        parent: Dict[Cell, Cell] = {}
        for row in maze.grid:
            for cell in row:
                parent[cell] = cell

        # Create list of all possible walls (edges)
        walls = self._get_all_walls(maze)
        random.shuffle(walls)

        # Process each wall
        for cell1, cell2 in walls:
            # Find roots of both cells
            root1 = self._find(parent, cell1)
            root2 = self._find(parent, cell2)

            # If cells are in different sets, remove wall and union
            if root1 != root2:
                maze.remove_wall_between(cell1, cell2)
                parent[root2] = root1

    def _get_all_walls(self, maze: Maze) -> List[Tuple[Cell, Cell]]:
        """Get all possible walls between cells.

        Args:
            maze: The maze

        Returns:
            List of (cell1, cell2) tuples representing walls
        """
        walls = []

        for row in maze.grid:
            for cell in row:
                # Add wall to the south
                south_cell = maze.get_cell(cell.row + 1, cell.col)
                if south_cell:
                    walls.append((cell, south_cell))

                # Add wall to the east
                east_cell = maze.get_cell(cell.row, cell.col + 1)
                if east_cell:
                    walls.append((cell, east_cell))

        return walls

    def _find(self, parent: Dict[Cell, Cell], cell: Cell) -> Cell:
        """Find the root of a cell's set with path compression.

        Args:
            parent: Union-find parent dictionary
            cell: Cell to find root for

        Returns:
            Root cell of the set
        """
        if parent[cell] != cell:
            parent[cell] = self._find(parent, parent[cell])
        return parent[cell]