import pygame, random, sys

# === GLOBAL SETUP ===
pygame.init()
w, h = 600, 400           # Window size
block = 20                # Snake block size
food_size = block + 6     # Bigger food size
win = pygame.display.set_mode((w, h))
pygame.display.set_caption("Snake Game")
clock = pygame.time.Clock()
font = pygame.font.SysFont(None, 35)

# === DRAW SNAKE ===
def draw_snake(snake):
    for x, y in snake:
        pygame.draw.rect(win, (0, 255, 0), (x, y, block, block))

# === SHOW SCORE ===
def show_score(score):
    text = font.render(f"Score: {score}", True, (255, 255, 255))
    win.blit(text, (10, 10))

# === MAIN GAME LOOP ===
def game_loop():
    x, y = w // 2, h // 2
    dx, dy = block, 0
    next_dx, next_dy = dx, dy
    snake = [(x, y)]
    length = 1
    food = get_new_food(snake)
    growth = 0  # This tracks how much to grow the snake

    while True:
        win.fill((0, 0, 0))
        offset = (food_size - block) // 2
        pygame.draw.rect(win, (255, 0, 0), (food[0] - offset, food[1] - offset, food_size, food_size))

        direction_changed = False
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit(); sys.exit()
            elif event.type == pygame.KEYDOWN and not direction_changed:
                if event.key == pygame.K_LEFT and dx != block:
                    next_dx, next_dy = -block, 0
                    direction_changed = True
                elif event.key == pygame.K_RIGHT and dx != -block:
                    next_dx, next_dy = block, 0
                    direction_changed = True
                elif event.key == pygame.K_UP and dy != block:
                    next_dx, next_dy = 0, -block
                    direction_changed = True
                elif event.key == pygame.K_DOWN and dy != -block:
                    next_dx, next_dy = 0, block
                    direction_changed = True

        dx, dy = next_dx, next_dy
        x += dx
        y += dy
        head = (x, y)

        if x < 0 or x >= w or y < 0 or y >= h or head in snake:
            break

        snake.append(head)

        # If snake eats the food, gradually grow it
        if head == food:
            growth += 1
            food = get_new_food(snake)

        # If there's growth left, add the new segment to the snake
        if growth > 0:
            snake.insert(0, snake[0])  # Add new segment at the front of the snake
            growth -= 1

        if len(snake) > length:
            snake.pop(0)

        draw_snake(snake)
        show_score(length - 1)
        pygame.display.update()
        clock.tick(10)

    game_over_screen()

# === GET NEW FOOD ===
def get_new_food(snake):
    while True:
        food = (random.randrange(0, w, block), random.randrange(0, h, block))
        if food not in snake:
            return food

# === GAME OVER SCREEN WITH CLICKABLE BUTTON ===
def game_over_screen():
    win.fill((0, 0, 0))
    msg = font.render("Game Over!", True, (255, 255, 255))
    win.blit(msg, (w // 2 - 80, h // 3))

    restart_btn = pygame.Rect(w // 2 - 80, h // 2, 160, 40)
    quit_btn = pygame.Rect(w // 2 - 80, h // 2 + 60, 160, 40)

    pygame.draw.rect(win, (0, 200, 0), restart_btn)
    pygame.draw.rect(win, (200, 0, 0), quit_btn)

    restart_text = font.render("Restart", True, (255, 255, 255))
    quit_text = font.render("Quit", True, (255, 255, 255))

    win.blit(restart_text, (restart_btn.x + 30, restart_btn.y + 5))
    win.blit(quit_text, (quit_btn.x + 50, quit_btn.y + 5))

    pygame.display.update()

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit(); sys.exit()
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_r:
                    game_loop(); return
                elif event.key == pygame.K_q:
                    pygame.quit(); sys.exit()
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if restart_btn.collidepoint(event.pos):
                    game_loop(); return
                elif quit_btn.collidepoint(event.pos):
                    pygame.quit(); sys.exit()

# === START GAME ===
game_loop()