main.py

import random
import pygame


pygame.init()
pygame.mixer.init()

sound_eat = pygame.mixer.Sound("eat.wav")
sound_move = pygame.mixer.Sound("move.wav")
sound_game_over = pygame.mixer.Sound("game_over.wav")
sound_eat.set_volume(0.5)
sound_move.set_volume(0.1)
sound_game_over.set_volume(0.5)

window = pygame.display.set_mode((1280, 720))
pygame.display.set_caption("Snake Game")
screen_width = window.get_width()
screen_height = window.get_height()
clock = pygame.time.Clock()

BLOCK_SIZE = 40
BUTTON_WIDTH = 320
BUTTON_HEIGHT = 60
BUTTON_SPACING = 20
columns = screen_width // BLOCK_SIZE
rows = screen_height // BLOCK_SIZE
FPS = 10

# keeping the directions as tuples makes the math manageable for me.
DIR_RIGHT = (1, 0)
DIR_LEFT = (-1, 0)
DIR_UP = (0, -1)
DIR_DOWN = (0, 1)

pygame.font.init()
font = pygame.font.SysFont(None, 36)
title_font = pygame.font.SysFont(None, 72)

# sorry future me, I keep almost everything global until I get brave enough for classes
running = True
snake = []
direction = DIR_RIGHT
next_direction = DIR_RIGHT
points = 0
pause = False
game_state = "menu"
game_over_flag = False
menu_selection_index = 0
game_over_selection_index = 0

menu_buttons = [
    {"text": "Play", "action": "play"},
    {"text": "Quit", "action": "quit"},
]

game_over_buttons = [
    {"text": "Play Again", "action": "play"},
    {"text": "Main Menu", "action": "menu"},
]


def setup_snake():
    global snake, direction, next_direction, points, pause, game_over_flag
    head_x = columns // 2
    head_y = rows // 2
    snake = [
        (head_x, head_y),
        (head_x - 1, head_y),
        (head_x - 2, head_y),
    ]
    direction = DIR_RIGHT
    next_direction = DIR_RIGHT
    points = 0
    pause = False
    game_over_flag = False


setup_snake()


def spawn_food():
    while True:
        random_pos = (
            random.randint(0, columns - 1),
            random.randint(0, rows - 1),
        )
        if random_pos not in snake:
            return random_pos


food_pos = spawn_food()


def toggle_pause():
    global pause
    if pause is True:
        pause = False
    else:
        pause = True


def handle_events():
    global running
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
            continue
        if game_state == "menu":
            handle_menu(event)
        elif game_state == "playing":
            handle_playing(event)
        elif game_state == "game_over":
            handle_game_over(event)


