import pygame
import sys

# Initialize Pygame
pygame.init()

# Set up the screen
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Two Characters Game")

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

# Player Class
class Player(pygame.sprite.Sprite):
    def __init__(self, x, y, color):
        super().__init__()
        self.image = pygame.Surface((50, 50))
        self.image.fill(color)
        self.rect = self.image.get_rect()
        self.rect.x = x
        self.rect.y = y
        self.velocity_y = 0
        self.gravity = 0.5  # Increased gravity for a faster fall
        self.max_health = 100
        self.health = self.max_health
        self.projectile_speed = 7
        self.projectile_cooldown = 10  # Reduced cooldown for faster firing rate
        self.jump_power = -15  # Increased jump power
        self.direction = "right"  # Initial direction

    def update(self):
        self.velocity_y += self.gravity
        self.rect.y += self.velocity_y

        # Keep the player above the platform
        if self.rect.bottom > SCREEN_HEIGHT:
            self.rect.bottom = SCREEN_HEIGHT
            self.velocity_y = 0
        elif self.rect.top < 0:
            self.rect.top = 0

        # Apply wrap-around movement for left and right edges
        if self.rect.right < 0:
            self.rect.left = SCREEN_WIDTH
        elif self.rect.left > SCREEN_WIDTH:
            self.rect.right = 0

        # Decrease projectile cooldown
        if self.projectile_cooldown > 0:
            self.projectile_cooldown -= 1

    def collide_with_platform(self, platform):
        if self.rect.colliderect(platform.rect):
            # Adjust player position to prevent sticking to platform
            if self.velocity_y > 0:  # Player is falling
                self.rect.bottom = platform.rect.top
                self.velocity_y = 0
            elif self.velocity_y < 0:  # Player is jumping
                self.rect.top = platform.rect.bottom
                self.velocity_y = 0

    def draw_health_bar(self):
        # Calculate width of health bar based on health percentage
        health_width = (self.health / self.max_health) * self.rect.width
        health_bar = pygame.Rect(self.rect.x, self.rect.y - 10, health_width, 5)
        pygame.draw.rect(screen, GREEN, health_bar)
        # Draw border around health bar
        pygame.draw.rect(screen, WHITE, health_bar, 1)

    def jump(self):
        # Allow jumping regardless of player's position
        self.velocity_y = self.jump_power

    def shoot_projectile(self):
        if self.projectile_cooldown == 0:
            projectile = Projectile(self.rect.centerx, self.rect.centery, RED, self.direction, self)
            all_sprites.add(projectile)
            projectiles.add(projectile)
            self.projectile_cooldown = 10  # Reset cooldown for next shot

# Projectile class
class Projectile(pygame.sprite.Sprite):
    def __init__(self, x, y, color, direction, source):
        super().__init__()
        self.image = pygame.Surface((10, 10))
        self.image.fill(color)
        self.rect = self.image.get_rect()
        self.rect.center = (x, y)
        self.speed = 7
        self.direction = direction
        self.source = source  # Player that shot the projectile
        self.lifespan = 5 * 60  # 5 seconds lifespan at 60 FPS

def update(self):
    global score_player1, score_player2  # Add this line

    # Move the projectile based on its direction
    if self.direction == "right":
        self.rect.x += self.speed
    elif self.direction == "left":
        self.rect.x -= self.speed

    # Wrap-around movement for bullets
    if self.rect.right < 0:
        self.rect.left = SCREEN_WIDTH
    elif self.rect.left > SCREEN_WIDTH:
        self.rect.right = 0

    # Decrease lifespan
    self.lifespan -= 1
    if self.lifespan <= 0:
        self.kill()  # Destroy the bullet if it reaches its lifespan

    # Check collision with players
    hits = pygame.sprite.spritecollide(self, players_group, False)
    for hit in hits:
        # Make sure the projectile doesn't hit its source player
        if hit != self.source:
            hit.health -= 10  # Decrease health
            if hit.health <= 0:  # Check if the player has lost all health
                if hit == player1:
                    score_player2 += 1
                else:
                    score_player1 += 1
            self.kill()  # Destroy projectile


# 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(WHITE)
        self.rect = self.image.get_rect()
        self.rect.x = x
        self.rect.y = y

# Create players
player1 = Player(100, 100, WHITE)
player2 = Player(600, 100, GREEN)

# Create platforms
platform1 = Platform(100, 300, 300, 20)
platform2 = Platform(400, 450, 200, 20)
platform3 = Platform(0, 200, 200, 20)
platform4 = Platform(600, 200, 200, 20)

# Group for all sprites
all_sprites = pygame.sprite.Group()
all_sprites.add(player1)
all_sprites.add(player2)
all_sprites.add(platform1)
all_sprites.add(platform2)
all_sprites.add(platform3)
all_sprites.add(platform4)

