import pygame
import random

# Initialize Pygame
pygame.init()

# Define colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GREEN = (0, 255, 0)
RED = (255, 0, 0)

# Set up the screen
WIDTH, HEIGHT = 800, 800
CELL_SIZE = 40
GRID_WIDTH, GRID_HEIGHT = WIDTH // CELL_SIZE, HEIGHT // CELL_SIZE
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Snake")

# Define directions
UP = (0, -1)
DOWN = (0, 1)
LEFT = (-1, 0)
RIGHT = (1, 0)

# Snake class
class Snake:
    def __init__(self):
        self.body = [(GRID_WIDTH // 2, GRID_HEIGHT // 2)]
        self.direction = RIGHT

    def move(self):
        head = self.body[0]
        new_head = (head[0] + self.direction[0], head[1] + self.direction[1])
        self.body.insert(0, new_head)
        if new_head != fruit.position:
            self.body.pop()
        else:
            fruit.spawn()

    def grow(self):
        tail = self.body[-1]
        self.body.append(tail)

    def change_direction(self, direction):
        if direction == UP and self.direction != DOWN:
            self.direction = direction
        elif direction == DOWN and self.direction != UP:
            self.direction = direction
        elif direction == LEFT and self.direction != RIGHT:
            self.direction = direction
        elif direction == RIGHT and self.direction != LEFT:
            self.direction = direction

    def draw(self):
        for segment in self.body:
            x, y = segment
            pygame.draw.rect(screen, GREEN, (x * CELL_SIZE, y * CELL_SIZE, CELL_SIZE, CELL_SIZE))

# Fruit class
class Fruit:
    def __init__(self):
        self.position = self.randomize_position()

    def randomize_position(self):
        return random.randrange(GRID_WIDTH), random.randrange(GRID_HEIGHT)

    def spawn(self):
        self.position = self.randomize_position()

    def draw(self):
        x, y = self.position
        pygame.draw.rect(screen, RED, (x * CELL_SIZE, y * CELL_SIZE, CELL_SIZE, CELL_SIZE))

# Main game loop
snake = Snake()
fruit = Fruit()
clock = pygame.time.Clock()
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:
                snake.change_direction(UP)
            elif event.key == pygame.K_DOWN:
                snake.change_direction(DOWN)
            elif event.key == pygame.K_LEFT:
                snake.change_direction(LEFT)
            elif event.key == pygame.K_RIGHT:
                snake.change_direction(RIGHT)

    snake.move()

    if (snake.body[0][0] < 0 or snake.body[0][0] >= GRID_WIDTH or
            snake.body[0][1] < 0 or snake.body[0][1] >= GRID_HEIGHT or
            snake.body[0] in snake.body[1:]):
        print("Game Over!")
        running = False

    screen.fill(BLACK)
    snake.draw()
    fruit.draw()
    pygame.display.flip()
    clock.tick(10)  # Adjust snake speed

pygame.quit()