#!/usr/bin/env python3
import pygame
import sys
import random
import math

pygame.init()

# -------------------- Settings --------------------
WIDTH, HEIGHT = 600, 800
FPS = 60

# Player settings
SHIP_SPEED = 7
SHIP_WIDTH, SHIP_HEIGHT = 40, 40
BULLET_SPEED = 10
FIRE_RATE = 15  # frames between shots

# Asteroid settings
ASTEROID_MIN_SIZE = 30
ASTEROID_MAX_SIZE = 80
ASTEROID_SPEED = 2
ASTEROID_SPAWN_RATE = 50  # frames between asteroid spawns

# Difficulty ramp
ASTEROID_SPEED_INCREASE = 0.0005
SPAWN_RATE_DECREASE = 0.001

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GREEN = (0, 255, 0)
DARK_GREEN = (0, 200, 0)
RED = (255, 0, 0)
DARK_RED = (150, 0, 0)
GREY = (100, 100, 100)

# Lives, health, etc.
STARTING_LIVES = 3
MAX_HEALTH = 100
MAX_SHIELD = 100
MAX_ENERGY = 100

# Fonts
FONT_NAME = pygame.font.match_font('arial')
def draw_text(surface, text, size, x, y, color=WHITE, center=True):
    font = pygame.font.Font(FONT_NAME, size)
    text_surface = font.render(text, True, color)
    text_rect = text_surface.get_rect()
    if center:
        text_rect.center = (x, y)
    else:
        text_rect.topleft = (x, y)
    surface.blit(text_surface, text_rect)

# Screen
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Asteroid Run")
clock = pygame.time.Clock()

# -------------------- Classes --------------------
class Starfield:
    """A class to represent a moving starfield background."""
    def __init__(self, num_stars=200):
        self.stars = []
        for _ in range(num_stars):
            x = random.randint(0, WIDTH)
            y = random.randint(0, HEIGHT)
            speed = random.random() * 0.5 + 0.1
            self.stars.append([x, y, speed])

    def update(self):
        for star in self.stars:
            # Move the star downward
            star[1] += star[2]
            # If off-screen, reset to top
            if star[1] > HEIGHT:
                star[1] = 0
                star[0] = random.randint(0, WIDTH)

    def draw(self, surface):
        for star in self.stars:
            surface.set_at((int(star[0]), int(star[1])), WHITE)

# Load your spaceship image
# Make sure you have "ship.png" in the same folder or adjust the path
ship_image = pygame.image.load("ship.png").convert_alpha()
ship_image = pygame.transform.scale(ship_image, (SHIP_WIDTH, SHIP_HEIGHT))

class Spaceship:
    def __init__(self):
        self.x = WIDTH // 2
        self.y = HEIGHT - SHIP_HEIGHT - 30
        self.width = SHIP_WIDTH
        self.height = SHIP_HEIGHT
        self.rect = pygame.Rect(self.x, self.y, self.width, self.height)
        self.can_fire = True
        self.fire_delay = 0
        self.lives = STARTING_LIVES
        self.health = MAX_HEALTH
        self.shield = MAX_SHIELD
        self.energy = MAX_ENERGY

    def move(self, dx):
        self.x += dx * SHIP_SPEED
        self.x = max(0, min(self.x, WIDTH - self.width))
        self.rect.x = self.x

    def shoot(self, bullets):
        if self.fire_delay == 0:
            bx = self.x + self.width // 2
            by = self.y
            bullets.append(Bullet(bx, by))
            self.fire_delay = FIRE_RATE

    def update(self):
        if self.fire_delay > 0:
            self.fire_delay -= 1

    def draw(self, surface):
        # Draw the spaceship image
        surface.blit(ship_image, (self.x, self.y))


class Bullet:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.speed = BULLET_SPEED
        self.rect = pygame.Rect(x-2, y-10, 4, 10)

    def update(self):
        self.y -= self.speed
        self.rect.y = self.y

    def draw(self, surface):
        pygame.draw.rect(surface, GREEN, self.rect)

    def off_screen(self):
        return self.y < 0

# Load the asteroid image
asteroid_image = pygame.image.load("asteroid.png").convert_alpha()

