import pygame
import random
import time

# Initialize Pygame
pygame.init()

# Set up the screen
SCREEN_WIDTH, SCREEN_HEIGHT = 800, 600
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Space Invaders")

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

# Game variables
player_size = 50
player_speed = 7
enemy_size = 50
enemy_speed = 2
enemy_spawn_delay = 50
bullet_speed = 7
bullet_cooldown = 20
score = 0

# Player
player = pygame.Rect(SCREEN_WIDTH // 2 - player_size // 2, SCREEN_HEIGHT - 50, player_size, player_size)

# Enemies
enemies = []

# Bullets
bullets = []

# Clock
clock = pygame.time.Clock()

# Fonts
font = pygame.font.SysFont(None, 36)

# Function to spawn enemies
def spawn_enemy():
    enemy_x = random.randint(0, SCREEN_WIDTH - enemy_size)
    enemy_y = random.randint(-SCREEN_HEIGHT, 0)
    enemy = pygame.Rect(enemy_x, enemy_y, enemy_size, enemy_size)
    enemies.append(enemy)

# Function to display text
def display_text(text, color, x, y):
    text_surface = font.render(text, True, color)
    screen.blit(text_surface, (x, y))

# Function to reset the game
def reset_game():
    global player, enemies, bullets, score
    player = pygame.Rect(SCREEN_WIDTH // 2 - player_size // 2, SCREEN_HEIGHT - 50, player_size, player_size)
    enemies = []
    bullets = []
    score = 0

# Main game loop
running = True
game_over = False
show_death_screen = False
death_screen_start_time = 0

while running:
    screen.fill(BLACK)

    if game_over:
        if not show_death_screen:
            display_text("Game Over!", RED, SCREEN_WIDTH // 2 - 100, SCREEN_HEIGHT // 2 - 50)
            display_text("Press 'R' to restart", WHITE, SCREEN_WIDTH // 2 - 130, SCREEN_HEIGHT // 2)
        else:
            display_text("You died!", RED, SCREEN_WIDTH // 2 - 80, SCREEN_HEIGHT // 2 - 50)
            display_text("Respawning in {:.1f} seconds".format(3 - (time.time() - death_screen_start_time)), WHITE, SCREEN_WIDTH // 2 - 150, SCREEN_HEIGHT // 2)
            if time.time() - death_screen_start_time >= 3:
                reset_game()
                game_over = False
                show_death_screen = False
    else:
        # Handle events
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE and len(bullets) < 3:
                    bullet = pygame.Rect(player.centerx - 2, player.top, 4, 10)
                    bullets.append(bullet)
                elif event.key == pygame.K_r:
                    reset_game()
                    game_over = False
                    show_death_screen = False

        # Player movement with WASD keys
        keys = pygame.key.get_pressed()
        if keys[pygame.K_a] and player.left > 0:
            player.x -= player_speed
        if keys[pygame.K_d] and player.right < SCREEN_WIDTH:
            player.x += player_speed
        if keys[pygame.K_w] and player.top > 0:
            player.y -= player_speed
        if keys[pygame.K_s] and player.bottom < SCREEN_HEIGHT:
            player.y += player_speed

        # Spawn enemies
        if random.randint(1, enemy_spawn_delay) == 1:
            spawn_enemy()

        # Update enemy positions
        for enemy in enemies:
            enemy.y += enemy_speed
            if enemy.y > SCREEN_HEIGHT:
                enemies.remove(enemy)
                score += 1

        # Draw player
        pygame.draw.rect(screen, WHITE, player)

        # Draw enemies
        for enemy in enemies:
            pygame.draw.rect(screen, GREEN, enemy)

        # Draw and update bullets
        for bullet in bullets:
            pygame.draw.rect(screen, RED, bullet)
            bullet.y -= bullet_speed
            if bullet.y < 0:
                bullets.remove(bullet)

        # Display score
        display_text("Score: " + str(score), WHITE, 10, 10)

        # Check for collision between player and enemies
        for enemy in enemies:
            if player.colliderect(enemy):
                game_over = True
                show_death_screen = True
                death_screen_start_time = time.time()

        # Check for collision between bullets and enemies
        for bullet in bullets:
            for enemy in enemies:
                if bullet.colliderect(enemy):
                    bullets.remove(bullet)
                    enemies.remove(enemy)
                    score += 10

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

pygame.quit()