import pygame
import sys
import random

# Initialize Pygame and mixer
pygame.init()
pygame.mixer.init()

# Screen setup
WIDTH, HEIGHT = 800, 400
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Jump Game+ Enhanced")
clock = pygame.time.Clock()

# Colors
WHITE = (255, 255, 255)
GROUND_COLOR = (50, 200, 50)
PLAYER_COLOR = (0, 100, 255)
OBSTACLE_COLOR = (200, 50, 50)
COIN_COLOR = (255, 215, 0)
POWERUP_COLOR = (255, 223, 0)
CRAZY_COLOR = (255, 0, 0)
ENEMY_COLOR = (255, 165, 0)
HEALTH_COLOR = (255, 0, 0)

# Fonts
font = pygame.font.SysFont(None, 48)
small_font = pygame.font.SysFont(None, 24)

# Game variables
gravity = 0.8
jump_strength = -15
double_jump_strength = -12
fly_jump_strength = -10
game_active = True
score = 0
on_ground = True
can_double_jump = True
is_flying = False
fly_duration = 0
health = 3
level = 1
combo_count = 0

# Player setup
player = pygame.Rect(100, HEIGHT - 100, 40, 50)
player_velocity = 0

# Obstacles setup
obstacles = []
obstacle_speed = 5

# Power-ups setup
powerups = []

# Enemies setup
enemies = []

# Coins setup
coins = []

# Sounds
jump_sound = pygame.mixer.Sound('die.wav')
coin_sound = pygame.mixer.Sound('die.wav')
powerup_sound = pygame.mixer.Sound('die.wav')
crazy_sound = pygame.mixer.Sound('die.wav')
enemy_hit_sound = pygame.mixer.Sound('die.wav')
background_music = pygame.mixer.music.load('background_music.mp3')
pygame.mixer.music.play(-1, 0.0)

# Visual effects (screen shake)
shake_time = 0

def draw_ground():
    pygame.draw.rect(screen, GROUND_COLOR, (0, HEIGHT - 50, WIDTH, 50))

def draw_text(text, x, y, color=(0, 0, 0), font_type=font):
    render = font_type.render(text, True, color)
    screen.blit(render, (x, y))

def create_obstacle():
    height = random.randint(30, 70)
    obstacle = pygame.Rect(WIDTH, HEIGHT - 50 - height, 40, height)
    obstacles.append(obstacle)

def create_coin():
    coin = pygame.Rect(WIDTH + random.randint(100, 500), random.randint(100, HEIGHT - 50 - 50), 20, 20)
    coins.append(coin)

def create_powerup():
    powerup_type = random.choice(["fly", "crazy", "shield", "speed"])
    x_position = WIDTH + random.randint(100, 500)
    y_position = random.randint(100, HEIGHT - 50 - 50)
    powerup = pygame.Rect(x_position, y_position, 30, 30)
    powerups.append((powerup, powerup_type))

def create_enemy():
    enemy = pygame.Rect(WIDTH, HEIGHT - 50 - random.randint(30, 70), 40, 50)
    enemies.append(enemy)

def reset_game():
    global score, obstacles, coins, powerups, player_velocity, on_ground, can_double_jump, fly_powerup, crazy_powerup, fly_duration, game_active, health, level, combo_count
    player.x = 100
    player.y = HEIGHT - 100
    player_velocity = 0
    obstacles.clear()
    coins.clear()
    powerups.clear()
    enemies.clear()
    score = 0
    health = 3
    level = 1
    combo_count = 0
    on_ground = True
    can_double_jump = True
    is_flying = False
    fly_duration = 0
    crazy_powerup = False
    obstacle_speed = 5
    game_active = True

def shake_screen():
    global shake_time
    if shake_time > 0:
        offset = random.randint(-5, 5)
        shake_time -= 1
        return offset
    return 0

