import pygame
import random
import sys

# Initialize Pygame
pygame.init()

# Constants
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
GRAVITY = 0.5
JUMP_SPEED = 10  # Base jump speed
MOVE_SPEED = 3
MAX_JUMP_CHARGE = 150  # Maximum jump charge time (milliseconds)

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
BLUE = (0, 0, 255)
GREEN = (0, 255, 0)

# Create the screen
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Jump King (Simplified)")

# Player Class
class Player(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.image = pygame.Surface([20, 30])  # Player rectangle
        self.image.fill(BLUE)
        self.rect = self.image.get_rect()
        self.rect.x = x
        self.rect.y = y
        self.y_velocity = 0
        self.x_velocity = 0
        self.grounded = False
        self.jump_charging = False
        self.jump_charge_time = 0

    def update(self, platforms):
        # Gravity
        self.y_velocity += GRAVITY
        self.rect.y += self.y_velocity

        # Collision detection with platforms
        platform_collisions = pygame.sprite.spritecollide(self, platforms, False)
        if platform_collisions:
            # Resolve collision (basic)
            self.rect.bottom = platform_collisions[0].rect.top
            self.y_velocity = 0
            self.grounded = True
        else:
            self.grounded = False

        # Keep player within screen bounds (horizontal)
        if self.rect.left < 0:
            self.rect.left = 0
            self.x_velocity = 0  # Stop movement at edge
        if self.rect.right > SCREEN_WIDTH:
            self.rect.right = SCREEN_WIDTH
            self.x_velocity = 0  # Stop movement at edge

        # Horizontal movement (only during jump charge)
        self.rect.x += self.x_velocity

    def start_jump(self):
        if self.grounded:
            self.jump_charging = True
            self.jump_charge_time = 0

    def stop_jump(self):
        if self.jump_charging:
            self.jump_charging = False
            jump_power = min(self.jump_charge_time / 100, 1.0)  # Normalize jump power
            self.y_velocity = -JUMP_SPEED * jump_power
            self.x_velocity *= jump_power # Apply jump power to horizontal velocity

    def move_left(self):
        if self.jump_charging:
            self.x_velocity = -MOVE_SPEED

    def move_right(self):
        if self.jump_charging:
            self.x_velocity = MOVE_SPEED

    def stop_moving(self):
        self.x_velocity = 0


# Platform Class
class Platform(pygame.sprite.Sprite):
    def __init__(self, x, y, width, height):
        super().__init__()
        self.image = pygame.Surface([width, height])
        self.image.fill(GREEN)
        self.rect = self.image.get_rect()
        self.rect.x = x
        self.rect.y = y

# Camera Class (Simple)
class Camera:
    def __init__(self, width, height):
        self.camera_rect = pygame.Rect(0, 0, width, height)
        self.width = width
        self.height = height

    def apply(self, entity):
        return entity.rect.move(-self.camera_rect.x, -self.camera_rect.y)

    def update(self, target):
        # Center the camera around the target
        x = -target.rect.x + int(SCREEN_WIDTH / 2)
        y = -target.rect.y + int(SCREEN_HEIGHT / 2)

        # Limit scrolling to map boundaries (if you have them)
        x = min(0, x)  # Left boundary
        y = min(0, y)  # Top boundary
        # Assuming you have map_width and map_height defined
        # x = max(-(map_width - SCREEN_WIDTH), x) # Right boundary
        # y = max(-(map_height - SCREEN_HEIGHT), y) # Bottom boundary

        self.camera_rect = pygame.Rect(x, y, self.width, self.height)


# --- Game Setup ---
player = Player(50, 50)
platform_list = pygame.sprite.Group()
all_sprites = pygame.sprite.Group()
all_sprites.add(player)

# Create some platforms
platform1 = Platform(0, 500, 200, 50)
platform2 = Platform(300, 400, 150, 50)
platform3 = Platform(600, 300, 200, 50)
platform_list.add(platform1, platform2, platform3)
all_sprites.add(platform1, platform2, platform3)

camera = Camera(SCREEN_WIDTH * 2, SCREEN_HEIGHT * 2)  # Example camera size

# Game loop
running = True
clock = pygame.time.Clock()

while running:
    # --- Event Handling ---
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                player.start_jump()
            if event.key == pygame.K_LEFT:
                player.move_left()
            if event.key == pygame.K_RIGHT:
                player.move_right()
        if event.type == pygame.KEYUP:
            if event.key == pygame.K_SPACE:
                player.stop_jump()
            if event.key == pygame.K_LEFT or event.key == pygame.K_RIGHT:
                player.stop_moving()

    # --- Game Logic ---
    if player.jump_charging:
        player.jump_charge_time += clock.get_time()  # Increment charge time
        player.jump_charge_time = min(player.jump_charge_time, MAX_JUMP_CHARGE) # Limit charge time

    player.update(platform_list)
    camera.update(player)

    # Game Over Condition (fall off screen)
    if player.rect.top > SCREEN_HEIGHT:
        # Reset player position (or restart level)
        player.rect.x = 50
        player.rect.y = 50
        player.y_velocity = 0

    # --- Drawing ---
    screen.fill(BLACK)

    for entity in all_sprites:
        screen.blit(entity.image, camera.apply(entity))  # Apply camera offset

    pygame.display.flip()

    # --- FPS ---
    clock.tick(60)

pygame.quit()
sys.exit()