import pygame
import sys
from enum import Enum, auto
from snake import Snake, Direction
from food import Food

class GameState(Enum):
    MENU = auto()
    PLAYING = auto()
    GAME_OVER = auto()
    PAUSED = auto()

class Game:
    def __init__(self):
        # Initialize display
        self.width = 800
        self.height = 600
        self.screen = pygame.display.set_mode((self.width, self.height))
        pygame.display.set_caption("AI Snake Game")
        
        # Initialize clock
        self.clock = pygame.time.Clock()
        self.fps = 60
        
        # Game objects
        self.block_size = 20
        self.snake = Snake((self.width // 2, self.height // 2), self.block_size)
        self.food = Food(self.block_size)
        
        # Spawn initial food
        self.food.spawn(self.width, self.height, self.snake.body)
        
        # Game state
        self.state = GameState.PLAYING  # Changed to start in playing state for now
        self.running = True
        self.score = 0

    def handle_events(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                self.running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    if self.state == GameState.PLAYING:
                        self.state = GameState.PAUSED
                    elif self.state == GameState.PAUSED:
                        self.state = GameState.PLAYING
                elif self.state == GameState.PLAYING:
                    # Handle snake direction
                    if event.key == pygame.K_UP:
                        self.snake.change_direction(Direction.UP)
                    elif event.key == pygame.K_DOWN:
                        self.snake.change_direction(Direction.DOWN)
                    elif event.key == pygame.K_LEFT:
                        self.snake.change_direction(Direction.LEFT)
                    elif event.key == pygame.K_RIGHT:
                        self.snake.change_direction(Direction.RIGHT)

    def update(self):
        if self.state == GameState.PLAYING:
            # Move snake
            if self.snake.move(pygame.time.get_ticks()):
                # Check wall collision
                if self.snake.check_collision(self.width, self.height):
                    self.state = GameState.GAME_OVER
                    return
                
                # Check food collision
                if self.food.check_collision(self.snake.body[0]):
                    self.snake.grow()
                    self.score += 1
                    self.food.spawn(self.width, self.height, self.snake.body)

    def render(self):
        # Clear screen
        self.screen.fill((0, 0, 0))  # Black background
        
        if self.state == GameState.PLAYING or self.state == GameState.PAUSED:
            # Draw game objects
            self.snake.draw(self.screen)
            self.food.draw(self.screen)
            
            # Draw score
            font = pygame.font.Font(None, 36)
            score_text = font.render(f'Score: {self.score}', True, (255, 255, 255))
            self.screen.blit(score_text, (10, 10))
            
            # Draw pause indicator
            if self.state == GameState.PAUSED:
                pause_text = font.render('PAUSED', True, (255, 255, 255))
                text_rect = pause_text.get_rect(center=(self.width//2, self.height//2))
                self.screen.blit(pause_text, text_rect)
        
        elif self.state == GameState.GAME_OVER:
            font = pygame.font.Font(None, 74)
            text = font.render('Game Over!', True, (255, 0, 0))
            score_text = font.render(f'Score: {self.score}', True, (255, 255, 255))
            
            text_rect = text.get_rect(center=(self.width//2, self.height//2 - 50))
            score_rect = score_text.get_rect(center=(self.width//2, self.height//2 + 50))
            
            self.screen.blit(text, text_rect)
            self.screen.blit(score_text, score_rect)
        
        # Update display
        pygame.display.flip()

    def run(self):
        while self.running:
            self.handle_events()
            self.update()
            self.render()
            self.clock.tick(self.fps)