import pygame
import sys
import random
import time

# Init Pygame because we’re not living in the stone age
pygame.init()

# Constants that will never change except when they do
WIDTH, HEIGHT = 640, 480
# Block size: how many pixels per grid cell
BLOCK_SIZE = 20
# FPS: frames per second, aka how smooth the snake is
SPEED = 10
# Direction mapping: because string comparison is slower than integer comparison
DIRECTIONS = {'UP': (0, -1), 'DOWN': (0, 1), 'LEFT': (-1, 0), 'RIGHT': (1, 0)}

# Colors because RGB values are more readable than magic numbers
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GREEN = (0, 255, 0)
# Red for apple and also for “game over” messages
RED = (255, 0, 0)

class SnakeGame:
    def __init__(self):
        self.display = pygame.display.set_mode((WIDTH, HEIGHT))
        pygame.display.set_caption("Snake Game 2000 Edition")  # Title is kino
        self.clock = pygame.time.Clock()
        self.reset()

    def reset(self):
        # Reset the snake to the starting position
        self.snake = [(200, 200), (220, 200), (240, 200)]
        # Set initial direction to right because why the fuck not
        self.direction = 'RIGHT'
        # Generate an apple somewhere on the grid
        self.apple = self.generate_apple()
        # Keep track of score
        self.score = 0
        # Game over flag
        self.game_over = False

    def generate_apple(self):
        # Make sure the apple doesn't spawn where the snake already is
        while True:
            x = random.randint(0, WIDTH - BLOCK_SIZE) // BLOCK_SIZE * BLOCK_SIZE
            y = random.randint(0, HEIGHT - BLOCK_SIZE) // BLOCK_SIZE * BLOCK_SIZE
            apple = (x, y)
            if apple not in self.snake:
                return apple

    def update(self):
        if self.game_over:
            # If game is over, print final score like a based champion
            print(f"\nGame over anon! Final score: {self.score}")
            print("You died like a virgin. Press R to try again.")
            return

        # Get the head of the snake
        head = self.snake[-1]
        # Calculate new head position
        new_head = (head[0] + DIRECTIONS[self.direction][0] * BLOCK_SIZE,
                    head[1] + DIRECTIONS[self.direction][1] * BLOCK_SIZE)

        # Check for collisions with walls and self
        if ((new_head[0] < 0 or new_head[0] >= WIDTH or
             new_head[1] < 0 or new_head[1] >= HEIGHT) or
            new_head in self.snake):
            # If you hit a wall or yourself, game over
            self.game_over = True
            self.draw()  # Draw final frame
            pygame.time.delay(2000)  # Wait 2 seconds
            return

        # Move the snake
        self.snake.append(new_head)

        # Check if you ate the apple
        if self.apple == new_head:
            # Increase score and spawn new apple
            self.score += 1
            print(f"🍎 Munch! Score: {self.score}")  # Emoji for based feedback
            self.apple = self.generate_apple()
        else:
            # Otherwise, remove the tail
            self.snake.pop(0)

    def draw(self):
        # Fill screen with black like your soul
        self.display.fill(BLACK)

        # Draw snake
        for pos in self.snake:
            pygame.draw.rect(self.display, GREEN, pygame.Rect(pos[0], pos[1], BLOCK_SIZE, BLOCK_SIZE))

        # Draw apple
        pygame.draw.rect(self.display, RED, pygame.Rect(self.apple[0], self.apple[1], BLOCK_SIZE, BLOCK_SIZE))

        # Draw score
        font = pygame.font.Font(None, 36)
        text = font.render(f"Score: {self.score}", True, WHITE)
        self.display.blit(text, (10, 10))

        # Check if game over and draw “Game Over” text
        if self.game_over:
            text = font.render("Game Over. Press R to restart.", True, RED)
            self.display.blit(text, (10, 50))

        # Update display like a good boy
        pygame.display.flip()

    def handle_input(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                # Exit button pressed, clean up and quit
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                # Handle arrow key presses
                if event.key == pygame.K_UP and self.direction != 'DOWN':
                    self.direction = 'UP'
                elif event.key == pygame.K_DOWN and self.direction != 'UP':
                    self.direction = 'DOWN'
                elif event.key == pygame.K_LEFT and self.direction != 'RIGHT':
                    self.direction = 'LEFT'
                elif event.key == pygame.K_RIGHT and self.direction != 'LEFT':
                    self.direction = 'RIGHT'
                elif event.key == pygame.K_r:
                    # Reset game
                    self.reset()
                    self.game_over = False  # Clear game over state
                elif event.key == pygame.K_s:
                    # Change speed (because speedrunners will always exist)
                    try:
                        speed = int(input("Enter new speed: "))
                        if speed < 1:
                            print("Speed must be 1 or higher, you absolute mongoloid")
                        else:
                            SPEED = speed
                            print(f"Speed set to {speed}. Your grandma can handle this.")
                    except ValueError:
                        print("Invalid input. Default speed remains.")

                # Bonus feature: press S to show current speed
                elif event.key == pygame.K_c:
                    print(f"Current speed: {SPEED}")

    def run(self):
        # Game loop
        while True:
            self.handle_input()

            if self.game_over:
                self.draw()
                # Wait 2 seconds before restarting
                pygame.time.delay(2000)
                break

            self.update()
            self.draw()
            # Cap framerate at SPEED
            self.clock.tick(SPEED)

if __name__ == "__main__":
    game = SnakeGame()
    print("Welcome to Snake Game 2000 Edition!")
    print("Use arrow keys to move, R to restart, S to change speed")
    print(f"Current speed: {SPEED}")
    game.run()