import pygame
from enum import Enum, auto
from typing import List, Tuple, Callable

class GameMode(Enum):
    PLAYER = auto()
    AI_EASY = auto()
    AI_MEDIUM = auto()
    AI_HARD = auto()
    SETTINGS = auto()
    QUIT = auto()

class MenuItem:
    def __init__(self, text, position, action, size=36, color=(255, 255, 255)):
        self.text = text
        self.position = position
        self.action = action
        self.size = size
        self.default_color = color
        self.color = color
        self.hover = False
        self._setup_font()
    
    def _setup_font(self):
        """Update the font surface with current color"""
        self.font = pygame.font.Font(None, self.size)
        color = (255, 255, 0) if self.hover else self.default_color  # Yellow when hovered
        self.surface = self.font.render(self.text, True, color)
        self.rect = self.surface.get_rect(center=self.position)
    
    def update_hover(self, mouse_pos):
        """Update hover state based on mouse position"""
        old_hover = self.hover
        self.hover = self.rect.collidepoint(mouse_pos)
        if old_hover != self.hover:
            self._setup_font()
    
    def draw(self, screen):
        """Draw the menu item on the screen"""
        screen.blit(self.surface, self.rect)

class Menu:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.setup_menu_items()
        self.title_font = pygame.font.Font(None, 72)
        self.subtitle_font = pygame.font.Font(None, 36)
        
        # Create title surfaces
        self.title_surface = self.title_font.render("AI Snake Game", True, (0, 255, 0))
        self.title_rect = self.title_surface.get_rect(center=(width//2, height//4))
        
        self.subtitle_surface = self.subtitle_font.render("Choose Game Mode", True, (200, 200, 200))
        self.subtitle_rect = self.subtitle_surface.get_rect(center=(width//2, height//4 + 50))
        
        # Initialize first item as selected
        self.selected_index = 0
        self.menu_items[0].hover = True
        self.menu_items[0]._setup_font()
    
    def setup_menu_items(self):
        # Calculate positions for menu items
        start_y = self.height // 2
        spacing = 50
        center_x = self.width // 2
        
        self.menu_items = [
            MenuItem("Player Mode", (center_x, start_y), GameMode.PLAYER),
            MenuItem("AI Easy", (center_x, start_y + spacing), GameMode.AI_EASY),
            MenuItem("AI Medium", (center_x, start_y + spacing * 2), GameMode.AI_MEDIUM),
            MenuItem("AI Hard", (center_x, start_y + spacing * 3), GameMode.AI_HARD),
            MenuItem("Settings", (center_x, start_y + spacing * 4), GameMode.SETTINGS),
            MenuItem("Quit", (center_x, start_y + spacing * 5), GameMode.QUIT, color=(255, 100, 100))
        ]
    
    def handle_input(self, event):
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            # Handle mouse clicks
            mouse_pos = pygame.mouse.get_pos()
            for i, item in enumerate(self.menu_items):
                if item.rect.collidepoint(mouse_pos):
                    self.selected_index = i
                    return item.action
        
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_RETURN:
                # Return the selected item's action
                return self.menu_items[self.selected_index].action
            
            elif event.key in (pygame.K_UP, pygame.K_DOWN):
                # Update selected index
                if event.key == pygame.K_UP:
                    self.selected_index = (self.selected_index - 1) % len(self.menu_items)
                else:
                    self.selected_index = (self.selected_index + 1) % len(self.menu_items)
                
                # Update hover states
                for i, item in enumerate(self.menu_items):
                    item.hover = (i == self.selected_index)
                    item._setup_font()
        
        return None
    
    def update(self):
        # Handle mouse hover
        mouse_pos = pygame.mouse.get_pos()
        for i, item in enumerate(self.menu_items):
            if item.rect.collidepoint(mouse_pos):
                # Update selected index when mouse hovers
                self.selected_index = i
                item.hover = True
                item._setup_font()
            else:
                # Keep keyboard selection visible
                item.hover = (i == self.selected_index)
                item._setup_font()
    
    def draw(self, screen):
        # Draw background
        screen.fill((0, 0, 0))
        
        # Draw title and subtitle
        screen.blit(self.title_surface, self.title_rect)
        screen.blit(self.subtitle_surface, self.subtitle_rect)
        
        # Draw menu items
        for item in self.menu_items:
            item.draw(screen)
        
        # Draw version and controls
        version_text = "v0.1.0"
        controls_text = "Arrow keys or mouse to navigate, Enter to select"
        font = pygame.font.Font(None, 24)
        
        version_surface = font.render(version_text, True, (100, 100, 100))
        controls_surface = font.render(controls_text, True, (100, 100, 100))
        
        screen.blit(version_surface, (10, self.height - 30))
        screen.blit(controls_surface, (self.width - controls_surface.get_width() - 10, self.height - 30))