import pygame
import sys
import math
import random

# ============================================================
# FROGGY DASH
# Geometry-Dash-inspired platform game
# Everything is contained in this single file.
# ============================================================

pygame.init()

# -------------------- WINDOW --------------------

WIDTH = 1100
HEIGHT = 650
FPS = 60

screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Froggy Dash")
clock = pygame.time.Clock()

# -------------------- COLORS --------------------

WHITE = (245, 248, 255)
BLACK = (15, 18, 30)
DARK = (25, 29, 48)
DARKER = (14, 17, 29)

GREEN = (80, 230, 120)
LIGHT_GREEN = (135, 255, 155)
LIME = (190, 255, 70)

BLUE = (75, 180, 255)
CYAN = (60, 245, 230)
PURPLE = (180, 95, 255)
PINK = (255, 90, 170)
RED = (255, 75, 90)
ORANGE = (255, 165, 60)
YELLOW = (255, 225, 70)

GRAY = (100, 110, 135)
LIGHT_GRAY = (170, 180, 200)

# -------------------- FONTS --------------------

FONT_BIG = pygame.font.SysFont("arial", 64, bold=True)
FONT_TITLE = pygame.font.SysFont("arial", 82, bold=True)
FONT_MED = pygame.font.SysFont("arial", 32, bold=True)
FONT_SMALL = pygame.font.SysFont("arial", 22)
FONT_TINY = pygame.font.SysFont("arial", 17)

# -------------------- GAME STATES --------------------

MENU = "menu"
FROG_SELECT = "frog_select"
LEVEL_SELECT = "level_select"
PLAYING = "playing"
PAUSED = "paused"
GAME_OVER = "game_over"
LEVEL_COMPLETE = "level_complete"

game_state = MENU

# -------------------- FROGS --------------------

frogs = [
    {
        "name": "Greeny",
        "color": (70, 220, 100),
        "dark": (35, 130, 60),
        "eye": WHITE,
        "accent": LIME
    },
    {
        "name": "Frogo",
        "color": (70, 180, 255),
        "dark": (30, 90, 160),
        "eye": WHITE,
        "accent": CYAN
    },
    {
        "name": "Berry",
        "color": (220, 80, 180),
        "dark": (130, 35, 110),
        "eye": WHITE,
        "accent": PINK
    },
    {
        "name": "Lemon",
        "color": (245, 220, 60),
        "dark": (170, 130, 20),
        "eye": BLACK,
        "accent": ORANGE
    },
    {
        "name": "Purple",
        "color": (165, 90, 255),
        "dark": (90, 40, 160),
        "eye": WHITE,
        "accent": CYAN
    },
    {
        "name": "Fire Frog",
        "color": (255, 90, 55),
        "dark": (155, 35, 30),
        "eye": YELLOW,
        "accent": ORANGE
    },
]

selected_frog = 0

# -------------------- LEVELS --------------------

LEVELS = [
    {
        "name": "Froggy Fields",
        "difficulty": "EASY",
        "speed": 6.0,
        "color": (50, 190, 100),
        "bg": (20, 65, 48),
        "length": 6500,
    },
    {
        "name": "Frog Factory",
        "difficulty": "NORMAL",
        "speed": 6.7,
        "color": (60, 160, 220),
        "bg": (20, 45, 75),
        "length": 7500,
    },
    {
        "name": "Neon Froggy",
        "difficulty": "HARD",
        "speed": 7.4,
        "color": (180, 70, 240),
        "bg": (45, 20, 70),
        "length": 8500,
    },
    {
        "name": "Frog Galaxy",
        "difficulty": "INSANE",
        "speed": 8.2,
        "color": (230, 75, 130),
        "bg": (20, 20, 55),
        "length": 9500,
    },
]

selected_level = 0

# -------------------- PARTICLES --------------------

particles = []


class Particle:
    def __init__(self, x, y, color, vx=None, vy=None, life=30, size=5):
        self.x = x
        self.y = y
        self.color = color
        self.vx = random.uniform(-2.5, 2.5) if vx is None else vx
        self.vy = random.uniform(-3.5, 1) if vy is None else vy
        self.life = life
        self.max_life = life
        self.size = size

    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.vy += 0.12
        self.life -= 1

    def draw(self):
        if self.life <= 0:
            return

        alpha = max(0, min(255, int(255 * self.life / self.max_life)))

        surf = pygame.Surface(
            (self.size * 2 + 4, self.size * 2 + 4),
            pygame.SRCALPHA
        )

        pygame.draw.circle(
            surf,
            (*self.color, alpha),
            (self.size + 2, self.size + 2),
            self.size
        )

        screen.blit(
            surf,
            (self.x - self.size - 2, self.y - self.size - 2)
        )


# -------------------- LEVEL OBJECTS --------------------