# Group for projectiles
projectiles = pygame.sprite.Group()

# Group for players
players_group = pygame.sprite.Group()
players_group.add(player1)
players_group.add(player2)
score_player1 = 0
score_player2 = 0

# Menu options
options = ["Return to Title Page", "Exit"]

def return_to_title_page():
    global score_player1, score_player2
    score_player1 = 0
    score_player2 = 0
    title_page_loop()

def display_menu():
    screen.fill(WHITE)
    title_text = font.render("Game Over - Options", True, BLACK)
    title_rect = title_text.get_rect(center=(SCREEN_WIDTH // 2, 100))
    screen.blit(title_text, title_rect)
    
    for i, option in enumerate(options):
        text = font.render(option, True, BLACK)
        text_rect = text.get_rect(center=(SCREEN_WIDTH // 2, 250 + i * 50))
        screen.blit(text, text_rect)

def title_page_loop():
    # Main loop for the title page
    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
                pygame.quit()
                sys.exit()
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 1:
                    pos = pygame.mouse.get_pos()
                    for i, _ in enumerate(options):
                        if (SCREEN_WIDTH // 2 - 100) < pos[0] < (SCREEN_WIDTH // 2 + 100) and (225 + 50 * i) < pos[1] < (275 + 50 * i):
                            selected_option = options[i]
                            if selected_option == "Exit":
                                pygame.quit()
                                sys.exit()
                            else:
                                return_to_title_page()

        display_menu()
        pygame.display.flip()

# Main loop
running = True
clock = pygame.time.Clock()
game_over = False  # Flag to indicate game over

while running:
    # Event handling
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:  # Check if Escape key is pressed
                return_to_title_page()  # Call the function to return to the title page

            if event.key == pygame.K_SPACE and game_over:  # Restart the game if spacebar is pressed and game is over
                for player in players_group:  # Reset player health
                    player.health = player.max_health
                all_sprites.empty()  # Clear all sprites
                projectiles.empty()  # Clear projectiles
                all_sprites.add(player1)  # Add players and platform back to all_sprites
                all_sprites.add(player2)
                all_sprites.add(platform1)
                all_sprites.add(platform2)
                all_sprites.add(platform3)
                all_sprites.add(platform4)
                game_over = False  # Reset game over flag

            # Player controls for jumping
            if event.key == pygame.K_UP:  # Player 2 jump
                player2.jump()
            if event.key == pygame.K_w:  # Player 1 jump
                player1.jump()

    if not game_over:  # Only update game state if the game is not over
        # Player movement
        keys = pygame.key.get_pressed()
        if keys[pygame.K_LEFT]:
            player2.rect.x -= 5
            player2.direction = "left"
        if keys[pygame.K_RIGHT]:
            player2.rect.x += 5
            player2.direction = "right"
        if keys[pygame.K_RCTRL]:
            player2.shoot_projectile()

        if keys[pygame.K_a]:
            player1.rect.x -= 5
            player1.direction = "left"
        if keys[pygame.K_d]:
            player1.rect.x += 5
            player1.direction = "right"
        if keys[pygame.K_q]:
            player1.shoot_projectile()

        # Apply gravity and handle collisions with platforms
        player1.update()
        for platform in [platform1, platform2, platform3, platform4]:
            player1.collide_with_platform(platform)
        player2.update()
        for platform in [platform1, platform2, platform3, platform4]:
            player2.collide_with_platform(platform)

        # Update projectiles
        projectiles.update()

        # Check if a player has lost all health
        for player in players_group:
            if player.health <= 0:
                game_over = True
                if player == player1:
                    score_player2 += 1
                else:
                    score_player1 += 1

    # Draw everything
    screen.fill(BLACK)
    all_sprites.draw(screen)
    projectiles.draw(screen)

    # Draw health bars
    player1.draw_health_bar()
    player2.draw_health_bar()

    # Display scores
    font = pygame.font.SysFont(None, 36)
    score_text1 = font.render(f"Player 1: {score_player1}", True, WHITE)
    score_text2 = font.render(f"Player 2: {score_player2}", True, WHITE)
    screen.blit(score_text1, (10, 10))
    screen.blit(score_text2, (SCREEN_WIDTH - score_text2.get_width() - 10, 10))

    if game_over:  # Display game over message
        font = pygame.font.SysFont(None, 48)
        text = font.render("Game Over", True, WHITE)
        text_rect = text.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2))
        screen.blit(text, text_rect)

        # Display "Press space to restart"
        restart_text = font.render("(Press space to restart)", True, WHITE)
        restart_rect = restart_text.get_rect(center=(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2 + 50))
        screen.blit(restart_text, restart_rect)

    # Update the display
    pygame.display.flip()

    # Cap the frame rate
    clock.tick(60)

pygame.quit()
sys.exit()