import pygame
import math
import random
import sys

# 初始化Pygame
pygame.init()

# 游戏设置
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
FPS = 60

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
YELLOW = (255, 255, 0)
PURPLE = (128, 0, 128)

class Player:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.width = 32
        self.height = 32
        self.speed = 5
        self.health = 100
        self.max_health = 100
        self.color = BLUE
        
    def move(self, keys):
        if keys[pygame.K_LEFT] or keys[pygame.K_a]:
            self.x -= self.speed
        if keys[pygame.K_RIGHT] or keys[pygame.K_d]:
            self.x += self.speed
        if keys[pygame.K_UP] or keys[pygame.K_w]:
            self.y -= self.speed
        if keys[pygame.K_DOWN] or keys[pygame.K_s]:
            self.y += self.speed
            
        # 边界限制
        self.x = max(0, min(self.x, SCREEN_WIDTH - self.width))
        self.y = max(0, min(self.y, SCREEN_HEIGHT - self.height))
    
    def draw(self, screen):
        # 绘制玩家（三角形指向鼠标方向）
        mouse_x, mouse_y = pygame.mouse.get_pos()
        angle = math.atan2(mouse_y - (self.y + self.height//2), 
                          mouse_x - (self.x + self.width//2))
        
        # 三角形的三个顶点
        center_x = self.x + self.width//2
        center_y = self.y + self.height//2
        
        tip_x = center_x + 20 * math.cos(angle)
        tip_y = center_y + 20 * math.sin(angle)
        
        left_angle = angle + math.radians(140)
        right_angle = angle - math.radians(140)
        
        left_x = center_x + 15 * math.cos(left_angle)
        left_y = center_y + 15 * math.sin(left_angle)
        right_x = center_x + 15 * math.cos(right_angle)
        right_y = center_y + 15 * math.sin(right_angle)
        
        pygame.draw.polygon(screen, self.color, [(tip_x, tip_y), (left_x, left_y), (right_x, right_y)])
        
        # 绘制血条
        bar_width = 40
        bar_height = 6
        bar_x = self.x + self.width//2 - bar_width//2
        bar_y = self.y - 10
        health_percent = self.health / self.max_health
        pygame.draw.rect(screen, RED, (bar_x, bar_y, bar_width, bar_height))
        pygame.draw.rect(screen, GREEN, (bar_x, bar_y, bar_width * health_percent, bar_height))
    
    def get_rect(self):
        return pygame.Rect(self.x, self.y, self.width, self.height)

class Bullet:
    def __init__(self, x, y, angle):
        self.x = x
        self.y = y
        self.angle = angle
        self.speed = 10
        self.radius = 5
        self.damage = 25
        
    def update(self):
        self.x += self.speed * math.cos(self.angle)
        self.y += self.speed * math.sin(self.angle)
        
    def draw(self, screen):
        pygame.draw.circle(screen, YELLOW, (int(self.x), int(self.y)), self.radius)
        
    def get_rect(self):
        return pygame.Rect(self.x - self.radius, self.y - self.radius, 
                          self.radius * 2, self.radius * 2)

class Enemy:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.width = 30
        self.height = 30
        self.speed = 2
        self.health = 50
        self.max_health = 50
        self.color = RED
        self.damage = 10
        
    def update(self, player_x, player_y):
        # 向玩家移动
        dx = player_x - self.x
        dy = player_y - self.y
        distance = math.sqrt(dx**2 + dy**2)
        if distance > 0:
            self.x += (dx / distance) * self.speed
            self.y += (dy / distance) * self.speed
    
    def draw(self, screen):
        # 绘制敌人为正方形
        pygame.draw.rect(screen, self.color, (self.x, self.y, self.width, self.height))
        
        # 绘制血条
        bar_width = self.width
        bar_height = 4
        bar_x = self.x
        bar_y = self.y - 8
        health_percent = self.health / self.max_health
        pygame.draw.rect(screen, RED, (bar_x, bar_y, bar_width, bar_height))
        pygame.draw.rect(screen, GREEN, (bar_x, bar_y, bar_width * health_percent, bar_height))
    
    def get_rect(self):
        return pygame.Rect(self.x, self.y, self.width, self.height)

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.reset_game()
        
    def reset_game(self):
        self.player = Player(SCREEN_WIDTH//2, SCREEN_HEIGHT//2)
        self.bullets = []
        self.enemies = []
        self.score = 0
        self.game_over = False
        self.shoot_cooldown = 0
        self.enemy_spawn_timer = 0
        
        # 初始生成几个敌人
        for _ in range(3):
            self.spawn_enemy()
    
    def spawn_enemy(self):
        side = random.choice(['top', 'bottom', 'left', 'right'])
        if side == 'top':
            x = random.randint(0, SCREEN_WIDTH - 30)
            y = -30
        elif side == 'bottom':
            x = random.randint(0, SCREEN_WIDTH - 30)
            y = SCREEN_HEIGHT
        elif side == 'left':
            x = -30
            y = random.randint(0, SCREEN_HEIGHT - 30)
        else:
            x = SCREEN_WIDTH
            y = random.randint(0, SCREEN_HEIGHT - 30)
        
        self.enemies.append(Enemy(x, y))
    
    def handle_input(self):
        keys = pygame.key.get_pressed()
        self.player.move(keys)
        
        # 射击（鼠标左键或空格）
        if self.shoot_cooldown <= 0:
            shoot = False
            if pygame.mouse.get_pressed()[0] or keys[pygame.K_SPACE]:
                shoot = True
            
            if shoot:
                mouse_x, mouse_y = pygame.mouse.get_pos()
                angle = math.atan2(mouse_y - (self.player.y + self.player.height//2),
                                  mouse_x - (self.player.x + self.player.width//2))
                bullet = Bullet(self.player.x + self.player.width//2, 
                               self.player.y + self.player.height//2, angle)
                self.bullets.append(bullet)
                self.shoot_cooldown = 10
        
        if self.shoot_cooldown > 0:
            self.shoot_cooldown -= 1
    
    def update(self):
        if self.game_over:
            return
        
        # 更新子弹
        for bullet in self.bullets[:]:
            bullet.update()
            # 超出边界删除
            if (bullet.x < -50 or bullet.x > SCREEN_WIDTH + 50 or
                bullet.y < -50 or bullet.y > SCREEN_HEIGHT + 50):
                self.bullets.remove(bullet)
        
        # 更新敌人
        for enemy in self.enemies[:]:
            enemy.update(self.player.x + self.player.width//2, 
                        self.player.y + self.player.height//2)
            
            # 检查敌人与玩家的碰撞
            if enemy.get_rect().colliderect(self.player.get_rect()):
                self.player.health -= enemy.damage
                self.enemies.remove(enemy)
                if self.player.health <= 0:
                    self.game_over = True
                    self.player.health = 0
        
        # 检查子弹与敌人的碰撞
        for bullet in self.bullets[:]:
            for enemy in self.enemies[:]:
                if bullet.get_rect().colliderect(enemy.get_rect()):
                    enemy.health -= bullet.damage
                    if bullet in self.bullets:
                        self.bullets.remove(bullet)
                    if enemy.health <= 0:
                        if enemy in self.enemies:
                            self.enemies.remove(enemy)
                            self.score += 10
                    break
        
        # 生成敌人
        self.enemy_spawn_timer += 1
        if self.enemy_spawn_timer >= 60:  # 每秒生成一个
            self.spawn_enemy()
            self.enemy_spawn_timer = 0
    
    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))
        
        health_text = self.font.render(f"生命: {self.player.health}", True, WHITE)
        self.screen.blit(health_text, (10, 50))
        
        # 游戏结束画面
        if self.game_over:
            overlay = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT))
            overlay.set_alpha(128)
            overlay.fill(BLACK)
            self.screen.blit(overlay, (0, 0))
            
            game_over_text = self.font.render("游戏结束!", True, RED)
            text_rect = game_over_text.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//2 - 50))
            self.screen.blit(game_over_text, text_rect)
            
            final_score_text = self.font.render(f"最终分数: {self.score}", True, WHITE)
            score_rect = final_score_text.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//2))
            self.screen.blit(final_score_text, score_rect)
            
            restart_text = self.small_font.render("按 R 键重新开始", True, WHITE)
            restart_rect = restart_text.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//2 + 50))
            self.screen.blit(restart_text, restart_rect)
            
            quit_text = self.small_font.render("按 ESC 退出游戏", True, WHITE)
            quit_rect = quit_text.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//2 + 80))
            self.screen.blit(quit_text, quit_rect)
        
        # 显示控制说明
        controls = [
            "WASD / 方向键: 移动",
            "鼠标 / 空格: 射击",
            "R: 重新开始",
            "ESC: 退出"
        ]
        y_offset = SCREEN_HEIGHT - 100
        for control in controls:
            control_text = self.small_font.render(control, True, WHITE)
            self.screen.blit(control_text, (10, y_offset))
            y_offset += 25
        
        pygame.display.flip()
    
    def run(self):
        running = True
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                elif event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_ESCAPE:
                        running = False
                    elif event.key == pygame.K_r and self.game_over:
                        self.reset_game()
            
            self.handle_input()
            self.update()
            self.draw()
            self.clock.tick(FPS)
        
        pygame.quit()
        sys.exit()

if __name__ == "__main__":
    game = Game()
    game.run()