class Spike:
    def __init__(self, x, ground_y, size=35):
        self.x = x
        self.y = ground_y
        self.size = size

    def rect(self):
        return pygame.Rect(
            self.x + 7,
            self.y - self.size + 10,
            self.size - 14,
            self.size - 10
        )

    def draw(self, camera_x, color):
        x = self.x - camera_x
        y = self.y

        points = [
            (x, y),
            (x + self.size / 2, y - self.size),
            (x + self.size, y)
        ]

        pygame.draw.polygon(screen, color, points)

        pygame.draw.line(
            screen,
            WHITE,
            points[0],
            points[1],
            2
        )

        pygame.draw.line(
            screen,
            WHITE,
            points[1],
            points[2],
            2
        )


class Block:
    def __init__(self, x, y, width, height):
        self.x = x
        self.y = y
        self.width = width
        self.height = height

    def rect(self):
        return pygame.Rect(
            self.x,
            self.y,
            self.width,
            self.height
        )

    def draw(self, camera_x, color):
        rect = pygame.Rect(
            self.x - camera_x,
            self.y,
            self.width,
            self.height
        )

        pygame.draw.rect(
            screen,
            color,
            rect,
            border_radius=5
        )

        pygame.draw.rect(
            screen,
            WHITE,
            rect,
            2,
            border_radius=5
        )

        # Decorative grid
        for xx in range(rect.x + 15, rect.right, 25):
            pygame.draw.line(
                screen,
                (*WHITE[:2], 90),
                (xx, rect.y),
                (xx, rect.bottom),
                1
            )


class Coin:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.collected = False
        self.rotation = random.random() * math.pi

    def update(self):
        self.rotation += 0.08

    def draw(self, camera_x):
        if self.collected:
            return

        x = self.x - camera_x
        scale = abs(math.cos(self.rotation))
        width = max(4, int(15 * scale))

        pygame.draw.ellipse(
            screen,
            YELLOW,
            pygame.Rect(
                x - width,
                self.y - 15,
                width * 2,
                30
            )
        )

        pygame.draw.circle(
            screen,
            ORANGE,
            (int(x), int(self.y)),
            5
        )


# -------------------- FROG PLAYER --------------------

class Frog:
    def __init__(self):
        self.x = 150
        self.y = 450
        self.size = 42

        self.vel_y = 0
        self.gravity = 0.72
        self.jump_power = -13

        self.on_ground = False
        self.rotation = 0

        self.alive = True

    def reset(self):
        self.x = 150
        self.y = 450
        self.vel_y = 0
        self.rotation = 0
        self.on_ground = False
        self.alive = True

    def rect(self):
        return pygame.Rect(
            int(self.x + 4),
            int(self.y + 5),
            self.size - 8,
            self.size - 7
        )

    def jump(self):
        if self.on_ground:
            self.vel_y = self.jump_power
            self.on_ground = False

            for _ in range(8):
                particles.append(
                    Particle(
                        self.x + self.size / 2,
                        self.y + self.size,
                        frogs[selected_frog]["accent"],
                        vx=random.uniform(-2, 2),
                        vy=random.uniform(-2, 0),
                        life=25,
                        size=random.randint(2, 4)
                    )
                )

    def update(self, blocks):
        self.vel_y += self.gravity
        self.y += self.vel_y

        self.on_ground = False

        # Ground
        ground_y = 540

        if self.y + self.size >= ground_y:
            self.y = ground_y - self.size
            self.vel_y = 0
            self.on_ground = True

        # Blocks
        for block in blocks:
            r = self.rect()
            br = block.rect()

            if r.colliderect(br):
                previous_bottom = r.bottom - self.vel_y

                if (
                    self.vel_y >= 0
                    and previous_bottom <= br.top + 10
                ):
                    self.y = br.top - self.size
                    self.vel_y = 0
                    self.on_ground = True

        if not self.on_ground:
            self.rotation += 7
        else:
            self.rotation = round(self.rotation / 90) * 90

    def draw(self):
        frog = frogs[selected_frog]

        x = int(self.x + self.size / 2)
        y = int(self.y + self.size / 2)

        # Create frog surface
        surf = pygame.Surface(
            (70, 70),
            pygame.SRCALPHA
        )

        cx = 35
        cy = 35

        # Body
        pygame.draw.circle(
            surf,
            frog["color"],
            (cx, cy + 5),
            21
        )

        # Feet
        pygame.draw.ellipse(
            surf,
            frog["dark"],
            (7, 40, 25, 14)
        )

        pygame.draw.ellipse(
            surf,
            frog["dark"],
            (38, 40, 25, 14)
        )

        # Eyes
        pygame.draw.circle(
            surf,
            frog["color"],
            (cx - 13, cy - 14),
            12
        )

        pygame.draw.circle(
            surf,
            frog["color"],
            (cx + 13, cy - 14),
            12
        )

        pygame.draw.circle(
            surf,
            frog["eye"],
            (cx - 13, cy - 14),
            7
        )

        pygame.draw.circle(
            surf,
            frog["eye"],
            (cx + 13, cy - 14),
            7
        )

        # Pupils
        pupil = BLACK if frog["eye"] == WHITE else DARK

        pygame.draw.circle(
            surf,
            pupil,
            (cx - 13, cy - 14),
            3
        )

        pygame.draw.circle(
            surf,
            pupil,
            (cx + 13, cy - 14),
            3
        )

        # Mouth
        pygame.draw.arc(
            surf,
            frog["dark"],
            (18, 27, 34, 18),
            0,
            math.pi,
            3
        )

        # Rotate
        rotated = pygame.transform.rotate(
            surf,
            -self.rotation
        )

        rect = rotated.get_rect(
            center=(x, y)
        )

        screen.blit(rotated, rect)