def handle_menu(event):
    global menu_selection_index
    menu_rects = build_button_rects(menu_buttons, screen_height // 2 + 80)
    if event.type == pygame.KEYDOWN:
        if (
            event.key == pygame.K_UP
            or event.key == pygame.K_w
            or event.key == pygame.K_LEFT
        ):
            menu_selection_index -= 1
            if menu_selection_index < 0:
                menu_selection_index = len(menu_buttons) - 1
        elif (
            event.key == pygame.K_DOWN
            or event.key == pygame.K_s
            or event.key == pygame.K_RIGHT
        ):
            menu_selection_index += 1
            if menu_selection_index >= len(menu_buttons):
                menu_selection_index = 0
        elif event.key == pygame.K_RETURN or event.key == pygame.K_KP_ENTER:
            press_button(menu_buttons[menu_selection_index]["action"])
    elif event.type == pygame.MOUSEMOTION:
        idx = button_index_from_pos(menu_rects, event.pos)
        if idx is not None:
            menu_selection_index = idx
    elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
        idx = button_index_from_pos(menu_rects, event.pos)
        if idx is not None:
            menu_selection_index = idx
            press_button(menu_buttons[idx]["action"])


def handle_playing(event):
    global next_direction
    if event.type != pygame.KEYDOWN:
        return
    # I like to keep each direction check separate so I can see them at a glance.
    if event.key == pygame.K_w or event.key == pygame.K_UP:
        if direction != DIR_DOWN:
            next_direction = DIR_UP
            sound_move.play()
    if event.key == pygame.K_s or event.key == pygame.K_DOWN:
        if direction != DIR_UP:
            next_direction = DIR_DOWN
            sound_move.play()
    if event.key == pygame.K_a or event.key == pygame.K_LEFT:
        if direction != DIR_RIGHT:
            next_direction = DIR_LEFT
            sound_move.play()
    if event.key == pygame.K_d or event.key == pygame.K_RIGHT:
        if direction != DIR_LEFT:
            next_direction = DIR_RIGHT
            sound_move.play()
    if event.key == pygame.K_p:
        toggle_pause()


def handle_game_over(event):
    global game_over_selection_index
    go_rects = build_button_rects(game_over_buttons, screen_height // 2 + 120)
    if event.type == pygame.KEYDOWN:
        if (
            event.key == pygame.K_UP
            or event.key == pygame.K_w
            or event.key == pygame.K_LEFT
        ):
            game_over_selection_index -= 1
            if game_over_selection_index < 0:
                game_over_selection_index = len(game_over_buttons) - 1
        elif (
            event.key == pygame.K_DOWN
            or event.key == pygame.K_s
            or event.key == pygame.K_RIGHT
        ):
            game_over_selection_index += 1
            if game_over_selection_index >= len(game_over_buttons):
                game_over_selection_index = 0
        elif event.key == pygame.K_RETURN or event.key == pygame.K_KP_ENTER:
            press_button(game_over_buttons[game_over_selection_index]["action"])
    elif event.type == pygame.MOUSEMOTION:
        idx = button_index_from_pos(go_rects, event.pos)
        if idx is not None:
            game_over_selection_index = idx
    elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
        idx = button_index_from_pos(go_rects, event.pos)
        if idx is not None:
            game_over_selection_index = idx
            press_button(game_over_buttons[idx]["action"])


def draw_game_screen():
    window.fill("black")
    draw_snake()
    draw_food()
    draw_hud()
    if pause is True and game_over_flag is False:
        pause_text = font.render("PAUSE", True, "white")
        window.blit(
            pause_text,
            (screen_width // 2 - pause_text.get_width() // 2, screen_height // 2),
        )
    pygame.display.flip()


def draw_snake():
    for block in snake:
        pygame.draw.rect(
            window,
            "red",
            (
                block[0] * BLOCK_SIZE,
                block[1] * BLOCK_SIZE,
                BLOCK_SIZE,
                BLOCK_SIZE,
            ),
        )


def draw_food():
    pygame.draw.rect(
        window,
        "blue",
        (food_pos[0] * BLOCK_SIZE, food_pos[1] * BLOCK_SIZE, BLOCK_SIZE, BLOCK_SIZE),
    )


def draw_hud():
    # I keep the HUD super simple for now because I just need the score.
    text_surface = font.render("Score: " + str(points), True, "white")
    window.blit(text_surface, (10, 10))


def draw_menu_screen():
    window.fill("black")
    title_surface = title_font.render("Snake Game", True, "white")
    window.blit(
        title_surface,
        (screen_width // 2 - title_surface.get_width() // 2, screen_height // 2 - 80),
    )
    menu_rects = build_button_rects(menu_buttons, screen_height // 2 + 80)
    draw_button_group(menu_buttons, menu_rects, menu_selection_index)
    pygame.display.flip()


def draw_game_over_screen():
    window.fill("black")
    text = font.render("GAME OVER", True, "red")
    score_text = font.render("Score: " + str(points), True, "red")
    window.blit(
        text, (screen_width // 2 - text.get_width() // 2, screen_height // 2 - 60)
    )
    window.blit(
        score_text,
        (screen_width // 2 - score_text.get_width() // 2, screen_height // 2 - 20),
    )
    go_rects = build_button_rects(game_over_buttons, screen_height // 2 + 120)
    draw_button_group(game_over_buttons, go_rects, game_over_selection_index)


def draw_button_group(buttons, rects, selected_idx):
    # this probably could be smaller, but I understand it this way.
    for idx in range(len(buttons)):
        rect = rects[idx]
        is_selected = idx == selected_idx
        button_color = "#ffb347" if is_selected else "#000000"
        text_color = "black" if is_selected else "white"
        pygame.draw.rect(window, button_color, rect, border_radius=10)
        pygame.draw.rect(window, "white", rect, width=2, border_radius=10)
        label = font.render(buttons[idx]["text"], True, text_color)
        label_rect = label.get_rect(center=rect.center)
        window.blit(label, label_rect)


def build_button_rects(buttons, base_y):
    total_height = len(buttons) * BUTTON_HEIGHT + (len(buttons) - 1) * BUTTON_SPACING
    start_y = base_y - total_height // 2 + BUTTON_HEIGHT // 2
    rects = []
    for _ in buttons:
        rect = pygame.Rect(0, 0, BUTTON_WIDTH, BUTTON_HEIGHT)
        rect.center = (screen_width // 2, start_y)
        rects.append(rect)
        start_y += BUTTON_HEIGHT + BUTTON_SPACING
    return rects


def button_index_from_pos(rects, pos):
    for idx, rect in enumerate(rects):
        if rect.collidepoint(pos):
            return idx
    return None


def update_game_logic():
    global direction, snake, points, food_pos
    if pause or game_state != "playing" or game_over_flag:
        return
    direction = next_direction
    new_head_x = snake[0][0] + direction[0]
    new_head_y = snake[0][1] + direction[1]
    if new_head_x < 0 or new_head_x >= columns:
        game_over()
        return
    if new_head_y < 0 or new_head_y >= rows:
        game_over()
        return
    if (new_head_x, new_head_y) in snake:
        game_over()
        return
    snake.insert(0, (new_head_x, new_head_y))
    if (new_head_x, new_head_y) == food_pos:
        points += 1
        #  print("Score is now", points)
        food_pos = spawn_food()
        sound_eat.play()
    else:
        snake.pop(len(snake) - 1)


def press_button(action):
    global running
    if action == "play":
        restart_game()
    elif action == "quit":
        running = False
    elif action == "menu":
        go_to_menu()


def restart_game():

    global food_pos, game_state, game_over_selection_index
    # print("Restarting the whole game state now.")
    setup_snake()
    food_pos = spawn_food()
    game_state = "playing"
    game_over_selection_index = 0


def go_to_menu():
    global game_state, menu_selection_index, pause, game_over_flag
    #  print("Heading back to the menu.")
    game_state = "menu"
    menu_selection_index = 0
    pause = False
    game_over_flag = False


def game_over():
    global pause, game_over_flag, game_state, game_over_selection_index
    if game_state == "game_over":
        return
    #   print("the snake crashed.")
    pause = True
    game_over_flag = True
    game_state = "game_over"
    game_over_selection_index = 0
    sound_game_over.play()


# player_pos = pygame.Vector2(screen_width / 2, screen_height / 2)

while running:
    handle_events()
    if game_state == "menu":
        draw_menu_screen()
    elif game_state == "playing":
        update_game_logic()
        draw_game_screen()
    elif game_state == "game_over":
        draw_game_over_screen()
        pygame.display.flip()
    clock.tick(FPS)

pygame.quit()