import pygame
import sys
import random

# Initialize Pygame
pygame.init()

# Constants
WIDTH, HEIGHT = 600, 400
CELL_SIZE = 20
SNAKE_SIZE = 20
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)

# Snake Class
class Snake:
    def __init__(self):
        self.length = 1
        self.positions = [((WIDTH // 2), (HEIGHT // 2))]
        self.direction = random.choice([0, 1, 2, 3])  # 0: UP, 1: RIGHT, 2: DOWN, 3: LEFT
        self.head_color = GREEN
        self.body_color = WHITE

    def get_head_position(self):
        return self.positions[0]

    def update(self):
        cur = self.get_head_position()
        x, y = self.direction_to_coord(self.direction)

        new = (((cur[0] + (x * CELL_SIZE)) % WIDTH), (cur[1] + (y * CELL_SIZE)) % HEIGHT)

        self.positions.insert(0, new)
        if len(self.positions) > self.length:
            self.positions.pop()

    def render(self, surface):
        for i, p in enumerate(self.positions):
            color = self.head_color if i == 0 else self.body_color
            pygame.draw.rect(surface, color, (p[0], p[1], CELL_SIZE, CELL_SIZE))

    def direction_to_coord(self, direction):
        if direction == 0:
            return 0, -1  # UP
        elif direction == 1:
            return 1, 0  # RIGHT
        elif direction == 2:
            return 0, 1  # DOWN
        elif direction == 3:
            return -1, 0  # LEFT

# Food Class
class Food:
    def __init__(self):
        self.position = (0, 0)
        self.color = RED
        self.randomize_position()

    def randomize_position(self):
        self.position = (random.randint(0, (WIDTH // CELL_SIZE) - 1) * CELL_SIZE,
                         random.randint(0, (HEIGHT // CELL_SIZE) - 1) * CELL_SIZE)

    def render(self, surface):
        pygame.draw.rect(surface, self.color, (self.position[0], self.position[1], CELL_SIZE, CELL_SIZE))

# Initialize game window
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Snake")
clock = pygame.time.Clock()

snake = Snake()
food = Food()

# Game loop
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP and not snake.direction == 2:
                snake.direction = 0
            elif event.key == pygame.K_DOWN and not snake.direction == 0:
                snake.direction = 2
            elif event.key == pygame.K_LEFT and not snake.direction == 1:
                snake.direction = 3
            elif event.key == pygame.K_RIGHT and not snake.direction == 3:
                snake.direction = 1

    snake.update()

    # Check if snake eats food
    if snake.get_head_position() == food.position:
        snake.length += 5
        food.randomize_position()

    # Draw everything
    screen.fill((0, 0, 0))

    snake.render(screen)
    food.render(screen)

    pygame.display.flip()

    # Cap the frame rate
    clock.tick(10)

# Quit the game
pygame.quit()
sys.exit()