import pygame
import random
import math

# --- Game Constants ---
WIDTH, HEIGHT = 800, 600
FPS = 60
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
SHIP_THRUST = 0.2
SHIP_MAX_SPEED = 5
BULLET_SPEED = 7

# --- Pygame Initialization ---
pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pygame Asteroids Clone")
clock = pygame.time.Clock()

# --- Utility Functions ---

def wrap_around(sprite):
    """Wraps sprite position when it leaves the screen edges."""
    if sprite.rect.left > WIDTH:
        sprite.rect.right = 0
    if sprite.rect.right < 0:
        sprite.rect.left = WIDTH
    if sprite.rect.top > HEIGHT:
        sprite.rect.bottom = 0
    if sprite.rect.bottom < 0:
        sprite.rect.top = HEIGHT

def distance(p1, p2):
    """Calculates the distance between two (x, y) points."""
    return math.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)

# --- Game Classes ---

class Ship(pygame.sprite.Sprite):
    def __init__(self):
        super().__init__()
        # The ship is a simple white triangle pointing up
        self.image_orig = pygame.Surface((30, 30), pygame.SRCALPHA)
        # Draw the triangle: (15, 0) is the nose
        pygame.draw.polygon(self.image_orig, WHITE, [(15, 0), (0, 30), (30, 30)])
        self.image = self.image_orig.copy()
        self.rect = self.image.get_rect(center=(WIDTH / 2, HEIGHT / 2))
        
        self.pos = pygame.math.Vector2(self.rect.center)
        self.vel = pygame.math.Vector2(0, 0)
        self.angle = 0 # 0 degrees is facing up
        self.rotation_speed = 4

    def update(self):
        # 1. Handle Rotation
        keys = pygame.key.get_pressed()
        if keys[pygame.K_LEFT]:
            self.angle += self.rotation_speed
        if keys[pygame.K_RIGHT]:
            self.angle -= self.rotation_speed
        
        # Keep angle within [0, 360)
        self.angle %= 360
        
        # Rotate the ship image
        self.image = pygame.transform.rotate(self.image_orig, self.angle)
        self.rect = self.image.get_rect(center=self.rect.center)
        
        # 2. Handle Thrust
        if keys[pygame.K_UP]:
            # Convert angle (degrees) to radians for trig functions
            angle_rad = math.radians(self.angle - 90) # Adjust angle for standard math orientation (0=right, 90=up)
            
            # Calculate thrust vector
            thrust_x = SHIP_THRUST * math.cos(angle_rad)
            thrust_y = SHIP_THRUST * math.sin(angle_rad)
            
            self.vel += pygame.math.Vector2(thrust_x, thrust_y)
        
        # 3. Apply Friction/Limit Speed (optional: implement max speed)
        if self.vel.length() > SHIP_MAX_SPEED:
            self.vel.scale_to_length(SHIP_MAX_SPEED)

        # 4. Apply Movement
        self.pos += self.vel
        self.rect.center = (int(self.pos.x), int(self.pos.y))

        # 5. Wrap Around Screen
        wrap_around(self)

    def shoot(self, all_sprites, bullets):
        # Calculate bullet starting position and velocity
        angle_rad = math.radians(self.angle - 90)
        
        # The nose of the ship is the spawn point. This is an approximation.
        spawn_x = self.pos.x + 15 * math.cos(angle_rad)
        spawn_y = self.pos.y + 15 * math.sin(angle_rad)
        
        bullet_vel_x = BULLET_SPEED * math.cos(angle_rad)
        bullet_vel_y = BULLET_SPEED * math.sin(angle_rad)
        
        bullet = Bullet(spawn_x, spawn_y, bullet_vel_x, bullet_vel_y)
        all_sprites.add(bullet)
        bullets.add(bullet)


