import pygame
import random

# Initialize the game
pygame.init()

# Set the screen dimensions
screen_width = 800
screen_height = 600
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption('Snake Game')

# Define colors
black = pygame.Color(0, 0, 0)
white = pygame.Color(255, 255, 255)
red = pygame.Color(255, 0, 0)
green = pygame.Color(0, 255, 0)

# Set the snake attributes
snake_pos = [[100, 50], [90, 50], [80, 50]]
snake_body = [100, 50]
snake_speed = 10
direction = "RIGHT"

# Set the food attributes
food_pos = [random.randrange(1, (screen_width//10)) * 10,
            random.randrange(1, (screen_height//10)) * 10]
food_spawn = True

# Set the game clock
clock = pygame.time.Clock()

# Function to display the snake
def show_snake(snake_body):
    for pos in snake_body:
        pygame.draw.rect(screen, green, pygame.Rect(
            pos[0], pos[1], 10, 10))

# Function to display the food
def show_food(food_pos):
    pygame.draw.rect(screen, red, pygame.Rect(
        food_pos[0], food_pos[1], 10, 10))

# Handle the game over scenario
def game_over():
    font = pygame.font.Font(None, 36)
    text = font.render("Game Over!", True, white)
    screen.blit(text, (screen_width/2 - text.get_width()/2,
                       screen_height/2 - text.get_height()/2))
    pygame.display.flip()
    pygame.time.wait(2000)
    pygame.quit()

# Start the game loop
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            exit(0)
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_RIGHT:
                direction = "RIGHT"
            elif event.key == pygame.K_LEFT:
                direction = "LEFT"
            elif event.key == pygame.K_UP:
                direction = "UP"
            elif event.key == pygame.K_DOWN:
                direction = "DOWN"

    # Update the snake position based on the direction
    if direction == "RIGHT":
        snake_body[0] += snake_speed
    elif direction == "LEFT":
        snake_body[0] -= snake_speed
    elif direction == "UP":
        snake_body[1] -= snake_speed
    elif direction == "DOWN":
        snake_body[1] += snake_speed

    # Insert new snake position at the front
    snake_pos.insert(0, list(snake_body))

    # Check if the snake collided with the food
    if snake_body[0] == food_pos[0] and snake_body[1] == food_pos[1]:
        food_spawn = False
    else:
        snake_pos.pop()

    # Respawn the food at a random position if eaten
    if not food_spawn:
        food_pos = [random.randrange(1, (screen_width//10)) * 10,
                    random.randrange(1, (screen_height//10)) * 10]
    food_spawn = True

    # Fill the screen with black color
    screen.fill(black)

    # Draw the snake and food
    show_snake(snake_pos)
    show_food(food_pos)

    # Game over conditions
    if snake_body[0] < 0 or snake_body[0] > screen_width-10:
        game_over()
    if snake_body[1] < 0 or snake_body[1] > screen_height-10:
        game_over()
    for block in snake_pos[1:]:
        if snake_body[0] == block[0] and snake_body[1] == block[1]:
            game_over()

    # Update the display
    pygame.display.flip()

    # Set the game speed
    clock.tick(30)