import pygame
import random

# Initialize Pygame
pygame.init()

# Constants
SCREEN_WIDTH = 600
SCREEN_HEIGHT = 600
BLOCK_SIZE = 20
FPS = 15
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)

# Screen setup
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Snake Maze Game")

# Clock to control the game speed
clock = pygame.time.Clock()

# Snake class
class Snake:
    def __init__(self):
        self.body = [(100, 100), (80, 100), (60, 100)]  # Snake starts here
        self.direction = (BLOCK_SIZE, 0)  # Snake is moving right initially

    def move(self):
        head_x, head_y = self.body[0]
        new_head = (head_x + self.direction[0], head_y + self.direction[1])
        self.body = [new_head] + self.body[:-1]  # Move the snake

    def grow(self):
        tail_x, tail_y = self.body[-1]
        self.body.append((tail_x, tail_y))  # Add new segment at the tail

    def draw(self):
        for segment in self.body:
            pygame.draw.rect(screen, GREEN, (segment[0], segment[1], BLOCK_SIZE, BLOCK_SIZE))

# Apple class
class Apple:
    def __init__(self):
        self.position = self.random_position()

    def random_position(self):
        return (random.randint(0, (SCREEN_WIDTH - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE,
                random.randint(0, (SCREEN_HEIGHT - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE)

    def draw(self):
        pygame.draw.rect(screen, RED, (self.position[0], self.position[1], BLOCK_SIZE, BLOCK_SIZE))

# Maze class (static walls and obstacles)
class Maze:
    def __init__(self):
        self.walls = [
            pygame.Rect(0, 0, SCREEN_WIDTH, BLOCK_SIZE),  # Top wall
            pygame.Rect(0, 0, BLOCK_SIZE, SCREEN_HEIGHT),  # Left wall
            pygame.Rect(SCREEN_WIDTH - BLOCK_SIZE, 0, BLOCK_SIZE, SCREEN_HEIGHT),  # Right wall
            pygame.Rect(0, SCREEN_HEIGHT - BLOCK_SIZE, SCREEN_WIDTH, BLOCK_SIZE),  # Bottom wall
            pygame.Rect(100, 100, 200, BLOCK_SIZE),  # Horizontal wall in the middle
            pygame.Rect(100, 100, BLOCK_SIZE, 200),  # Vertical wall in the middle
        ]

    def draw(self):
        for wall in self.walls:
            pygame.draw.rect(screen, BLUE, wall)

    def check_collision(self, rect):
        for wall in self.walls:
            if rect.colliderect(wall):
                return True
        return False

# Collision with Snake itself
def check_self_collision(snake):
    head = snake.body[0]
    if head in snake.body[1:]:  # Snake collides with itself
        return True
    return False

# Main game loop
def game_loop():
    snake = Snake()
    apple = Apple()
    maze = Maze()

    running = True
    game_over = False

    while running:
        screen.fill(BLACK)

        # Handle events
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_LEFT:
                    snake.direction = (-BLOCK_SIZE, 0)
                elif event.key == pygame.K_RIGHT:
                    snake.direction = (BLOCK_SIZE, 0)
                elif event.key == pygame.K_UP:
                    snake.direction = (0, -BLOCK_SIZE)
                elif event.key == pygame.K_DOWN:
                    snake.direction = (0, BLOCK_SIZE)

        if game_over:
            font = pygame.font.SysFont(None, 35)
            game_over_text = font.render("Game Over! Press Q to Quit or R to Restart", True, WHITE)
            screen.blit(game_over_text, (SCREEN_WIDTH // 4, SCREEN_HEIGHT // 2))
            pygame.display.update()

            # Wait for user input to quit or restart
            waiting_for_input = True
            while waiting_for_input:
                for event in pygame.event.get():
                    if event.type == pygame.QUIT:
                        running = False
                        waiting_for_input = False
                    if event.type == pygame.KEYDOWN:
                        if event.key == pygame.K_q:
                            running = False
                            waiting_for_input = False
                        if event.key == pygame.K_r:
                            game_loop()  # Restart the game
                            waiting_for_input = False
            continue

        # Move the snake
        snake.move()

        # Check if snake collides with walls or itself
        head_rect = pygame.Rect(snake.body[0][0], snake.body[0][1], BLOCK_SIZE, BLOCK_SIZE)
        if maze.check_collision(head_rect) or check_self_collision(snake):
            game_over = True

        # Check if snake eats an apple
        if head_rect.colliderect(pygame.Rect(apple.position[0], apple.position[1], BLOCK_SIZE, BLOCK_SIZE)):
            snake.grow()
            apple = Apple()  # Generate new apple

        # Draw the maze, snake, and apple
        maze.draw()
        snake.draw()
        apple.draw()

        # Display score
        font = pygame.font.SysFont(None, 35)
        score_text = font.render(f"Score: {len(snake.body) - 3}", True, WHITE)
        screen.blit(score_text, (10, 10))

        # Update the display
        pygame.display.update()

        # Control the game speed
        clock.tick(FPS)

    pygame.quit()

# Run the game
game_loop()