import pygame
import random
import sys

# Initialize pygame
pygame.init()

# Set up display
WIDTH, HEIGHT = 500, 600
win = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Dodge the Blocks")

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
BLOCK_COLOR = (255, 0, 0)
PLAYER_COLORS = [(0, 128, 255), (255, 100, 100), (100, 255, 100), (255, 255, 0), (255, 0, 255)]

# Fonts
font = pygame.font.SysFont("comicsansms", 28)
big_font = pygame.font.SysFont("comicsansms", 50)

# Clock
clock = pygame.time.Clock()


def detect_collision(player_pos, block_pos, player_size, block_size):
    px, py = player_pos
    bx, by = block_pos
    return (
        bx < px + player_size and px < bx + block_size and
        by < py + player_size and py < by + block_size
    )

def draw_text_centered(text, font, color, y_offset=0):
    text_surface = font.render(text, True, color)
    rect = text_surface.get_rect(center=(WIDTH // 2, HEIGHT // 2 + y_offset))
    win.blit(text_surface, rect)


def customization_screen():
    color_index = 0
    selecting = True

    while selecting:
        win.fill(WHITE)
        draw_text_centered("Choose Your Color", big_font, BLACK, -100)

        # Draw color preview
        pygame.draw.rect(win, PLAYER_COLORS[color_index],
                         (WIDTH // 2 - 25, HEIGHT // 2 - 25, 50, 50))

        draw_text_centered("← / → to change color", font, BLACK, 80)
        draw_text_centered("Press ENTER to start", font, BLACK, 120)
        pygame.display.update()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_LEFT:
                    color_index = (color_index - 1) % len(PLAYER_COLORS)
                elif event.key == pygame.K_RIGHT:
                    color_index = (color_index + 1) % len(PLAYER_COLORS)
                elif event.key == pygame.K_RETURN:
                    selecting = False

    return PLAYER_COLORS[color_index]


def game_over_screen(score):
    win.fill(WHITE)
    draw_text_centered("Game Over", big_font, BLOCK_COLOR, -50)
    draw_text_centered(f"Score: {int(score)}", font, BLACK)
    draw_text_centered("Press R to Restart or Q to Quit", font, BLACK, 50)
    pygame.display.update()

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_r:
                    return True
                if event.key == pygame.K_q:
                    pygame.quit()
                    sys.exit()


def run_game(player_color):
    player_size = 50
    player_pos = [WIDTH // 2, HEIGHT - player_size]
    player_speed = 7

    block_size = 50
    block_pos = [random.randint(0, WIDTH - block_size), 0]
    block_speed = 7

    gravity = 1
    jump_power = -15
    velocity_y = 0
    on_ground = True
    score = 0
    game_over = False

    while not game_over:
        clock.tick(60)

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

        keys = pygame.key.get_pressed()
        if keys[pygame.K_LEFT] and player_pos[0] > 0:
            player_pos[0] -= player_speed
        if keys[pygame.K_RIGHT] and player_pos[0] < WIDTH - player_size:
            player_pos[0] += player_speed
        if keys[pygame.K_SPACE] and on_ground:
            velocity_y = jump_power
            on_ground = False

        # Apply gravity
        velocity_y += gravity
        player_pos[1] += velocity_y
        if player_pos[1] >= HEIGHT - player_size:
            player_pos[1] = HEIGHT - player_size
            velocity_y = 0
            on_ground = True

        # Move block
        block_pos[1] += block_speed
        if block_pos[1] > HEIGHT:
            block_pos[1] = 0
            block_pos[0] = random.randint(0, WIDTH - block_size)
            score += 1
            block_speed += 0.3

        if detect_collision(player_pos, block_pos, player_size, block_size):
            game_over = True

        win.fill(WHITE)
        pygame.draw.rect(win, player_color, (*player_pos, player_size, player_size))
        pygame.draw.rect(win, BLOCK_COLOR, (*block_pos, block_size, block_size))
        win.blit(font.render(f"Score: {int(score)}", True, BLACK), (10, 10))
        pygame.display.update()

    if game_over_screen(score):
        main()


def main():
    player_color = customization_screen()
    run_game(player_color)


# --- Start Game ---
main()