# -------------------- LEVEL GENERATION --------------------

current_level = None
frog = Frog()

blocks = []
spikes = []
coins = []

camera_x = 0
level_distance = 0
level_coins = 0
level_start_time = 0


def generate_level(index):
    global blocks, spikes, coins

    blocks = []
    spikes = []
    coins = []

    level = LEVELS[index]

    random.seed(1000 + index)

    # Start area
    x = 600

    # Different patterns
    while x < level["length"]:
        pattern = random.randint(0, 7)

        if pattern == 0:
            # Single spike
            spikes.append(
                Spike(x, 540, 38)
            )
            x += random.randint(240, 320)

        elif pattern == 1:
            # Double spike
            spikes.append(
                Spike(x, 540, 35)
            )
            spikes.append(
                Spike(x + 38, 540, 35)
            )
            x += 300

        elif pattern == 2:
            # Platform
            h = random.choice([80, 110, 140])

            blocks.append(
                Block(
                    x,
                    540 - h,
                    130,
                    h
                )
            )

            coins.append(
                Coin(
                    x + 65,
                    540 - h - 35
                )
            )

            x += 320

        elif pattern == 3:
            # Jump over spike + coin
            spikes.append(
                Spike(x, 540, 38)
            )

            coins.append(
                Coin(
                    x + 80,
                    445
                )
            )

            x += 260

        elif pattern == 4:
            # Staircase
            blocks.append(
                Block(
                    x,
                    470,
                    100,
                    70
                )
            )

            blocks.append(
                Block(
                    x + 120,
                    410,
                    100,
                    130
                )
            )

            blocks.append(
                Block(
                    x + 240,
                    350,
                    100,
                    190
                )
            )

            coins.append(
                Coin(
                    x + 170,
                    350
                )
            )

            x += 440

        elif pattern == 5:
            # Spike sequence
            for i in range(3):
                spikes.append(
                    Spike(
                        x + i * 42,
                        540,
                        34
                    )
                )

            x += 360

        elif pattern == 6:
            # Platform with spikes
            blocks.append(
                Block(
                    x,
                    455,
                    170,
                    85
                )
            )

            spikes.append(
                Spike(
                    x + 210,
                    540,
                    40
                )
            )

            coins.append(
                Coin(
                    x + 85,
                    415
                )
            )

            x += 380

        else:
            # Coin trail
            for i in range(4):
                coins.append(
                    Coin(
                        x + i * 45,
                        430 - abs(i - 1.5) * 25
                    )
                )

            x += 330

    # Final area
    spikes.append(
        Spike(
            level["length"] - 420,
            540,
            40
        )
    )

    spikes.append(
        Spike(
            level["length"] - 360,
            540,
            40
        )
    )

    random.seed()


# -------------------- DRAWING --------------------

def draw_background():
    level = LEVELS[selected_level]

    # Gradient-like background
    screen.fill(level["bg"])

    # Large decorative circles
    for i in range(8):
        px = (
            i * 180
            - int(camera_x * 0.15)
        ) % (WIDTH + 300) - 100

        py = 80 + (i % 3) * 130

        pygame.draw.circle(
            screen,
            tuple(
                min(255, c + 15)
                for c in level["bg"]
            ),
            (px, py),
            50 + (i % 3) * 20,
            2
        )

    # Grid
    grid_offset = int(camera_x * 0.3) % 50

    for x in range(-50, WIDTH + 50, 50):
        pygame.draw.line(
            screen,
            (255, 255, 255, 20),
            (x - grid_offset, 0),
            (x - grid_offset, HEIGHT),
            1
        )

    for y in range(0, HEIGHT, 50):
        pygame.draw.line(
            screen,
            (255, 255, 255, 20),
            (0, y),
            (WIDTH, y),
            1
        )

    # Ground
    pygame.draw.rect(
        screen,
        DARKER,
        (0, 540, WIDTH, 110)
    )

    pygame.draw.line(
        screen,
        level["color"],
        (0, 540),
        (WIDTH, 540),
        5
    )

    # Ground decorations
    for x in range(
        -50,
        WIDTH + 100,
        70
    ):
        xx = (
            x
            - int(camera_x * 0.8)
        ) % (WIDTH + 70)

        pygame.draw.rect(
            screen,
            level["color"],
            (xx, 575, 35, 5),
            border_radius=2
        )


