import pygame
import sys
import math

# Initialize Pygame
pygame.init()

# Constants
WIDTH, HEIGHT = 800, 600
FPS = 60
PLAYER_SIZE = 50
DASH_SPEED = 20
NORMAL_SPEED = 5
DASH_COOLDOWN = 500  # milliseconds
BULLET_SPEED = 10
BULLET_SIZE = 10
ENEMY_SIZE = 50
ENEMY_SPEED = 2
RESPAWN_DELAY = 3000  # milliseconds (3 seconds respawn delay)

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

# Setup display
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Shooter with Dash, Cover, Enemies, and Respawn")

# Player class
class Player(pygame.sprite.Sprite):
    def __init__(self):
        super().__init__()
        self.image = pygame.Surface((PLAYER_SIZE, PLAYER_SIZE))
        self.image.fill(RED)
        self.rect = self.image.get_rect()
        # Move the player near the top center
        self.rect.center = (WIDTH // 2, 100)  # Centered horizontally, near the top of the screen
        
        self.speed = NORMAL_SPEED
        self.dash_speed = DASH_SPEED
        self.dashing = False
        self.last_dash_time = pygame.time.get_ticks()

    def update(self, walls):
        keys = pygame.key.get_pressed()
        
        # Movement (WASD)
        dx, dy = 0, 0
        if keys[pygame.K_a]:  # Left
            dx = -self.speed
        if keys[pygame.K_d]:  # Right
            dx = self.speed
        if keys[pygame.K_w]:  # Up
            dy = -self.speed
        if keys[pygame.K_s]:  # Down
            dy = self.speed

        # Check for collisions with walls and move around them (horizontal first, vertical second)
        if not self.collides_with_walls(dx, 0, walls):  # Check horizontal movement first
            self.rect.x += dx
        if not self.collides_with_walls(0, dy, walls):  # Check vertical movement second
            self.rect.y += dy
        
        # Dash ability
        current_time = pygame.time.get_ticks()
        if keys[pygame.K_SPACE] and not self.dashing and current_time - self.last_dash_time > DASH_COOLDOWN:
            self.dash()
            self.last_dash_time = current_time

    def dash(self):
        self.dashing = True
        self.speed = self.dash_speed
        pygame.time.set_timer(pygame.USEREVENT, 100)  # End dash after 100ms

    def stop_dash(self):
        self.dashing = False
        self.speed = NORMAL_SPEED
        pygame.time.set_timer(pygame.USEREVENT, 0)  # Stop dash timer

    def collides_with_walls(self, dx, dy, walls):
        """ Check for collision with walls after moving """
        new_rect = self.rect.move(dx, dy)
        for wall in walls:
            if new_rect.colliderect(wall.rect):
                return True
        return False

    def shoot(self):
        # Get mouse position
        mouse_x, mouse_y = pygame.mouse.get_pos()
        
        # Calculate the direction vector from the player to the mouse position
        direction_x = mouse_x - self.rect.centerx
        direction_y = mouse_y - self.rect.centery
        
        # Normalize the direction vector
        magnitude = math.hypot(direction_x, direction_y)
        if magnitude != 0:
            direction_x /= magnitude
            direction_y /= magnitude
        
        # Create a bullet and return it
        bullet = Bullet(self.rect.centerx, self.rect.centery, direction_x, direction_y)
        return bullet


# Bullet class
class Bullet(pygame.sprite.Sprite):
    def __init__(self, x, y, direction_x, direction_y):
        super().__init__()
        self.image = pygame.Surface((BULLET_SIZE, BULLET_SIZE))
        self.image.fill(YELLOW)
        self.rect = self.image.get_rect()
        self.rect.center = (x, y)
        
        # Set the movement direction
        self.direction_x = direction_x
        self.direction_y = direction_y

    def update(self, walls):
        # Move the bullet in the direction of the mouse
        self.rect.x += self.direction_x * BULLET_SPEED
        self.rect.y += self.direction_y * BULLET_SPEED
        
        # Check for collision with walls
        for wall in walls:
            if self.rect.colliderect(wall.rect):
                self.kill()  # Remove the bullet if it collides with a wall
                break
        
        # If the bullet goes off screen, kill it
        if self.rect.bottom < 0 or self.rect.top > HEIGHT or self.rect.right < 0 or self.rect.left > WIDTH:
            self.kill()


# Wall class (Cover)
class Wall(pygame.sprite.Sprite):
    def __init__(self, x, y, width, height):
        super().__init__()
        self.image = pygame.Surface((width, height))
        self.image.fill(GREEN)
        self.rect = self.image.get_rect()
        self.rect.topleft = (x, y)


# Enemy class
class Enemy(pygame.sprite.Sprite):
    def __init__(self):
        super().__init__()
        self.image = pygame.Surface((ENEMY_SIZE, ENEMY_SIZE))
        self.image.fill(BLUE)
        self.rect = self.image.get_rect()
        self.rect.center = (WIDTH // 2, HEIGHT - 50)  # Start at the bottom center of the screen
        
        self.speed = ENEMY_SPEED
        self.dead = False
        self.respawn_time = 0

    def update(self, player, walls):
        # If the enemy is dead, check if the respawn time has passed
        if self.dead and pygame.time.get_ticks() - self.respawn_time >= RESPAWN_DELAY:
            self.respawn()  # Respawn the enemy after the delay
        elif not self.dead:
            # Calculate the direction vector from the enemy to the player
            direction_x = player.rect.centerx - self.rect.centerx
            direction_y = player.rect.centery - self.rect.centery
            
            # Normalize the direction vector
            magnitude = math.hypot(direction_x, direction_y)
            if magnitude != 0:
                direction_x /= magnitude
                direction_y /= magnitude
            
            # Move the enemy towards the player (check for wall collisions)
            dx = direction_x * self.speed
            dy = direction_y * self.speed

            if not self.collides_with_walls(dx, 0, walls):  # Check horizontal movement first
                self.rect.x += dx
            if not self.collides_with_walls(0, dy, walls):  # Check vertical movement second
                self.rect.y += dy

    def die(self):
        self.dead = True
        self.respawn_time = pygame.time.get_ticks()
        self.kill()  # Remove the enemy when it dies

    def respawn(self):
        self.rect.center = (WIDTH // 2, HEIGHT - 50)  # Respawn at the bottom center
        self.dead = False
        self.image.fill(BLUE)  # Reset the color to blue


    def collides_with_walls(self, dx, dy, walls):
        """ Check for collision with walls after moving """
        new_rect = self.rect.move(dx, dy)
        for wall in walls:
            if new_rect.colliderect(wall.rect):
                return True
        return False


# Setup groups
all_sprites = pygame.sprite.Group()
player = Player()
all_sprites.add(player)

walls = pygame.sprite.Group()

# Create some walls (cover)
wall1 = Wall(200, 200, 100, 20)
wall2 = Wall(400, 300, 150, 20)
wall3 = Wall(600, 150, 100, 20)

walls.add(wall1, wall2, wall3)
all_sprites.add(wall1, wall2, wall3)

bullets = pygame.sprite.Group()

# Create enemies
enemy1 = Enemy()
enemy2 = Enemy()
all_sprites.add(enemy1, enemy2)

# Game loop
clock = pygame.time.Clock()
running = True
while running:
    screen.fill(BLACK)
    
    # Event handling
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        
        if event.type == pygame.USEREVENT:
            player.stop_dash()

        # Shooting bullets with the left mouse button
        if event.type == pygame.MOUSEBUTTONDOWN:
            if event.button == 1:  # Left click
                bullet = player.shoot()
                all_sprites.add(bullet)
                bullets.add(bullet)

    # Update
    player.update(walls)
    bullets.update(walls)  # Pass walls to the bullet update method for collision detection
    enemy1.update(player, walls)  # Move enemy 1 toward the player
    enemy2.update(player, walls)  # Move enemy 2 toward the player
    
    # Check for bullet collision with enemies
    for bullet in bullets:
        if pygame.sprite.spritecollide(bullet, all_sprites, False):
            for enemy in pygame.sprite.spritecollide(bullet, all_sprites, False):
                if isinstance(enemy, Enemy):  # Check if the colliding sprite is an enemy
                    enemy.die()  # Enemy dies
                    bullet.kill()  # Bullet is removed

    # Draw everything
    all_sprites.draw(screen)
    bullets.draw(screen)

    # Update the display
    pygame.display.flip()

    # Maintain FPS
    clock.tick(FPS)

pygame.quit()
sys.exit()