import pygame
import sys
import random
import math

# 初始化pygame
pygame.init()

# 游戏常量
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
FPS = 60

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 100, 255)
YELLOW = (255, 255, 0)


class Player:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.width = 40
        self.height = 30
        self.speed = 5
        self.health = 100
        self.max_health = 100

    def draw(self, screen):
        # 绘制玩家飞船（三角形）
        points = [
            (self.x, self.y - self.height // 2),  # 顶部
            (self.x - self.width // 2, self.y + self.height // 2),  # 左下
            (self.x + self.width // 2, self.y + self.height // 2)   # 右下
        ]
        pygame.draw.polygon(screen, BLUE, points)

        # 绘制血条
        bar_width = 50
        bar_height = 5
        pygame.draw.rect(screen, RED, (self.x - bar_width //
                         2, self.y - 25, bar_width, bar_height))
        health_width = int(bar_width * (self.health / self.max_health))
        pygame.draw.rect(screen, GREEN, (self.x - bar_width //
                         2, self.y - 25, health_width, bar_height))

    def move(self, dx, dy):
        self.x = max(self.width // 2, min(SCREEN_WIDTH -
                     self.width // 2, self.x + dx))
        self.y = max(self.height // 2, min(SCREEN_HEIGHT -
                     self.height // 2, self.y + dy))

    def get_rect(self):
        return pygame.Rect(self.x - self.width // 2, self.y - self.height // 2,
                           self.width, self.height)


class Bullet:
    def __init__(self, x, y, angle=0):
        self.x = x
        self.y = y
        self.radius = 3
        self.speed = 8
        self.angle = angle
        self.dx = math.sin(math.radians(angle)) * self.speed
        self.dy = -math.cos(math.radians(angle)) * self.speed

    def update(self):
        self.x += self.dx
        self.y += self.dy

    def draw(self, screen):
        pygame.draw.circle(
            screen, YELLOW, (int(self.x), int(self.y)), self.radius)

    def is_off_screen(self):
        return (self.x < 0 or self.x > SCREEN_WIDTH or
                self.y < 0 or self.y > SCREEN_HEIGHT)


class Enemy:
    def __init__(self):
        self.radius = 20
        # 从屏幕上方随机位置生成
        self.x = random.randint(self.radius, SCREEN_WIDTH - self.radius)
        self.y = -self.radius
        self.speed = random.uniform(1, 3)
        self.health = 20

    def update(self):
        self.y += self.speed

    def draw(self, screen):
        pygame.draw.circle(
            screen, RED, (int(self.x), int(self.y)), self.radius)
        # 绘制敌人血条
        if self.health < 20:
            bar_width = 40
            bar_height = 4
            pygame.draw.rect(screen, RED, (self.x - bar_width //
                             2, self.y - 30, bar_width, bar_height))
            health_width = int(bar_width * (self.health / 20))
            pygame.draw.rect(screen, GREEN, (self.x - bar_width //
                             2, self.y - 30, health_width, bar_height))

    def is_off_screen(self):
        return self.y > SCREEN_HEIGHT + self.radius

    def get_rect(self):
        return pygame.Rect(self.x - self.radius, self.y - self.radius,
                           self.radius * 2, self.radius * 2)


class Game:
    def __init__(self):
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
        pygame.display.set_caption("2D 射击游戏")
        self.clock = pygame.time.Clock()
        self.font = pygame.font.Font(None, 36)
        self.small_font = pygame.font.Font(None, 24)

        # 游戏对象
        self.player = Player(SCREEN_WIDTH // 2, SCREEN_HEIGHT - 50)
        self.bullets = []
        self.enemies = []

        # 游戏状态
        self.score = 0
        self.game_over = False
        self.enemy_spawn_timer = 0
        self.shoot_cooldown = 0

    def handle_events(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE and not self.game_over:
                    self.shoot()
                elif event.key == pygame.K_r and self.game_over:
                    self.restart_game()
        return True

    def shoot(self):
        if self.shoot_cooldown <= 0:
            # 从玩家位置发射子弹
            bullet = Bullet(self.player.x, self.player.y -
                            self.player.height // 2)
            self.bullets.append(bullet)
            self.shoot_cooldown = 10  # 射击冷却时间

    def update(self):
        if self.game_over:
            return

        # 处理按键输入
        keys = pygame.key.get_pressed()
        dx, dy = 0, 0
        if keys[pygame.K_LEFT] or keys[pygame.K_a]:
            dx = -self.player.speed
        if keys[pygame.K_RIGHT] or keys[pygame.K_d]:
            dx = self.player.speed
        if keys[pygame.K_UP] or keys[pygame.K_w]:
            dy = -self.player.speed
        if keys[pygame.K_DOWN] or keys[pygame.K_s]:
            dy = self.player.speed

        self.player.move(dx, dy)

        # 更新射击冷却
        if self.shoot_cooldown > 0:
            self.shoot_cooldown -= 1

        # 生成敌人
        self.enemy_spawn_timer += 1
        if self.enemy_spawn_timer >= 60:  # 每秒生成一个敌人
            self.enemies.append(Enemy())
            self.enemy_spawn_timer = 0

        # 更新子弹
        for bullet in self.bullets[:]:
            bullet.update()
            if bullet.is_off_screen():
                self.bullets.remove(bullet)

        # 更新敌人
        for enemy in self.enemies[:]:
            enemy.update()
            if enemy.is_off_screen():
                self.enemies.remove(enemy)

        # 碰撞检测 - 子弹击中敌人
        for bullet in self.bullets[:]:
            for enemy in self.enemies[:]:
                if self.check_collision(bullet, enemy):
                    if bullet in self.bullets:
                        self.bullets.remove(bullet)
                    enemy.health -= 10
                    if enemy.health <= 0:
                        self.enemies.remove(enemy)
                        self.score += 10

        # 碰撞检测 - 玩家与敌人
        player_rect = self.player.get_rect()
        for enemy in self.enemies[:]:
            if player_rect.colliderect(enemy.get_rect()):
                self.player.health -= 1
                if enemy in self.enemies:
                    self.enemies.remove(enemy)
                if self.player.health <= 0:
                    self.game_over = True

    def check_collision(self, bullet, enemy):
        distance = math.sqrt((bullet.x - enemy.x) ** 2 +
                             (bullet.y - enemy.y) ** 2)
        return distance < bullet.radius + enemy.radius

    def draw(self):
        self.screen.fill(BLACK)

        # 绘制游戏对象
        self.player.draw(self.screen)

        for bullet in self.bullets:
            bullet.draw(self.screen)

        for enemy in self.enemies:
            enemy.draw(self.screen)

        # 绘制UI
        score_text = self.font.render(f"分数: {self.score}", True, WHITE)
        self.screen.blit(score_text, (10, 10))

        # 绘制操作说明
        if not self.game_over:
            controls_text = self.small_font.render(
                "WASD/方向键移动，空格键射击", True, WHITE)
            self.screen.blit(controls_text, (SCREEN_WIDTH //
                             2 - controls_text.get_width() // 2, 10))

        # 游戏结束画面
        if self.game_over:
            game_over_text = self.font.render("游戏结束!", True, RED)
            score_text = self.font.render(f"最终分数: {self.score}", True, WHITE)
            restart_text = self.font.render("按 R 重新开始", True, WHITE)

            self.screen.blit(game_over_text, (SCREEN_WIDTH // 2 - game_over_text.get_width() // 2,
                                              SCREEN_HEIGHT // 2 - 60))
            self.screen.blit(score_text, (SCREEN_WIDTH // 2 - score_text.get_width() // 2,
                                          SCREEN_HEIGHT // 2))
            self.screen.blit(restart_text, (SCREEN_WIDTH // 2 - restart_text.get_width() // 2,
                                            SCREEN_HEIGHT // 2 + 60))

        pygame.display.flip()

    def restart_game(self):
        self.player = Player(SCREEN_WIDTH // 2, SCREEN_HEIGHT - 50)
        self.bullets = []
        self.enemies = []
        self.score = 0
        self.game_over = False
        self.enemy_spawn_timer = 0
        self.shoot_cooldown = 0

    def run(self):
        running = True
        while running:
            running = self.handle_events()
            self.update()
            self.draw()
            self.clock.tick(FPS)

        pygame.quit()
        sys.exit()


if __name__ == "__main__":
    game = Game()
    game.run()