def draw_text_center(text, font, color, y):
    surf = font.render(
        text,
        True,
        color
    )

    screen.blit(
        surf,
        (
            WIDTH // 2 - surf.get_width() // 2,
            y
        )
    )


def draw_button(rect, text, color, mouse_pos):
    hover = rect.collidepoint(mouse_pos)

    button_color = color

    if hover:
        button_color = tuple(
            min(255, c + 30)
            for c in color
        )

    pygame.draw.rect(
        screen,
        button_color,
        rect,
        border_radius=12
    )

    pygame.draw.rect(
        screen,
        WHITE,
        rect,
        2,
        border_radius=12
    )

    surf = FONT_MED.render(
        text,
        True,
        WHITE
    )

    screen.blit(
        surf,
        (
            rect.centerx - surf.get_width() // 2,
            rect.centery - surf.get_height() // 2
        )
    )


def draw_menu():
    screen.fill(DARKER)

    # Animated circles
    t = pygame.time.get_ticks() / 1000

    for i in range(12):
        x = (
            i * 110
            + math.sin(t + i) * 30
        ) % WIDTH

        y = (
            100
            + i * 37
            + math.cos(t * 0.7 + i) * 30
        ) % 520

        pygame.draw.circle(
            screen,
            (
                30 + i * 8,
                80 + i * 5,
                100 + i * 10
            ),
            (int(x), int(y)),
            4
        )

    draw_text_center(
        "FROGGY",
        FONT_TITLE,
        GREEN,
        70
    )

    draw_text_center(
        "DASH",
        FONT_TITLE,
        CYAN,
        145
    )

    draw_text_center(
        "A HOPPING RHYTHM ADVENTURE",
        FONT_SMALL,
        LIGHT_GRAY,
        245
    )

    mouse = pygame.mouse.get_pos()

    play_rect = pygame.Rect(
        WIDTH // 2 - 160,
        310,
        320,
        65
    )

    frog_rect = pygame.Rect(
        WIDTH // 2 - 160,
        395,
        320,
        65
    )

    quit_rect = pygame.Rect(
        WIDTH // 2 - 160,
        480,
        320,
        65
    )

    draw_button(
        play_rect,
        "PLAY",
        GREEN,
        mouse
    )

    draw_button(
        frog_rect,
        "CHOOSE FROG",
        PURPLE,
        mouse
    )

    draw_button(
        quit_rect,
        "QUIT",
        RED,
        mouse
    )

    # Mini frog
    draw_frog_preview(
        WIDTH - 150,
        90,
        selected_frog,
        0.9
    )

    draw_text_center(
        "SPACE / W / UP = JUMP",
        FONT_TINY,
        GRAY,
        595
    )


def draw_frog_preview(x, y, index, scale=1):
    frog_data = frogs[index]

    size = int(70 * scale)

    # Body
    pygame.draw.circle(
        screen,
        frog_data["color"],
        (x, y + 8),
        int(25 * scale)
    )

    # Eyes
    pygame.draw.circle(
        screen,
        frog_data["color"],
        (
            int(x - 16 * scale),
            int(y - 15 * scale)
        ),
        int(14 * scale)
    )

    pygame.draw.circle(
        screen,
        frog_data["color"],
        (
            int(x + 16 * scale),
            int(y - 15 * scale)
        ),
        int(14 * scale)
    )

    eye_color = frog_data["eye"]

    pygame.draw.circle(
        screen,
        eye_color,
        (
            int(x - 16 * scale),
            int(y - 15 * scale)
        ),
        int(8 * scale)
    )

    pygame.draw.circle(
        screen,
        eye_color,
        (
            int(x + 16 * scale),
            int(y - 15 * scale)
        ),
        int(8 * scale)
    )

    pupil = BLACK if eye_color != BLACK else WHITE

    pygame.draw.circle(
        screen,
        pupil,
        (
            int(x - 16 * scale),
            int(y - 15 * scale)
        ),
        int(3 * scale)
    )

    pygame.draw.circle(
        screen,
        pupil,
        (
            int(x + 16 * scale),
            int(y - 15 * scale)
        ),
        int(3 * scale)
    )


