import pygame
import sys

# Game setup
pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Squares")
clock = pygame.time.Clock()

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
BLUE = (50, 50, 255)
RED = (255, 50, 50)

# Constants
GRAVITY = 0.5
PLAYER_SPEED = 5
JUMP_STRENGTH = 15

# Player class
class Player(pygame.sprite.Sprite):
    def __init__(self):
        super().__init__()
        self.image = pygame.Surface((40, 60))
        self.image.fill(BLUE)
        self.rect = self.image.get_rect()
        self.rect.x, self.rect.y = 50, HEIGHT - 100
        self.vel_y = 0
        self.on_ground = False

    def update(self, platforms):
        keys = pygame.key.get_pressed()
        dx = 0
        if keys[pygame.K_LEFT]:
            dx -= PLAYER_SPEED
        if keys[pygame.K_RIGHT]:
            dx += PLAYER_SPEED

        self.vel_y += GRAVITY
        dy = self.vel_y

        # Move horizontally and check for collisions
        self.rect.x += dx
        for platform in platforms:
            if self.rect.colliderect(platform.rect):
                if dx > 0:
                    self.rect.right = platform.rect.left
                if dx < 0:
                    self.rect.left = platform.rect.right

        # Move vertically and check for collisions
        self.rect.y += dy
        self.on_ground = False
        for platform in platforms:
            if self.rect.colliderect(platform.rect):
                if dy > 0:
                    self.rect.bottom = platform.rect.top
                    self.vel_y = 0
                    self.on_ground = True
                elif dy < 0:
                    self.rect.top = platform.rect.bottom
                    self.vel_y = 0

    def jump(self):
        if self.on_ground:
            self.vel_y = -JUMP_STRENGTH

# Platform class
class Platform(pygame.sprite.Sprite):
    def __init__(self, x, y, w, h):
        super().__init__()
        self.image = pygame.Surface((w, h))
        self.image.fill(BLACK)
        self.rect = self.image.get_rect(topleft=(x, y))

# Enemy class (moves left and right)
class Enemy(pygame.sprite.Sprite):
    def __init__(self, x, y, range_):
        super().__init__()
        self.image = pygame.Surface((40, 40))
        self.image.fill(RED)
        self.rect = self.image.get_rect(topleft=(x, y))
        self.start_x = x
        self.range = range_
        self.direction = 1
        self.speed = 2

    def update(self):
        self.rect.x += self.direction * self.speed
        if self.rect.x < self.start_x or self.rect.x > self.start_x + self.range:
            self.direction *= -1

# Levels setup
levels = [
    # Level 1: Easy platforms, no enemies
    [(Platform(0, HEIGHT - 40, WIDTH, 40),), ()],

    # Level 2: Platforms floating
    [(Platform(0, HEIGHT - 40, WIDTH, 40), Platform(200, 400, 100, 20), Platform(400, 300, 100, 20)), ()],

    # Level 3: More platforms + enemy
    [(Platform(0, HEIGHT - 40, WIDTH, 40), Platform(200, 400, 100, 20), Platform(400, 300, 100, 20)), (Enemy(300, HEIGHT - 80, 200),)],

    # Level 4: Tight jumps and enemies
    [(Platform(0, HEIGHT - 40, WIDTH, 40), Platform(100, 450, 100, 20), Platform(300, 350, 100, 20), Platform(500, 250, 100, 20)),
     (Enemy(150, 410, 100), Enemy(320, 310, 80))],

    # Level 5: Hell mode
    [(Platform(0, HEIGHT - 40, WIDTH, 40), Platform(100, 500, 80, 20), Platform(250, 400, 80, 20), Platform(400, 300, 80, 20), Platform(550, 200, 80, 20)),
     (Enemy(110, 460, 60), Enemy(260, 360, 60), Enemy(410, 260, 60))],
]

level_index = 0

# Game loop
player = Player()
all_sprites = pygame.sprite.Group()
platforms = pygame.sprite.Group()
enemies = pygame.sprite.Group()

def load_level(index):
    global all_sprites, platforms, enemies
    all_sprites.empty()
    platforms.empty()
    enemies.empty()
    player.rect.x, player.rect.y = 50, HEIGHT - 100
    all_sprites.add(player)
    for plat in levels[index][0]:
        platforms.add(plat)
        all_sprites.add(plat)
    for en in levels[index][1]:
        enemies.add(en)
        all_sprites.add(en)

load_level(level_index)

running = True
while running:
    clock.tick(60)
    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_SPACE:
                player.jump()

    player.update(platforms)
    enemies.update()

    # Check collision with enemies
    if pygame.sprite.spritecollideany(player, enemies):
        print("Ouch! Restarting level.")
        load_level(level_index)

    # Check level complete
    if player.rect.right > WIDTH:
        level_index += 1
        if level_index >= len(levels):
            print("You beat all levels!")
            pygame.quit()
            sys.exit()
        else:
            print(f"Level {level_index + 1} starting!")
            load_level(level_index)

    # Draw
    screen.fill(WHITE)
    all_sprites.draw(screen)
    pygame.display.flip()