import pygame
import sys
import math
import random

# 初始化pygame
pygame.init()

# 游戏常量
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
FPS = 60

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

class Player:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.radius = 20
        self.speed = 5
        self.color = BLUE
        self.shoot_cooldown = 0
        self.max_health = 100
        self.health = self.max_health
    
    def update(self, keys):
        # 移动控制
        if keys[pygame.K_w] and self.y > self.radius:
            self.y -= self.speed
        if keys[pygame.K_s] and self.y < SCREEN_HEIGHT - self.radius:
            self.y += self.speed
        if keys[pygame.K_a] and self.x > self.radius:
            self.x -= self.speed
        if keys[pygame.K_d] and self.x < SCREEN_WIDTH - self.radius:
            self.x += self.speed
        
        # 射击冷却
        if self.shoot_cooldown > 0:
            self.shoot_cooldown -= 1
    
    def shoot(self, mouse_pos):
        if self.shoot_cooldown <= 0:
            self.shoot_cooldown = 15  # 射击间隔
            return Bullet(self.x, self.y, mouse_pos[0], mouse_pos[1])
        return None
    
    def draw(self, screen):
        pygame.draw.circle(screen, self.color, (int(self.x), int(self.y)), self.radius)
        # 绘制血条
        pygame.draw.rect(screen, RED, (self.x - 30, self.y - 35, 60, 8))
        pygame.draw.rect(screen, GREEN, (self.x - 30, self.y - 35, 
                                        int(60 * self.health / self.max_health), 8))

class Bullet:
    def __init__(self, x, y, target_x, target_y):
        self.x = x
        self.y = y
        self.radius = 5
        self.color = YELLOW
        self.speed = 10
        
        # 计算方向向量
        dx = target_x - x
        dy = target_y - y
        distance = max(1, math.sqrt(dx*dx + dy*dy))
        self.dx = dx / distance * self.speed
        self.dy = dy / distance * self.speed
    
    def update(self):
        self.x += self.dx
        self.y += self.dy
    
    def draw(self, screen):
        pygame.draw.circle(screen, self.color, (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, x, y):
        self.x = x
        self.y = y
        self.radius = 15
        self.color = RED
        self.speed = random.uniform(1, 3)
        self.max_health = 30
        self.health = self.max_health
        self.direction_x = random.choice([-1, 1])
        self.direction_y = random.choice([-1, 1])
    
    def update(self, player_x, player_y):
        # 简单的AI：向玩家移动
        dx = player_x - self.x
        dy = player_y - self.y
        distance = max(1, math.sqrt(dx*dx + dy*dy))
        self.x += dx / distance * self.speed
        self.y += dy / distance * self.speed
    
    def draw(self, screen):
        pygame.draw.circle(screen, self.color, (int(self.x), int(self.y)), self.radius)
        # 绘制血条
        pygame.draw.rect(screen, RED, (self.x - 20, self.y - 25, 40, 6))
        pygame.draw.rect(screen, GREEN, (self.x - 20, self.y - 25, 
                                        int(40 * self.health / self.max_health), 6))
    
    def is_off_screen(self):
        return (self.x < -50 or self.x > SCREEN_WIDTH + 50 or 
                self.y < -50 or self.y > SCREEN_HEIGHT + 50)

class Game:
    def __init__(self):
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
        pygame.display.set_caption("二维移动射击游戏")
        self.clock = pygame.time.Clock()
        
        self.player = Player(SCREEN_WIDTH // 2, SCREEN_HEIGHT // 2)
        self.bullets = []
        self.enemies = []
        self.score = 0
        self.enemy_spawn_timer = 0
        self.game_over = False
    
    def spawn_enemy(self):
        # 在屏幕边缘随机生成敌人
        side = random.randint(0, 3)
        if side == 0:  # 上边
            x = random.randint(0, SCREEN_WIDTH)
            y = -20
        elif side == 1:  # 右边
            x = SCREEN_WIDTH + 20
            y = random.randint(0, SCREEN_HEIGHT)
        elif side == 2:  # 下边
            x = random.randint(0, SCREEN_WIDTH)
            y = SCREEN_HEIGHT + 20
        else:  # 左边
            x = -20
            y = random.randint(0, SCREEN_HEIGHT)
        
        self.enemies.append(Enemy(x, y))
    
    def check_collisions(self):
        # 子弹和敌人的碰撞
        for bullet in self.bullets[:]:
            for enemy in self.enemies[:]:
                distance = math.sqrt((bullet.x - enemy.x)**2 + (bullet.y - enemy.y)**2)
                if distance < bullet.radius + enemy.radius:
                    enemy.health -= 25
                    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
    
    def check_player_hit(self):
        for enemy in self.enemies:
            distance = math.sqrt((self.player.x - enemy.x)**2 + (self.player.y - enemy.y)**2)
            if distance < self.player.radius + enemy.radius:
                self.player.health -= 1
                if self.player.health <= 0:
                    self.game_over = True
    
    def update(self):
        if self.game_over:
            return
        
        keys = pygame.key.get_pressed()
        self.player.update(keys)
        
        # 更新子弹
        for bullet in self.bullets[:]:
            bullet.update()
            if bullet.is_off_screen():
                self.bullets.remove(bullet)
        
        # 更新敌人
        for enemy in self.enemies[:]:
            enemy.update(self.player.x, self.player.y)
            if enemy.is_off_screen():
                self.enemies.remove(enemy)
        
        # 生成敌人
        self.enemy_spawn_timer += 1
        if self.enemy_spawn_timer >= 60:  # 每秒生成一个敌人
            self.spawn_enemy()
            self.enemy_spawn_timer = 0
        
        # 检查碰撞
        self.check_collisions()
        self.check_player_hit()
    
    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)
        
        # 绘制分数
        font = pygame.font.Font(None, 36)
        score_text = font.render(f"Score: {self.score}", True, WHITE)
        self.screen.blit(score_text, (10, 10))
        
        # 绘制健康值
        health_text = font.render(f"Health: {self.player.health}", True, WHITE)
        self.screen.blit(health_text, (10, 50))
        
        # 游戏结束显示
        if self.game_over:
            game_over_font = pygame.font.Font(None, 72)
            game_over_text = game_over_font.render("GAME OVER", True, RED)
            restart_text = font.render("Press R to Restart", True, WHITE)
            
            self.screen.blit(game_over_text, 
                           (SCREEN_WIDTH//2 - game_over_text.get_width()//2, 
                            SCREEN_HEIGHT//2 - 50))
            self.screen.blit(restart_text, 
                           (SCREEN_WIDTH//2 - restart_text.get_width()//2, 
                            SCREEN_HEIGHT//2 + 30))
        
        pygame.display.flip()
    
    def handle_events(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
            
            if event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 1 and not self.game_over:  # 左键点击
                    mouse_pos = pygame.mouse.get_pos()
                    bullet = self.player.shoot(mouse_pos)
                    if bullet:
                        self.bullets.append(bullet)
            
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_r and self.game_over:
                    # 重新开始游戏
                    self.__init__()
        
        return True
    
    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__":
    # 检查是否安装了pygame
    try:
        game = Game()
        game.run()
    except ImportError:
        print("请先安装pygame: pip install pygame")