def draw_frog_select():
    screen.fill(DARKER)

    draw_text_center(
        "CHOOSE YOUR FROG",
        FONT_BIG,
        GREEN,
        35
    )

    draw_text_center(
        "Each frog has a different color!",
        FONT_SMALL,
        LIGHT_GRAY,
        105
    )

    mouse = pygame.mouse.get_pos()

    cols = 3

    for i, f in enumerate(frogs):
        row = i // cols
        col = i % cols

        x = 160 + col * 300
        y = 200 + row * 175

        rect = pygame.Rect(
            x - 105,
            y - 45,
            210,
            145
        )

        if i == selected_frog:
            pygame.draw.rect(
                screen,
                f["accent"],
                rect,
                border_radius=18
            )

            pygame.draw.rect(
                screen,
                WHITE,
                rect,
                3,
                border_radius=18
            )
        else:
            pygame.draw.rect(
                screen,
                DARK,
                rect,
                border_radius=18
            )

            pygame.draw.rect(
                screen,
                GRAY,
                rect,
                2,
                border_radius=18
            )

        draw_frog_preview(
            x,
            y + 5,
            i,
            0.8
        )

        text = FONT_SMALL.render(
            f["name"],
            True,
            WHITE
        )

        screen.blit(
            text,
            (
                x - text.get_width() // 2,
                y + 52
            )
        )

    back = pygame.Rect(
        30,
        HEIGHT - 65,
        150,
        45
    )

    draw_button(
        back,
        "BACK",
        GRAY,
        mouse
    )


def draw_level_select():
    screen.fill(DARKER)

    draw_text_center(
        "SELECT LEVEL",
        FONT_BIG,
        CYAN,
        35
    )

    mouse = pygame.mouse.get_pos()

    for i, level in enumerate(LEVELS):
        y = 135 + i * 105

        rect = pygame.Rect(
            WIDTH // 2 - 330,
            y,
            660,
            80
        )

        hover = rect.collidepoint(mouse)

        color = level["color"]

        if hover:
            pygame.draw.rect(
                screen,
                tuple(
                    min(255, c + 20)
                    for c in color
                ),
                rect,
                border_radius=14
            )
        else:
            pygame.draw.rect(
                screen,
                DARK,
                rect,
                border_radius=14
            )

        pygame.draw.rect(
            screen,
            color,
            rect,
            3,
            border_radius=14
        )

        name = FONT_MED.render(
            level["name"],
            True,
            WHITE
        )

        screen.blit(
            name,
            (
                rect.x + 25,
                rect.y + 12
            )
        )

        diff_color = {
            "EASY": GREEN,
            "NORMAL": CYAN,
            "HARD": ORANGE,
            "INSANE": RED
        }[level["difficulty"]]

        diff = FONT_SMALL.render(
            level["difficulty"],
            True,
            diff_color
        )

        screen.blit(
            diff,
            (
                rect.right - diff.get_width() - 25,
                rect.y + 25
            )
        )

    back = pygame.Rect(
        30,
        HEIGHT - 65,
        150,
        45
    )

    draw_button(
        back,
        "BACK",
        GRAY,
        mouse
    )


def draw_hud():
    level = LEVELS[selected_level]

    # Progress
    progress = min(
        1,
        level_distance / level["length"]
    )

    pygame.draw.rect(
        screen,
        (40, 45, 65),
        (220, 18, 660, 15),
        border_radius=7
    )

    pygame.draw.rect(
        screen,
        level["color"],
        (
            220,
            18,
            int(660 * progress),
            15
        ),
        border_radius=7
    )

    # Level name
    text = FONT_SMALL.render(
        level["name"],
        True,
        WHITE
    )

    screen.blit(
        text,
        (20, 15)
    )

    # Coin counter
    coin_text = FONT_SMALL.render(
        f"COINS: {level_coins}",
        True,
        YELLOW
    )

    screen.blit(
        coin_text,
        (20, 48)
    )

    # Frog
    frog_text = FONT_TINY.render(
        frogs[selected_frog]["name"],
        True,
        frogs[selected_frog]["accent"]
    )

    screen.blit(
        frog_text,
        (
            WIDTH - frog_text.get_width() - 20,
            18
        )
    )

    # Pause
    pause_text = FONT_TINY.render(
        "ESC = PAUSE",
        True,
        GRAY
    )

    screen.blit(
        pause_text,
        (
            WIDTH - pause_text.get_width() - 20,
            45
        )
    )


def draw_pause():
    overlay = pygame.Surface(
        (WIDTH, HEIGHT),
        pygame.SRCALPHA
    )

    overlay.fill(
        (0, 0, 0, 165)
    )

    screen.blit(
        overlay,
        (0, 0)
    )

    draw_text_center(
        "PAUSED",
        FONT_BIG,
        WHITE,
        170
    )

    draw_text_center(
        "Press ESC to continue",
        FONT_SMALL,
        LIGHT_GRAY,
        250
    )

    mouse = pygame.mouse.get_pos()

    resume = pygame.Rect(
        WIDTH // 2 - 150,
        320,
        300,
        60
    )

    restart = pygame.Rect(
        WIDTH // 2 - 150,
        400,
        300,
        60
    )

    menu = pygame.Rect(
        WIDTH // 2 - 150,
        480,
        300,
        60
    )

    draw_button(
        resume,
        "RESUME",
        GREEN,
        mouse
    )

    draw_button(
        restart,
        "RESTART",
        ORANGE,
        mouse
    )

    draw_button(
        menu,
        "MAIN MENU",
        PURPLE,
        mouse
    )


