import pygame
import sys
import math

pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Tea Egg Adventure - Beautiful")
clock = pygame.time.Clock()
FPS = 60

# Colors
WHITE = (255,255,255)
BLACK = (0,0,0)
SKY_TOP = (135,206,235)
SKY_BOTTOM = (200,230,255)
BROWN = (170,120,70)
DARK_BROWN = (120,80,40)
LIGHT_BROWN = (210,170,120)
RED = (220,60,60)
GREEN = (80,180,60)
YELLOW = (255,220,0)
GRASS = (100,200,80)
DIRT = (140,100,60)
GRAY = (150,150,150)
DARK_GRAY = (80,80,80)

# Fonts
font_small = pygame.font.Font(None, 24)
font_medium = pygame.font.Font(None, 36)
font_large = pygame.font.Font(None, 48)

# Constants
PLAYER_SPEED = 5
JUMP_SPEED = -13
GRAVITY = 0.6
PLAYER_W = 36
PLAYER_H = 48
GROUND_Y = HEIGHT - 50

# Predefined platforms for 5 levels (x, y, width, height)
# Ground is always at GROUND_Y.
LEVELS = [
    # Level 1
    [
        (0, GROUND_Y, WIDTH, 30),          # ground
        (200, GROUND_Y-80, 120, 18),
        (400, GROUND_Y-80, 120, 18),
        (600, GROUND_Y-80, 120, 18),
    ],
    # Level 2
    [
        (0, GROUND_Y, WIDTH, 30),
        (150, GROUND_Y-70, 100, 18),
        (350, GROUND_Y-120, 100, 18),
        (550, GROUND_Y-70, 100, 18),
        (700, GROUND_Y-120, 100, 18),
    ],
    # Level 3
    [
        (0, GROUND_Y, WIDTH, 30),
        (100, GROUND_Y-60, 100, 18),
        (300, GROUND_Y-100, 100, 18),
        (500, GROUND_Y-140, 100, 18),
        (700, GROUND_Y-100, 100, 18),
    ],
    # Level 4
    [
        (0, GROUND_Y, WIDTH, 30),
        (120, GROUND_Y-80, 120, 18),
        (320, GROUND_Y-40, 120, 18),
        (520, GROUND_Y-120, 120, 18),
        (720, GROUND_Y-80, 120, 18),
    ],
    # Level 5
    [
        (0, GROUND_Y, WIDTH, 30),
        (80, GROUND_Y-50, 100, 18),
        (250, GROUND_Y-90, 100, 18),
        (420, GROUND_Y-130, 100, 18),
        (590, GROUND_Y-90, 100, 18),
        (750, GROUND_Y-50, 100, 18),
    ],
]

# Spikes for each level (x, y, width, height) placed on platforms
LEVEL_SPIKES = [
    [],  # Level 1 no spikes
    [(400, GROUND_Y-120-18, 26, 18)],  # Level2 spike on platform (350,120)
    [(550, GROUND_Y-140-18, 26, 18)],  # Level3 spike on platform (500,140)
    [(370, GROUND_Y-40-18, 26, 18)],   # Level4 spike on platform (320,40)
    [(300, GROUND_Y-90-18, 26, 18), (640, GROUND_Y-90-18, 26, 18)],  # Level5
]

# Goal positions for each level (x, y)
LEVEL_GOALS = [
    (680, GROUND_Y-80-40),  # Level1 on last platform
    (760, GROUND_Y-120-40), # Level2 on last platform
    (760, GROUND_Y-100-40), # Level3 on last platform
    (780, GROUND_Y-80-40),  # Level4 on last platform
    (790, GROUND_Y-50-40),  # Level5 on last platform
]