# Game loop
while True:
    screen.fill(WHITE)

    # Shake effect during crazy power-up
    shake_offset = shake_screen()

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        if game_active:
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:
                    if on_ground:
                        player_velocity = jump_strength
                        on_ground = False
                        jump_sound.play()
                    elif can_double_jump:
                        player_velocity = double_jump_strength
                        can_double_jump = False
                        jump_sound.play()
                    elif is_flying:
                        player_velocity = fly_jump_strength
                        fly_duration -= 1
                if event.key == pygame.K_r and not game_active:
                    reset_game()
        else:
            if event.type == pygame.KEYDOWN and event.key == pygame.K_r:
                reset_game()

    if game_active:
        # Player movement and gravity
        player_velocity += gravity
        if not is_flying:
            player.y += player_velocity
        else:
            player.y += player_velocity / 2  # Flying slows down gravity

        # Handle fly power-up
        if fly_duration <= 0:
            is_flying = False

        # Keep the player from falling through the ground
        if player.y >= HEIGHT - 50 - player.height:
            player.y = HEIGHT - 50 - player.height
            player_velocity = 0
            on_ground = True
            can_double_jump = True

        # Handle obstacles
        for obstacle in obstacles:
            obstacle.x -= obstacle_speed
            if obstacle.x < -obstacle.width:
                obstacles.remove(obstacle)
                create_obstacle()

        # Handle coins
        for coin in coins:
            coin.x -= obstacle_speed
            if coin.x < -coin.width:
                coins.remove(coin)
                create_coin()

        # Handle power-ups
        for powerup, powerup_type in powerups:
            powerup.x -= obstacle_speed
            if powerup.x < -powerup.width:
                powerups.remove((powerup, powerup_type))
                create_powerup()
            if player.colliderect(powerup):
                if powerup_type == "fly":
                    is_flying = True
                    fly_duration = 5  # Fly for 5 seconds
                    powerup_sound.play()
                elif powerup_type == "crazy":
                    crazy_powerup = True
                    obstacle_speed += 1  # Increase obstacle speed
                    crazy_sound.play()
                elif powerup_type == "shield":
                    health += 1  # Increase health
                    powerup_sound.play()
                elif powerup_type == "speed":
                    obstacle_speed += 2  # Increase player speed
                    powerup_sound.play()

        # Handle enemies
        for enemy in enemies:
            enemy.x -= obstacle_speed
            if enemy.x < -enemy.width:
                enemies.remove(enemy)
                create_enemy()
            if player.colliderect(enemy):
                health -= 1
                enemy_hit_sound.play()
                if health <= 0:
                    game_active = False

        # Handle collision with obstacles
        for obstacle in obstacles:
            if player.colliderect(obstacle):
                health -= 1
                if health <= 0:
                    game_active = False

        # Handle collision with coins
        for coin in coins:
            if player.colliderect(coin):
                coins.remove(coin)
                coin_sound.play()
                score += 1
                combo_count += 1
                create_coin()

        # Create new obstacles and items periodically
        if random.random() < 0.02:
            create_obstacle()
        if random.random() < 0.03:
            create_coin()
        if random.random() < 0.01:
            create_powerup()
        if random.random() < 0.01:
            create_enemy()

        # Draw player, obstacles, coins, power-ups, and enemies
        pygame.draw.rect(screen, PLAYER_COLOR, (player.x + shake_offset, player.y, player.width, player.height))  # Apply shake to player position
        for obstacle in obstacles:
            pygame.draw.rect(screen, OBSTACLE_COLOR, obstacle)
        for coin in coins:
            pygame.draw.circle(screen, COIN_COLOR, coin.center, coin.width // 2)
        for powerup, powerup_type in powerups:
            if powerup_type == "fly":
                pygame.draw.rect(screen, POWERUP_COLOR, powerup)
            elif powerup_type == "crazy":
                pygame.draw.rect(screen, CRAZY_COLOR, powerup)
            elif powerup_type == "shield":
                pygame.draw.rect(screen, (0, 255, 0), powerup)
            elif powerup_type == "speed":
                pygame.draw.rect(screen, (0, 0, 255), powerup)
        for enemy in enemies:
            pygame.draw.rect(screen, ENEMY_COLOR, enemy)

        # Draw health and score
        draw_text(f"Score: {score}", 10, 10)
        draw_text(f"Health: {health}", 10, 50)
        draw_text(f"Level: {level}", 10, 90)

    else:
        draw_text("Game Over! Press R to Restart", WIDTH // 2 - 200, HEIGHT // 2 - 30)
        draw_text(f"Final Score: {score}", WIDTH // 2 - 100, HEIGHT // 2 + 30)

    pygame.display.update()
    clock.tick(60)