def draw_game_over():
    overlay = pygame.Surface(
        (WIDTH, HEIGHT),
        pygame.SRCALPHA
    )

    overlay.fill(
        (5, 5, 15, 210)
    )

    screen.blit(
        overlay,
        (0, 0)
    )

    draw_text_center(
        "OH NO!",
        FONT_BIG,
        RED,
        140
    )

    draw_text_center(
        "The frog hit an obstacle!",
        FONT_SMALL,
        WHITE,
        215
    )

    mouse = pygame.mouse.get_pos()

    retry = pygame.Rect(
        WIDTH // 2 - 160,
        300,
        320,
        65
    )

    level_select = pygame.Rect(
        WIDTH // 2 - 160,
        385,
        320,
        65
    )

    menu = pygame.Rect(
        WIDTH // 2 - 160,
        470,
        320,
        65
    )

    draw_button(
        retry,
        "TRY AGAIN",
        GREEN,
        mouse
    )

    draw_button(
        level_select,
        "LEVEL SELECT",
        CYAN,
        mouse
    )

    draw_button(
        menu,
        "MAIN MENU",
        PURPLE,
        mouse
    )


def draw_level_complete():
    overlay = pygame.Surface(
        (WIDTH, HEIGHT),
        pygame.SRCALPHA
    )

    overlay.fill(
        (5, 15, 10, 200)
    )

    screen.blit(
        overlay,
        (0, 0)
    )

    draw_text_center(
        "FROG-TASTIC!",
        FONT_BIG,
        GREEN,
        110
    )

    draw_text_center(
        "LEVEL COMPLETE!",
        FONT_MED,
        CYAN,
        185
    )

    draw_text_center(
        f"Coins collected: {level_coins}",
        FONT_SMALL,
        YELLOW,
        240
    )

    mouse = pygame.mouse.get_pos()

    next_button = pygame.Rect(
        WIDTH // 2 - 170,
        310,
        340,
        60
    )

    levels_button = pygame.Rect(
        WIDTH // 2 - 170,
        390,
        340,
        60
    )

    menu_button = pygame.Rect(
        WIDTH // 2 - 170,
        470,
        340,
        60
    )

    draw_button(
        next_button,
        "NEXT LEVEL",
        GREEN,
        mouse
    )

    draw_button(
        levels_button,
        "LEVEL SELECT",
        CYAN,
        mouse
    )

    draw_button(
        menu_button,
        "MAIN MENU",
        PURPLE,
        mouse
    )


# -------------------- GAME CONTROL --------------------

def start_level(index):
    global selected_level
    global camera_x
    global level_distance
    global level_coins
    global level_start_time
    global game_state

    selected_level = index

    generate_level(index)

    frog.reset()

    camera_x = 0
    level_distance = 0
    level_coins = 0
    level_start_time = pygame.time.get_ticks()

    particles.clear()

    game_state = PLAYING


def kill_frog():
    global game_state

    if not frog.alive:
        return

    frog.alive = False

    for _ in range(30):
        particles.append(
            Particle(
                frog.x + frog.size / 2,
                frog.y + frog.size / 2,
                frogs[selected_frog]["color"],
                vx=random.uniform(-5, 5),
                vy=random.uniform(-6, 2),
                life=random.randint(25, 50),
                size=random.randint(2, 6)
            )
        )

    game_state = GAME_OVER


def update_game():
    global camera_x
    global level_distance
    global level_coins
    global game_state

    level = LEVELS[selected_level]

    frog.update(blocks)

    # Automatic forward movement
    frog.x += level["speed"]
    level_distance = frog.x

    # Camera follows frog
    target_camera = frog.x - 180

    camera_x += (
        target_camera - camera_x
    ) * 0.12

    if camera_x < 0:
        camera_x = 0

    # Spike collision
    frog_rect = frog.rect()

    for spike in spikes:
        if frog_rect.colliderect(
            spike.rect()
        ):
            kill_frog()
            return

    # Coin collection
    for coin in coins:
        if not coin.collected:
            coin.update()

            coin_rect = pygame.Rect(
                coin.x - 15,
                coin.y - 15,
                30,
                30
            )

            if frog_rect.colliderect(
                coin_rect
            ):
                coin.collected = True
                level_coins += 1

                for _ in range(10):
                    particles.append(
                        Particle(
                            coin.x,
                            coin.y,
                            YELLOW,
                            life=30,
                            size=random.randint(2, 4)
                        )
                    )

    # Update particles
    for p in particles[:]:
        p.update()

        if p.life <= 0:
            particles.remove(p)

    # Finish
    if frog.x >= level["length"]:
        for _ in range(45):
            particles.append(
                Particle(
                    frog.x - camera_x,
                    frog.y + 20,
                    level["color"],
                    vx=random.uniform(-4, 4),
                    vy=random.uniform(-6, 2),
                    life=random.randint(30, 70),
                    size=random.randint(2, 6)
                )
            )

        game_state = LEVEL_COMPLETE