class Player:
    def __init__(self, x, y):
        self.rect = pygame.Rect(x, y, PLAYER_W, PLAYER_H)
        self.vx = 0
        self.vy = 0
        self.on_ground = False
        self.facing = 1
        self.jump_particles = []

    def update(self, platforms):
        keys = pygame.key.get_pressed()
        self.vx = 0
        if keys[pygame.K_LEFT] or keys[pygame.K_a]:
            self.vx = -PLAYER_SPEED
            self.facing = -1
        if keys[pygame.K_RIGHT] or keys[pygame.K_d]:
            self.vx = PLAYER_SPEED
            self.facing = 1
        if (keys[pygame.K_UP] or keys[pygame.K_w] or keys[pygame.K_SPACE]) and self.on_ground:
            self.vy = JUMP_SPEED
            self.on_ground = False
            for _ in range(6):
                self.jump_particles.append(Particle(
                    self.rect.centerx + random.randint(-10, 10),
                    self.rect.bottom,
                    (200, 200, 200),
                    vx=random.uniform(-2, 2),
                    vy=random.uniform(-4, -1)
                ))

        self.vy += GRAVITY
        if self.vy > 15:
            self.vy = 15

        self.rect.x += self.vx
        for p in platforms:
            if self.rect.colliderect(p):
                if self.vx > 0:
                    self.rect.right = p.left
                elif self.vx < 0:
                    self.rect.left = p.right

        self.rect.y += self.vy
        self.on_ground = False
        for p in platforms:
            if self.rect.colliderect(p):
                if self.vy > 0:
                    self.rect.bottom = p.top
                    self.vy = 0
                    self.on_ground = True
                elif self.vy < 0:
                    self.rect.top = p.bottom
                    self.vy = 0

        self.rect.left = max(0, self.rect.left)
        self.rect.right = min(WIDTH, self.rect.right)

        self.jump_particles = [pt for pt in self.jump_particles if pt.update()]

    def draw(self, surface):
        # Shadow
        shadow = pygame.Surface((PLAYER_W+10, 10), pygame.SRCALPHA)
        pygame.draw.ellipse(shadow, (0,0,0,60), (0,0,PLAYER_W+10,10))
        surface.blit(shadow, (self.rect.x-5, self.rect.bottom-3))

        # Egg body
        egg_rect = self.rect.inflate(-2, -4)
        pygame.draw.ellipse(surface, BROWN, egg_rect)
        # Highlight
        hl_rect = pygame.Rect(egg_rect.x+5, egg_rect.y+3, egg_rect.width-10, egg_rect.height//3)
        pygame.draw.ellipse(surface, LIGHT_BROWN, hl_rect)
        pygame.draw.ellipse(surface, DARK_BROWN, egg_rect, 3)

        # Face
        cx = self.rect.centerx
        eye_y = self.rect.y + 18
        off = 5 * self.facing

        pygame.draw.circle(surface, WHITE, (cx-6+off, eye_y), 7)
        pygame.draw.circle(surface, WHITE, (cx+6+off, eye_y), 7)
        pygame.draw.circle(surface, BLACK, (cx-5+off, eye_y+1), 3)
        pygame.draw.circle(surface, BLACK, (cx+5+off, eye_y+1), 3)
        # shine
        pygame.draw.circle(surface, WHITE, (cx-6+off, eye_y-1), 1)
        pygame.draw.circle(surface, WHITE, (cx+6+off, eye_y-1), 1)
        # mouth
        pygame.draw.arc(surface, BLACK, (cx-7, self.rect.y+28, 14, 10), 0, math.pi, 2)
        # blush
        pygame.draw.circle(surface, (255,180,160), (cx-12+off, self.rect.y+27), 4)
        pygame.draw.circle(surface, (255,180,160), (cx+12+off, self.rect.y+27), 4)

        for pt in self.jump_particles:
            pt.draw(surface)

class Particle:
    def __init__(self, x, y, color, vx=0, vy=-3):
        self.x = x
        self.y = y
        self.vx = vx
        self.vy = vy
        self.life = 20
        self.max_life = 20
        self.color = color
        self.size = random.randint(3,6)

    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.vy += 0.3
        self.life -= 1
        return self.life > 0

    def draw(self, surface):
        alpha = 255 * self.life // self.max_life
        r = max(1, int(self.size * self.life / self.max_life))
        s = pygame.Surface((r*2, r*2), pygame.SRCALPHA)
        pygame.draw.circle(s, (*self.color, alpha), (r, r), r)
        surface.blit(s, (self.x-r, self.y-r))

class Game:
    def __init__(self):
        self.current_level = 0
        self.load_level()

    def load_level(self):
        platforms_data = LEVELS[self.current_level]
        self.platforms = [pygame.Rect(*p) for p in platforms_data]
        spike_data = LEVEL_SPIKES[self.current_level]
        self.spikes = [pygame.Rect(*s) for s in spike_data]
        self.goal_x, self.goal_y = LEVEL_GOALS[self.current_level]
        self.player = Player(50, GROUND_Y - PLAYER_H)
        self.win = False
        self.dead = False
        self.death_timer = 0
        self.death_particles = []

    def handle_events(self):
        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_ESCAPE:
                    pygame.quit(); sys.exit()
                if event.key == pygame.K_r:
                    self.load_level()

    def update(self):
        if self.dead:
            self.death_timer += 1
            for pt in self.death_particles:
                pt.update()
            if self.death_timer > 90:
                self.load_level()
            return

        if self.win:
            keys = pygame.key.get_pressed()
            if keys[pygame.K_SPACE] or keys[pygame.K_RETURN]:
                self.current_level = (self.current_level + 1) % len(LEVELS)
                self.load_level()
            return

        self.player.update(self.platforms)

        for spike in self.spikes:
            if self.player.rect.colliderect(spike):
                self.dead = True
                self.death_timer = 0
                for _ in range(25):
                    self.death_particles.append(Particle(
                        self.player.rect.centerx,
                        self.player.rect.centery,
                        random.choice([BROWN, DARK_BROWN, LIGHT_BROWN]),
                        vx=random.uniform(-5,5),
                        vy=random.uniform(-8,-2)
                    ))
                break

        if not self.dead:
            if abs(self.player.rect.centerx - self.goal_x) < 30 and abs(self.player.rect.centery - self.goal_y) < 60:
                self.win = True

    def draw_background(self):
        # Sky gradient
        for y in range(HEIGHT):
            t = y / HEIGHT
            r = int(SKY_TOP[0] + (SKY_BOTTOM[0]-SKY_TOP[0])*t)
            g = int(SKY_TOP[1] + (SKY_BOTTOM[1]-SKY_TOP[1])*t)
            b = int(SKY_TOP[2] + (SKY_BOTTOM[2]-SKY_TOP[2])*t)
            pygame.draw.line(screen, (r,g,b), (0,y), (WIDTH,y))

        # Sun
        pygame.draw.circle(screen, (255,240,150), (700,70), 35)
        pygame.draw.circle(screen, (255,250,200), (700,70), 25)

        # Clouds
        for cx, cy in [(120,70), (350,50), (550,90)]:
            pygame.draw.ellipse(screen, WHITE, (cx,cy,90,35))
            pygame.draw.ellipse(screen, WHITE, (cx+20,cy-15,70,30))
            pygame.draw.ellipse(screen, WHITE, (cx+55,cy+5,60,25))

        # Mountains
        for mx in range(0, WIDTH, 200):
            pts = [(mx, GROUND_Y), (mx+100, GROUND_Y-80), (mx+200, GROUND_Y)]
            pygame.draw.polygon(screen, (150,170,190), pts)

    def draw_platforms(self):
        for p in self.platforms:
            # shadow
            pygame.draw.rect(screen, (0,0,0,30), (p.x+3, p.y+3, p.width, p.height))
            # dirt
            pygame.draw.rect(screen, DIRT, p)
            # grass top
            pygame.draw.rect(screen, GRASS, (p.x, p.y, p.width, 7))
            pygame.draw.rect(screen, (50,140,40), (p.x, p.y+7, p.width, 3))
            # outline
            pygame.draw.rect(screen, BLACK, p, 2)

    def draw_spikes(self):
        for spike in self.spikes:
            pygame.draw.rect(screen, GRAY, (spike.x-2, spike.y+spike.height-4, spike.width+4, 4))
            for i in range(3):
                sx = spike.x + i*(spike.width//3)
                sw = spike.width//3
                pts = [(sx, spike.y+spike.height), (sx+sw//2, spike.y), (sx+sw, spike.y+spike.height)]
                pygame.draw.polygon(screen, RED, pts)
                pygame.draw.polygon(screen, BLACK, pts, 1)

    def draw_goal(self):
        fx, fy = self.goal_x, self.goal_y
        pygame.draw.line(screen, BLACK, (fx, fy), (fx, fy+55), 3)
        wave = math.sin(pygame.time.get_ticks()*0.008)*6
        pygame.draw.polygon(screen, YELLOW, [(fx, fy), (fx+38, fy+12+wave), (fx, fy+28)])
        pygame.draw.circle(screen, YELLOW, (fx, fy), 5)

    def draw_ui(self):
        level_text = font_medium.render(f"Level {self.current_level+1}", True, BLACK)
        screen.blit(level_text, (12, 12))
        if self.dead:
            overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
            overlay.fill((0,0,0,140))
            screen.blit(overlay, (0,0))
            dead_text = font_large.render("You Broke!", True, (255,100,100))
            screen.blit(dead_text, (WIDTH//2 - dead_text.get_width()//2, HEIGHT//2-30))
        elif self.win:
            overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
            overlay.fill((0,0,0,140))
            screen.blit(overlay, (0,0))
            win_text = font_large.render("Complete!", True, (100,255,100))
            screen.blit(win_text, (WIDTH//2 - win_text.get_width()//2, HEIGHT//2-30))
            next_text = font_small.render("SPACE for next level", True, WHITE)
            screen.blit(next_text, (WIDTH//2 - next_text.get_width()//2, HEIGHT//2+15))
        else:
            hint = font_small.render("Arrows / WASD move   |   Space jump", True, BLACK)
            screen.blit(hint, (12, HEIGHT-28))

    def draw(self):
        self.draw_background()
        self.draw_platforms()
        self.draw_spikes()
        self.draw_goal()
        if not self.dead:
            self.player.draw(screen)
        else:
            for pt in self.death_particles:
                pt.draw(screen)
        self.draw_ui()
        pygame.display.flip()

    def run(self):
        while True:
            self.handle_events()
            self.update()
            self.draw()
            clock.tick(FPS)

if __name__ == "__main__":
    import random
    game = Game()
    game.run()