פונקציה אחת לכל מצב

import random
import pygame

pygame.init()
screen = pygame.display.set_mode((480, 360))
pygame.display.set_caption("Meteor Dodge")
clock = pygame.time.Clock()
font = pygame.font.Font(None, 32)
WHITE = (240, 240, 240)


def reset():
    global player, rock, score
    player = pygame.Rect(220, 300, 40, 40)
    rock = pygame.Rect(random.randint(0, 456), -24, 24, 24)
    score = 0


def draw_menu():
    screen.fill((30, 30, 46))
    title = font.render("Meteor Dodge", True, WHITE)
    screen.blit(title, title.get_rect(center=(240, 140)))
    hint = font.render("Press Enter to play", True, WHITE)
    screen.blit(hint, hint.get_rect(center=(240, 200)))


def draw_game():
    screen.fill((30, 30, 46))
    pygame.draw.rect(screen, (255, 196, 0), player)
    pygame.draw.rect(screen, (235, 80, 80), rock)
    score_text = font.render(f"Score: {score}", True, WHITE)
    screen.blit(score_text, (10, 10))


def play_frame():
    global score, state
    keys = pygame.key.get_pressed()
    if keys[pygame.K_LEFT]:
        player.x -= 6
    if keys[pygame.K_RIGHT]:
        player.x += 6
    player.clamp_ip(screen.get_rect())

    rock.y += 4
    if rock.top > 360:
        rock.y = -24
        rock.x = random.randint(0, 456)
        score += 1
    if rock.colliderect(player):
        state = "over"
    draw_game()


def draw_over():
    draw_game()
    message = font.render("Game over", True, WHITE)
    screen.blit(message, message.get_rect(center=(240, 160)))
    hint = font.render("Press R to play again", True, WHITE)
    screen.blit(hint, hint.get_rect(center=(240, 200)))


reset()
state = "menu"

running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_RETURN and state == "menu":
                state = "play"
            elif event.key == pygame.K_r and state == "over":
                reset()
                state = "play"

    if state == "menu":
        draw_menu()
    elif state == "play":
        play_frame()
    else:
        draw_over()
    pygame.display.flip()
    clock.tick(60)

pygame.quit()