def draw_game():
    draw_background()

    # Blocks
    for block in blocks:
        if (
            block.x - camera_x < WIDTH
            and block.x + block.width - camera_x > 0
        ):
            block.draw(
                camera_x,
                LEVELS[selected_level]["color"]
            )

    # Spikes
    for spike in spikes:
        if (
            spike.x - camera_x < WIDTH
            and spike.x + spike.size - camera_x > 0
        ):
            spike.draw(
                camera_x,
                RED
            )

    # Coins
    for coin in coins:
        if (
            coin.x - camera_x < WIDTH
            and coin.x - camera_x > -50
        ):
            coin.draw(camera_x)

    # Particles behind frog
    for p in particles:
        p.draw()

    frog.draw()

    draw_hud()


# -------------------- EVENTS --------------------

def handle_menu_click(pos):
    global game_state

    mouse_pos = pos

    play_rect = pygame.Rect(
        WIDTH // 2 - 160,
        310,
        320,
        65
    )

    frog_rect = pygame.Rect(
        WIDTH // 2 - 160,
        395,
        320,
        65
    )

    quit_rect = pygame.Rect(
        WIDTH // 2 - 160,
        480,
        320,
        65
    )

    if play_rect.collidepoint(mouse_pos):
        game_state = LEVEL_SELECT

    elif frog_rect.collidepoint(mouse_pos):
        game_state = FROG_SELECT

    elif quit_rect.collidepoint(mouse_pos):
        pygame.quit()
        sys.exit()


def handle_frog_select_click(pos):
    global selected_frog
    global game_state

    mouse_pos = pos

    cols = 3

    for i in range(len(frogs)):
        row = i // cols
        col = i % cols

        x = 160 + col * 300
        y = 200 + row * 175

        rect = pygame.Rect(
            x - 105,
            y - 45,
            210,
            145
        )

        if rect.collidepoint(mouse_pos):
            selected_frog = i
            return

    back = pygame.Rect(
        30,
        HEIGHT - 65,
        150,
        45
    )

    if back.collidepoint(mouse_pos):
        game_state = MENU


def handle_level_select_click(pos):
    global game_state

    mouse_pos = pos

    for i in range(len(LEVELS)):
        y = 135 + i * 105

        rect = pygame.Rect(
            WIDTH // 2 - 330,
            y,
            660,
            80
        )

        if rect.collidepoint(mouse_pos):
            start_level(i)
            return

    back = pygame.Rect(
        30,
        HEIGHT - 65,
        150,
        45
    )

    if back.collidepoint(mouse_pos):
        game_state = MENU


def handle_pause_click(pos):
    global game_state

    mouse_pos = pos

    resume = pygame.Rect(
        WIDTH // 2 - 150,
        320,
        300,
        60
    )

    restart = pygame.Rect(
        WIDTH // 2 - 150,
        400,
        300,
        60
    )

    menu = pygame.Rect(
        WIDTH // 2 - 150,
        480,
        300,
        60
    )

    if resume.collidepoint(mouse_pos):
        game_state = PLAYING

    elif restart.collidepoint(mouse_pos):
        start_level(selected_level)

    elif menu.collidepoint(mouse_pos):
        game_state = MENU


def handle_game_over_click(pos):
    global game_state

    mouse_pos = pos

    retry = pygame.Rect(
        WIDTH // 2 - 160,
        300,
        320,
        65
    )

    level_select = pygame.Rect(
        WIDTH // 2 - 160,
        385,
        320,
        65
    )

    menu = pygame.Rect(
        WIDTH // 2 - 160,
        470,
        320,
        65
    )

    if retry.collidepoint(mouse_pos):
        start_level(selected_level)

    elif level_select.collidepoint(mouse_pos):
        game_state = LEVEL_SELECT

    elif menu.collidepoint(mouse_pos):
        game_state = MENU


