import pygame
import time
import random

# Define the colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (213, 50, 80)
GREEN = (0, 255, 0)
BLUE = (50, 153, 213)

# Set the width and height of each segment of the snake
SEGMENT_WIDTH = 20
SEGMENT_HEIGHT = 20
# Set the margin between each segment of the snake
SEGMENT_MARGIN = 2

# Set the width and height of the screen [width, height]
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600

# Set initial speed of the snake
SNAKE_SPEED = 10

class Segment(pygame.sprite.Sprite):
    """ Class to represent one segment of the snake """

    def __init__(self, x, y):
        # Call the parent's constructor
        super().__init__()

        # Set the segment's dimensions
        self.image = pygame.Surface([SEGMENT_WIDTH, SEGMENT_HEIGHT])
        self.image.fill(WHITE)

        # Set the position of the segment
        self.rect = self.image.get_rect()
        self.rect.x = x
        self.rect.y = y

class Snake:
    """ Class to represent the snake """

    def __init__(self):
        # Create the initial snake with 3 segments
        self.segments = []
        self.length = 3

        for i in range(self.length):
            segment = Segment(250 - (SEGMENT_WIDTH + SEGMENT_MARGIN) * i, 30)
            self.segments.append(segment)

        # Set the initial direction of the snake
        self.direction = "RIGHT"

    def move(self):
        """ Move the snake """
        # Move each segment of the snake to the position of the segment in front of it
        for i in range(len(self.segments) - 1, 0, -1):
            self.segments[i].rect.x = self.segments[i - 1].rect.x
            self.segments[i].rect.y = self.segments[i - 1].rect.y

        # Move the head of the snake in the direction it is facing
        if self.direction == "UP":
            self.segments[0].rect.y -= SEGMENT_HEIGHT + SEGMENT_MARGIN
        elif self.direction == "DOWN":
            self.segments[0].rect.y += SEGMENT_HEIGHT + SEGMENT_MARGIN
        elif self.direction == "LEFT":
            self.segments[0].rect.x -= SEGMENT_WIDTH + SEGMENT_MARGIN
        elif self.direction == "RIGHT":
            self.segments[0].rect.x += SEGMENT_WIDTH + SEGMENT_MARGIN

    def change_direction(self, direction):
        """ Change the direction of the snake """
        if direction == "UP" and self.direction != "DOWN":
            self.direction = "UP"
        elif direction == "DOWN" and self.direction != "UP":
            self.direction = "DOWN"
        elif direction == "LEFT" and self.direction != "RIGHT":
            self.direction = "LEFT"
        elif direction == "RIGHT" and self.direction != "LEFT":
            self.direction = "RIGHT"

    def grow(self):
        """ Add segments to the snake """
        for _ in range(3):
            last_segment = self.segments[-1]
            new_segment = Segment(last_segment.rect.x, last_segment.rect.y)
            self.segments.append(new_segment)

class Food(pygame.sprite.Sprite):
    """ Class to represent the food """

    def __init__(self):
        # Call the parent's constructor
        super().__init__()

        # Set the food's dimensions and color
        self.image = pygame.Surface([SEGMENT_WIDTH, SEGMENT_HEIGHT])
        self.image.fill(RED)

        # Set the position of the food
        self.rect = self.image.get_rect()
        self.spawn()

    def spawn(self):
        """ Randomly spawn the food on the screen """
        self.rect.x = random.randrange(0, SCREEN_WIDTH - SEGMENT_WIDTH, SEGMENT_WIDTH + SEGMENT_MARGIN)
        self.rect.y = random.randrange(0, SCREEN_HEIGHT - SEGMENT_HEIGHT, SEGMENT_HEIGHT + SEGMENT_MARGIN)

def main():
    """ Main function for the game """

    pygame.init()

    # Set the screen dimensions
    screen = pygame.display.set_mode([SCREEN_WIDTH, SCREEN_HEIGHT])
    pygame.display.set_caption("stefan is a retard")

    # Create the snake and food
    snake = Snake()
    food = Food()

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

    # Initialize game variables
    game_over = False
    score = 0

    while not game_over:
        # Process events
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                game_over = True
            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")

        # Move the snake
        snake.move()

        # Check if the snake has collided with itself
        for segment in snake.segments[1:]:
            if pygame.sprite.collide_rect(snake.segments[0], segment):
                game_over = True

        # Check if the snake has collided with the walls
        if (snake.segments[0].rect.x < 0 or
            snake.segments[0].rect.x >= SCREEN_WIDTH or
            snake.segments[0].rect.y < 0 or
            snake.segments[0].rect.y >= SCREEN_HEIGHT):
            game_over = True

        # Check if the snake has eaten the food
        if pygame.sprite.collide_rect(snake.segments[0], food):
            snake.grow()
            food.spawn()
            score += 1

        # Clear the screen
        screen.fill(BLACK)

        # Draw the snake
        for segment in snake.segments:
            pygame.draw.rect(screen, GREEN, segment.rect)

        # Draw the food
        pygame.draw.rect(screen, RED, food.rect)

        # Display score
        font = pygame.font.SysFont(None, 25)
        text = font.render("Score: " + str(score), True, WHITE)
        screen.blit(text, [10, 10])

        # Update the screen
        pygame.display.flip()

        # Cap the frame rate
        clock.tick(SNAKE_SPEED)

    # Game over message
    font = pygame.font.SysFont(None, 100)
    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()
    time.sleep(2)

    pygame.quit()

if __name__ == "__main__":
    main()