class Asteroid(pygame.sprite.Sprite):
    def __init__(self, size, center):
        super().__init__()
        self.size = size # 3 (large), 2 (medium), 1 (small)
        self.radius = self.size * 15
        
        self.image = pygame.Surface((self.radius * 2, self.radius * 2), pygame.SRCALPHA)
        # Draw a white circle for simplicity; complex asteroids use polygon shapes
        pygame.draw.circle(self.image, WHITE, (self.radius, self.radius), self.radius, 1)
        
        self.rect = self.image.get_rect(center=center)
        
        # Random initial velocity
        angle = random.uniform(0, 2 * math.pi)
        speed = random.uniform(0.5, 2)
        self.vel = pygame.math.Vector2(speed * math.cos(angle), speed * math.sin(angle))

    def update(self):
        self.rect.centerx += self.vel.x
        self.rect.centery += self.vel.y
        wrap_around(self)

    def spawn_smaller(self):
        """Splits the asteroid into two smaller ones."""
        if self.size > 1:
            new_size = self.size - 1
            # Spawn two smaller asteroids moving in opposite directions
            return [
                Asteroid(new_size, self.rect.center),
                Asteroid(new_size, self.rect.center)
            ]
        return []

class Bullet(pygame.sprite.Sprite):
    def __init__(self, x, y, vel_x, vel_y):
        super().__init__()
        self.image = pygame.Surface((4, 4))
        self.image.fill(WHITE)
        self.rect = self.image.get_rect(center=(x, y))
        self.vel = pygame.math.Vector2(vel_x, vel_y)
        # Bullets have a limited lifespan
        self.lifespan = 90 # 90 frames (1.5 seconds at 60 FPS)

    def update(self):
        self.rect.centerx += self.vel.x
        self.rect.centery += self.vel.y
        
        # Check if bullet goes off screen (optional: wrap, but standard is destroy)
        if not (0 < self.rect.x < WIDTH and 0 < self.rect.y < HEIGHT):
            self.kill()
        
        # Decrease lifespan
        self.lifespan -= 1
        if self.lifespan <= 0:
            self.kill()

# --- Setup Game State ---
all_sprites = pygame.sprite.Group()
asteroids = pygame.sprite.Group()
bullets = pygame.sprite.Group()

player = Ship()
all_sprites.add(player)

# Spawn initial asteroids
for i in range(4):
    # Spawn large asteroids away from the center
    spawn_pos = (random.choice([random.randrange(0, 100), random.randrange(WIDTH - 100, WIDTH)]),
                 random.choice([random.randrange(0, 100), random.randrange(HEIGHT - 100, HEIGHT)]))
    asteroid = Asteroid(3, spawn_pos)
    all_sprites.add(asteroid)
    asteroids.add(asteroid)

# --- Game Loop ---
running = True
last_shot_time = 0
SHOT_DELAY = 250 # milliseconds

while running:
    # --- 1. Process Input/Events ---
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                # Add shot delay to prevent rapid-fire on key-hold
                now = pygame.time.get_ticks()
                if now - last_shot_time > SHOT_DELAY:
                    player.shoot(all_sprites, bullets)
                    last_shot_time = now

    # --- 2. Update Game State ---
    all_sprites.update()

    # --- 3. Check Collisions ---
    
    # Bullet hits Asteroid
    hits = pygame.sprite.groupcollide(bullets, asteroids, True, False)
    for bullet, hit_asteroids in hits.items():
        for asteroid in hit_asteroids:
            # Create smaller asteroids
            new_asteroids = asteroid.spawn_smaller()
            for new_ast in new_asteroids:
                all_sprites.add(new_ast)
                asteroids.add(new_ast)
            
            asteroid.kill() # Destroy the original asteroid

    # Ship hits Asteroid
    # 'pygame.sprite.collide_mask' is better for irregular shapes but 'pygame.sprite.collide_circle' is fine here
    if pygame.sprite.spritecollide(player, asteroids, True, pygame.sprite.collide_mask):
        # Game Over logic (reset or show message)
        print("Ship Destroyed! Game Over.")
        running = False # End the game for simplicity

    # --- 4. Render/Draw ---
    screen.fill(BLACK) # Clear the screen
    all_sprites.draw(screen) # Draw all sprites

    # Update the full display surface
    pygame.display.flip()

    # Wait for the next frame
    clock.tick(FPS)

pygame.quit()