def handle_complete_click(pos):
    global game_state

    mouse_pos = pos

    next_button = pygame.Rect(
        WIDTH // 2 - 170,
        310,
        340,
        60
    )

    levels_button = pygame.Rect(
        WIDTH // 2 - 170,
        390,
        340,
        60
    )

    menu_button = pygame.Rect(
        WIDTH // 2 - 170,
        470,
        340,
        60
    )

    if next_button.collidepoint(mouse_pos):
        if selected_level < len(LEVELS) - 1:
            start_level(selected_level + 1)
        else:
            game_state = LEVEL_SELECT

    elif levels_button.collidepoint(mouse_pos):
        game_state = LEVEL_SELECT

    elif menu_button.collidepoint(mouse_pos):
        game_state = MENU


# -------------------- MAIN LOOP --------------------

running = True

while running:

    # --------------------------------------------------------
    # EVENTS
    # --------------------------------------------------------

    for event in pygame.event.get():

        if event.type == pygame.QUIT:
            running = False

        # Keyboard
        if event.type == pygame.KEYDOWN:

            # Universal quit
            if event.key == pygame.K_q:
                if game_state != PLAYING:
                    running = False

            # MENU
            if game_state == MENU:

                if event.key in (
                    pygame.K_RETURN,
                    pygame.K_SPACE
                ):
                    game_state = LEVEL_SELECT

                elif event.key == pygame.K_f:
                    game_state = FROG_SELECT

            # FROG SELECT
            elif game_state == FROG_SELECT:

                if event.key == pygame.K_ESCAPE:
                    game_state = MENU

                elif event.key == pygame.K_LEFT:
                    selected_frog = (
                        selected_frog - 1
                    ) % len(frogs)

                elif event.key == pygame.K_RIGHT:
                    selected_frog = (
                        selected_frog + 1
                    ) % len(frogs)

                elif event.key == pygame.K_RETURN:
                    game_state = LEVEL_SELECT

            # LEVEL SELECT
            elif game_state == LEVEL_SELECT:

                if event.key == pygame.K_ESCAPE:
                    game_state = MENU

                elif event.key == pygame.K_1:
                    start_level(0)

                elif event.key == pygame.K_2:
                    start_level(1)

                elif event.key == pygame.K_3:
                    start_level(2)

                elif event.key == pygame.K_4:
                    start_level(3)

            # PLAYING
            elif game_state == PLAYING:

                if event.key in (
                    pygame.K_SPACE,
                    pygame.K_UP,
                    pygame.K_w
                ):
                    frog.jump()

                elif event.key == pygame.K_ESCAPE:
                    game_state = PAUSED

            # PAUSED
            elif game_state == PAUSED:

                if event.key == pygame.K_ESCAPE:
                    game_state = PLAYING

                elif event.key == pygame.K_r:
                    start_level(selected_level)

            # GAME OVER
            elif game_state == GAME_OVER:

                if event.key in (
                    pygame.K_RETURN,
                    pygame.K_r,
                    pygame.K_SPACE
                ):
                    start_level(selected_level)

                elif event.key == pygame.K_ESCAPE:
                    game_state = LEVEL_SELECT

            # LEVEL COMPLETE
            elif game_state == LEVEL_COMPLETE:

                if event.key == pygame.K_RETURN:
                    if selected_level < len(LEVELS) - 1:
                        start_level(selected_level + 1)
                    else:
                        game_state = LEVEL_SELECT

                elif event.key == pygame.K_ESCAPE:
                    game_state = LEVEL_SELECT

        # Mouse
        if event.type == pygame.MOUSEBUTTONDOWN:

            if event.button == 1:

                if game_state == MENU:
                    handle_menu_click(event.pos)

                elif game_state == FROG_SELECT:
                    handle_frog_select_click(
                        event.pos
                    )

                elif game_state == LEVEL_SELECT:
                    handle_level_select_click(
                        event.pos
                    )

                elif game_state == PAUSED:
                    handle_pause_click(
                        event.pos
                    )

                elif game_state == GAME_OVER:
                    handle_game_over_click(
                        event.pos
                    )

                elif game_state == LEVEL_COMPLETE:
                    handle_complete_click(
                        event.pos
                    )

                elif game_state == PLAYING:
                    frog.jump()

    # --------------------------------------------------------
    # UPDATE
    # --------------------------------------------------------

    if game_state == PLAYING:
        update_game()

    # --------------------------------------------------------
    # DRAW
    # --------------------------------------------------------

    if game_state == MENU:
        draw_menu()

    elif game_state == FROG_SELECT:
        draw_frog_select()

    elif game_state == LEVEL_SELECT:
        draw_level_select()

    elif game_state == PLAYING:
        draw_game()

    elif game_state == PAUSED:
        draw_game()
        draw_pause()

    elif game_state == GAME_OVER:
        draw_game()
        draw_game_over()

    elif game_state == LEVEL_COMPLETE:
        draw_game()
        draw_level_complete()

    pygame.display.flip()
    clock.tick(FPS)


pygame.quit()
sys.exit()