import pygame
import maze
import sys

# Initialize Pygame
pygame.init()

# Constants
# CELL_SIZE = 10
# CLOCK_TICK = 20
HIGHLIGHT_COLOR = (0, 255, 255)     # Cyan for highlighted cells
SECONDARY_COLOR = (0, 0, 255)       # Blue for secondary highlighted cells
WALL_COLOR = (0, 0, 0)              # Black for walls
WORKING_WALL_COLOR = (32, 32, 32)   # Grey for walls while working
PATH_COLOR = (255, 255, 255)        # White for paths

class MazeVisualizer:

    def __init__(self, maze: maze.Maze):
        self.maze = maze
        self.screen_width = maze.width * CELL_SIZE
        self.screen_height = maze.height * CELL_SIZE
        self.screen = pygame.display.set_mode(
            (self.screen_width, self.screen_height))
        self.working = False
        pygame.display.set_caption("Maze Generator Visualization")

    def draw_maze(self):
        self.screen.fill(PATH_COLOR)
        for y in range(self.maze.height):
            for x in range(self.maze.width):
                if self.maze[x, y]:
                    pygame.draw.rect(
                        self.screen,
                        WALL_COLOR if self.working else WORKING_WALL_COLOR,
                        pygame.Rect(
                            x * CELL_SIZE, y * CELL_SIZE, CELL_SIZE, CELL_SIZE)
                    )
                    if not self.maze.secondary((x, y)):
                        pygame.draw.rect(
                            self.screen,
                            SECONDARY_COLOR,
                            pygame.Rect(
                                x * CELL_SIZE,
                                y * CELL_SIZE,
                                CELL_SIZE,
                                CELL_SIZE
                            )
                        )

        if self.maze.highlighted:
            hx, hy = self.maze.highlighted
            pygame.draw.rect(
                self.screen,
                HIGHLIGHT_COLOR,
                pygame.Rect(
                    hx * CELL_SIZE, hy * CELL_SIZE, CELL_SIZE, CELL_SIZE)
            )

    def run(self):
        clock = pygame.time.Clock()
        running = True
        self.working = True

        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                if event.type == pygame.KEYUP:
                    if event.key == pygame.K_SPACE:
                        self.working = not self.working
                    if event.key == pygame.K_o:
                        if isinstance(self.maze, maze.VectorWrapper):
                            self.maze = maze.VectorWrapper(
                                maze.OriginShift.clone(self.maze.maze))
                    if event.key == pygame.K_s:
                        if self.working: self.maze.step()
                    if event.key == pygame.K_h:
                        for _ in range(100):
                            if not self.maze.step(): break
                    if event.key == pygame.K_t:
                        for _ in range(1000):
                            if not self.maze.step(): break
                    if event.key == pygame.K_m:
                        for _ in range(10000):
                            if not self.maze.step(): break
                    if event.key == pygame.K_v:
                        self.maze = maze.VectorWrapper.convert(self.maze)

            # Only step through the algorithm if it's not finished
            if not self.working:
                self.working = not self.maze.step()

            # Draw the maze regardless of generation state
            self.draw_maze()
            pygame.display.flip()
            clock.tick(CLOCK_TICK)

        pygame.quit()
        sys.exit()

CELL_SIZE = 5
CLOCK_TICK = 20

MAZE_SIZE = 127
VEC_MAZE_SIZE = (MAZE_SIZE + 1) // 2

# my_maze = maze.RecursiveBacktracker(MAZE_SIZE)
# my_maze = maze.VectorWrapper(maze.OriginShift(VEC_MAZE_SIZE))
# my_maze = maze.VectorWrapper(maze.BinaryTree(
#     VEC_MAZE_SIZE, bias=maze.VectorEnum.Up | maze.VectorEnum.Left))
# my_maze = maze.VectorWrapper(maze.BinaryTree(
#     VEC_MAZE_SIZE, bias=maze.VectorEnum.Up | maze.VectorEnum.Right))
# my_maze = maze.VectorWrapper(maze.BinaryTree(
#     VEC_MAZE_SIZE, bias=maze.VectorEnum.Down | maze.VectorEnum.Left))
# my_maze = maze.VectorWrapper(maze.BinaryTree(
#     VEC_MAZE_SIZE, bias=maze.VectorEnum.Down | maze.VectorEnum.Right))
# my_maze = maze.Sidewinder(MAZE_SIZE, run_param=0.2)
# my_maze = maze.Sidewinder(MAZE_SIZE, run_param=0.5)
# my_maze = maze.Sidewinder(MAZE_SIZE, run_param=0.8)
# my_maze = maze.RecursiveDivision(MAZE_SIZE)
# my_maze = maze.RecursiveDivision(MAZE_SIZE, uniform=False)
# my_maze = maze.RecursiveDivision(MAZE_SIZE, depth_first=False)
# my_maze = maze.RecursiveDivision(MAZE_SIZE, binary=True)
# my_maze = maze.VectorWrapper(maze.Prim(VEC_MAZE_SIZE))
# my_maze = maze.VectorWrapper(maze.Wilson(VEC_MAZE_SIZE))
# my_maze = maze.VectorWrapper(maze.AldousBroder(VEC_MAZE_SIZE))
my_maze = maze.VectorWrapper(maze.AldousBroderWilson(VEC_MAZE_SIZE))

# for _ in range(512): my_maze.step()
# for _ in range(2048): my_maze.step()
# for _ in range(16384): my_maze.step()

visualizer = MazeVisualizer(my_maze)
visualizer.run()