class Asteroid:
    def __init__(self, x, y, size, speed):
        self.x = x
        self.y = y
        self.size = size
        self.speed = speed
        self.rect = pygame.Rect(self.x - self.size//2, self.y - self.size//2, self.size, self.size)
        # Create a scaled image for each asteroid
        self.image = pygame.transform.scale(asteroid_image, (self.size, self.size))

    def update(self):
        self.y += self.speed
        self.rect.x = self.x - self.size//2
        self.rect.y = self.y - self.size//2

    def draw(self, surface):
        # Draw the asteroid image
        surface.blit(self.image, (self.rect.x, self.rect.y))

    def off_screen(self):
        return self.y > HEIGHT + self.size

def draw_status_bars(surface, ship):
    # Draw multiple bars at bottom-right corner:
    # For example: Health (Red), Shield (Green), Energy (Green) - side by side
    bar_width = 100
    bar_height = 10
    spacing = 5
    # Bottom-right corner starting position
    base_x = WIDTH - bar_width - 20
    base_y = HEIGHT - (bar_height * 3 + spacing * 2) - 20

    # Health bar (Red)
    # Background
    pygame.draw.rect(surface, DARK_RED, (base_x, base_y, bar_width, bar_height))
    # Foreground
    current_health_width = int((ship.health / MAX_HEALTH) * bar_width)
    pygame.draw.rect(surface, RED, (base_x, base_y, current_health_width, bar_height))

    # Shield bar (Green)
    shield_y = base_y + bar_height + spacing
    pygame.draw.rect(surface, (0, 100, 0), (base_x, shield_y, bar_width, bar_height))
    current_shield_width = int((ship.shield / MAX_SHIELD) * bar_width)
    pygame.draw.rect(surface, GREEN, (base_x, shield_y, current_shield_width, bar_height))

    # Energy bar (Greenish)
    energy_y = shield_y + bar_height + spacing
    pygame.draw.rect(surface, (0, 100, 100), (base_x, energy_y, bar_width, bar_height))
    current_energy_width = int((ship.energy / MAX_ENERGY) * bar_width)
    pygame.draw.rect(surface, (0, 200, 200), (base_x, energy_y, current_energy_width, bar_height))

def split_asteroid(asteroid, asteroids):
    # If we wanted to implement splitting, we could do so here.
    # For simplicity, let's not split. Just remove the asteroid.
    # But here's how you might do it if you wanted:
    # If asteroid is large, spawn two smaller ones
    # if asteroid.size > ASTEROID_MIN_SIZE * 1.5:
    #     new_size = asteroid.size // 2
    #     for i in range(2):
    #         ax = asteroid.x + random.randint(-new_size, new_size)
    #         ay = asteroid.y
    #         new = Asteroid(ax, ay, new_size, asteroid.speed * 1.2)
    #         asteroids.append(new)
    pass

def main_game():
    starfield = Starfield()
    ship = Spaceship()
    bullets = []
    asteroids = []
    frame_count = 0

    score = 0
    asteroid_timer = 0
    asteroid_speed = ASTEROID_SPEED
    spawn_rate = ASTEROID_SPAWN_RATE

    running = True
    while running:
        clock.tick(FPS)
        frame_count += 1

        # Increase difficulty
        asteroid_speed += ASTEROID_SPEED_INCREASE
        spawn_rate = max(10, spawn_rate - SPAWN_RATE_DECREASE)

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

        # Input
        keys = pygame.key.get_pressed()
        dx = 0
        if keys[pygame.K_LEFT] or keys[pygame.K_a]:
            dx = -1
        if keys[pygame.K_RIGHT] or keys[pygame.K_d]:
            dx = 1
        if keys[pygame.K_SPACE]:
            ship.shoot(bullets)

        # Update ship
        ship.move(dx)
        ship.update()

        # Update bullets
        for b in bullets[:]:
            b.update()
            if b.off_screen():
                bullets.remove(b)

        # Spawn asteroids
        asteroid_timer += 1
        if asteroid_timer > spawn_rate:
            asteroid_timer = 0
            # Random position, random size
            size = random.randint(ASTEROID_MIN_SIZE, ASTEROID_MAX_SIZE)
            ax = random.randint(size//2, WIDTH - size//2)
            ay = -size
            asteroids.append(Asteroid(ax, ay, size, asteroid_speed))

        # Update asteroids
        for a in asteroids[:]:
            a.update()
            if a.off_screen():
                asteroids.remove(a)
                # Maybe penalize player if asteroid escapes? (Not required)
                # ship.health = max(0, ship.health - 10)

        # Check collisions: bullet vs asteroid
        for a in asteroids[:]:
            for b in bullets[:]:
                if b.rect.colliderect(a.rect):
                    # Destroy asteroid
                    asteroids.remove(a)
                    bullets.remove(b)
                    score += 10
                    # Optionally split asteroid:
                    # split_asteroid(a, asteroids)
                    break

        # Check collisions: asteroid vs ship
        for a in asteroids[:]:
            if ship.rect.colliderect(a.rect):
                # Damage the ship
                damage = a.size // 2
                # If shield is available, use it first
                if ship.shield > 0:
                    absorbed = min(ship.shield, damage)
                    ship.shield -= absorbed
                    damage -= absorbed
                if damage > 0:
                    ship.health -= damage
                # Remove asteroid
                asteroids.remove(a)
                # Check if ship is destroyed
                if ship.health <= 0:
                    ship.lives -= 1
                    ship.health = MAX_HEALTH
                    ship.shield = MAX_SHIELD
                    ship.energy = MAX_ENERGY
                    if ship.lives <= 0:
                        running = False

        # Update starfield
        starfield.update()

        # Draw everything
        screen.fill(BLACK)
        starfield.draw(screen)

        # Draw ship
        ship.draw(screen)
        # Draw bullets
        for b in bullets:
            b.draw(screen)
        # Draw asteroids
        for a in asteroids:
            a.draw(screen)

        # Draw UI
        draw_text(screen, f"Lives: {ship.lives}", 24, 10, 10, color=GREEN, center=False)
        draw_text(screen, f"Score: {score}", 24, WIDTH - 10, 10, color=GREEN, center=False)
        draw_status_bars(screen, ship)

        pygame.display.flip()

    return score

def game_over_screen(score):
    screen.fill(BLACK)
    draw_text(screen, "GAME OVER", 64, WIDTH//2, HEIGHT//2 - 50, color=RED)
    draw_text(screen, f"Your Score: {score}", 36, WIDTH//2, HEIGHT//2 + 10, color=WHITE)
    draw_text(screen, "Press Space to Play Again or Q to Quit", 24, WIDTH//2, HEIGHT//2 + 70, color=WHITE)
    pygame.display.flip()

    waiting = True
    while waiting:
        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:
                    waiting = False
                elif event.key == pygame.K_q:
                    pygame.quit()
                    sys.exit()

def main():
    while True:
        score = main_game()
        game_over_screen(score)

if __name__ == "__main__":
    main()