import pygame
import random
import math
import sys

pygame.init()

# 游戏窗口设置
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("二维移动射击 - 敌人可攻击")
clock = pygame.time.Clock()
FPS = 60

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 50, 50)
GREEN = (50, 255, 50)
BLUE = (50, 150, 255)
YELLOW = (255, 255, 50)
PURPLE = (200, 50, 255)
ORANGE = (255, 150, 50)
DARK_BLUE = (0, 0, 100)
PARTICLE_COLORS = [(255, 200, 50), (255, 100, 50), (255, 255, 200)]


class Player:
    def __init__(self):
        self.width = 40
        self.height = 40
        self.x = WIDTH // 2
        self.y = HEIGHT - 100
        self.speed = 5
        self.color = BLUE
        self.health = 100
        self.max_health = 100
        self.shoot_cooldown = 0
        self.score = 0
        self.invincible = 0
        self.weapon_level = 1
        self.bullet_damage = 20

    def move(self, keys):
        if keys[pygame.K_LEFT] and self.x > 0:
            self.x -= self.speed
        if keys[pygame.K_RIGHT] and self.x < WIDTH - self.width:
            self.x += self.speed
        if keys[pygame.K_UP] and self.y > 0:
            self.y -= self.speed
        if keys[pygame.K_DOWN] and self.y < HEIGHT - self.height:
            self.y += self.speed

    def shoot(self, bullets):
        if self.shoot_cooldown <= 0:
            if self.weapon_level == 1:
                bullets.append(PlayerBullet(self.x + self.width //
                               2 - 2, self.y, -10, self.bullet_damage))
            elif self.weapon_level == 2:
                bullets.append(PlayerBullet(self.x + self.width //
                               2 - 8, self.y, -10, self.bullet_damage))
                bullets.append(PlayerBullet(self.x + self.width //
                               2 + 4, self.y, -10, self.bullet_damage))
            elif self.weapon_level >= 3:
                bullets.append(PlayerBullet(self.x + self.width //
                               2 - 12, self.y, -10, self.bullet_damage))
                bullets.append(PlayerBullet(self.x + self.width //
                               2 - 4, self.y, -10, self.bullet_damage))
                bullets.append(PlayerBullet(self.x + self.width //
                               2 + 4, self.y, -10, self.bullet_damage))
                bullets.append(PlayerBullet(self.x + self.width //
                               2 + 12, self.y, -10, self.bullet_damage))

            self.shoot_cooldown = 5

    def update(self):
        if self.shoot_cooldown > 0:
            self.shoot_cooldown -= 1
        if self.invincible > 0:
            self.invincible -= 1

    def draw(self, surface):
        # 绘制玩家飞船
        if self.invincible > 0 and self.invincible % 4 < 2:
            pygame.draw.rect(
                surface, YELLOW, (self.x, self.y, self.width, self.height))
        else:
            pygame.draw.rect(surface, self.color,
                             (self.x, self.y, self.width, self.height))

        # 飞船细节
        pygame.draw.polygon(surface, RED, [
            (self.x + self.width//2, self.y - 10),
            (self.x + 10, self.y + 5),
            (self.x + self.width - 10, self.y + 5)
        ])

    def get_rect(self):
        return pygame.Rect(self.x, self.y, self.width, self.height)

    def take_damage(self, amount):
        if self.invincible <= 0:
            self.health -= amount
            self.invincible = 30
            return True
        return False


class PlayerBullet:
    def __init__(self, x, y, speed, damage=20):
        self.x = x
        self.y = y
        self.width = 4
        self.height = 12
        self.speed = speed
        self.color = GREEN
        self.damage = damage

    def update(self):
        self.y += self.speed

    def draw(self, surface):
        pygame.draw.rect(surface, self.color,
                         (self.x, self.y, self.width, self.height))
        pygame.draw.rect(surface, YELLOW, (self.x-1,
                         self.y-2, self.width+2, 2))

    def is_off_screen(self):
        return self.y < 0 or self.y > HEIGHT


class EnemyBullet:
    def __init__(self, x, y, target_x, target_y):
        self.x = x
        self.y = y
        self.width = 6
        self.height = 6
        self.color = ORANGE
        self.speed = 3
        self.damage = 10

        # 计算子弹方向（追踪玩家）
        dx = target_x - x
        dy = target_y - y
        distance = max(math.sqrt(dx*dx + dy*dy), 0.1)
        self.vx = (dx / distance) * self.speed
        self.vy = (dy / distance) * self.speed

    def update(self):
        self.x += self.vx
        self.y += self.vy

    def draw(self, surface):
        pygame.draw.circle(surface, self.color,
                           (int(self.x), int(self.y)), self.width)
        pygame.draw.circle(surface, YELLOW, (int(
            self.x), int(self.y)), self.width-2)

    def is_off_screen(self):
        return (self.x < 0 or self.x > WIDTH or
                self.y < 0 or self.y > HEIGHT)


class Enemy:
    def __init__(self, enemy_type="normal"):
        self.enemy_type = enemy_type

        if enemy_type == "normal":
            self.width = 30
            self.height = 30
            self.speed = random.uniform(1.0, 3.0)
            self.color = RED
            self.health = 30
            self.max_health = 30
            self.shoot_cooldown = random.randint(60, 120)  # 1-2秒射击一次
            self.bullet_damage = 10

        elif enemy_type == "fast":
            self.width = 20
            self.height = 20
            self.speed = random.uniform(3.0, 5.0)
            self.color = (255, 100, 100)
            self.health = 20
            self.max_health = 20
            self.shoot_cooldown = random.randint(30, 60)  # 0.5-1秒射击一次
            self.bullet_damage = 8

        elif enemy_type == "tank":
            self.width = 50
            self.height = 50
            self.speed = random.uniform(0.5, 1.5)
            self.color = (100, 100, 255)
            self.health = 100
            self.max_health = 100
            self.shoot_cooldown = random.randint(90, 150)  # 1.5-2.5秒射击一次
            self.bullet_damage = 15

        self.x = random.randint(0, WIDTH - self.width)
        self.y = random.randint(-100, -40)

    def update(self, player_x, player_y, enemy_bullets):
        # 移动
        self.y += self.speed
        self.x += random.uniform(-0.5, 0.5)
        self.x = max(0, min(WIDTH - self.width, self.x))

        # 射击冷却
        if self.shoot_cooldown > 0:
            self.shoot_cooldown -= 1

        # 射击玩家
        if self.shoot_cooldown <= 0 and random.random() < 0.02:  # 2%几率射击
            target_x = player_x + 20
            target_y = player_y + 20
            bullet_x = self.x + self.width // 2
            bullet_y = self.y + self.height

            enemy_bullets.append(EnemyBullet(
                bullet_x, bullet_y, target_x, target_y))
            self.shoot_cooldown = random.randint(60, 120)  # 重置冷却

    def draw(self, surface):
        # 绘制敌人
        pygame.draw.rect(surface, self.color,
                         (self.x, self.y, self.width, self.height))

        # 敌人细节
        if self.enemy_type == "normal":
            pygame.draw.rect(surface, (255, 100, 100),
                             (self.x+5, self.y+5, self.width-10, self.height-10))
        elif self.enemy_type == "fast":
            pygame.draw.rect(surface, (255, 150, 150),
                             (self.x+3, self.y+3, self.width-6, self.height-6))
        elif self.enemy_type == "tank":
            pygame.draw.rect(surface, (150, 150, 255),
                             (self.x+8, self.y+8, self.width-16, self.height-16))

        # 敌人生命条
        health_width = self.width
        health_height = 4
        health_ratio = self.health / self.max_health

        pygame.draw.rect(surface, (50, 50, 50),
                         (self.x, self.y-8, health_width, health_height))
        health_color = GREEN if health_ratio > 0.5 else YELLOW if health_ratio > 0.2 else RED
        pygame.draw.rect(surface, health_color, (self.x, self.y-8,
                         int(health_width * health_ratio), health_height))

    def is_off_screen(self):
        return self.y > HEIGHT

    def get_rect(self):
        return pygame.Rect(self.x, self.y, self.width, self.height)


class Particle:
    def __init__(self, x, y, color=None):
        self.x = x
        self.y = y
        self.vx = random.uniform(-3, 3)
        self.vy = random.uniform(-3, 3)
        self.color = color if color else random.choice(PARTICLE_COLORS)
        self.life = random.randint(20, 40)
        self.size = random.randint(2, 6)

    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.life -= 1
        self.size = max(0, self.size - 0.1)
        self.vx *= 0.95
        self.vy *= 0.95

    def draw(self, surface):
        if self.life > 0 and self.size > 0:
            pygame.draw.circle(surface, self.color,
                               (int(self.x), int(self.y)), int(self.size))

    def is_dead(self):
        return self.life <= 0 or self.size <= 0


class PowerUp:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.width = 20
        self.height = 20
        self.speed = 2
        self.type = random.choice(["health", "weapon", "shield"])
        self.color = GREEN if self.type == "health" else PURPLE if self.type == "weapon" else BLUE
        self.life = 300

    def update(self):
        self.y += self.speed
        self.life -= 1
        self.x += math.sin(pygame.time.get_ticks() * 0.01) * 0.5

    def draw(self, surface):
        if self.life > 0:
            if self.type == "health":
                pygame.draw.rect(surface, self.color,
                                 (self.x, self.y, self.width, self.height))
                pygame.draw.polygon(surface, WHITE, [
                    (self.x + self.width//2, self.y + 5),
                    (self.x + 5, self.y + self.height - 5),
                    (self.x + self.width - 5, self.y + self.height - 5)
                ])
            elif self.type == "weapon":
                pygame.draw.rect(surface, self.color,
                                 (self.x, self.y, self.width, self.height))
                pygame.draw.rect(surface, YELLOW, (self.x + 3,
                                 self.y + 5, self.width - 6, 3))
                pygame.draw.rect(surface, YELLOW, (self.x + 3,
                                 self.y + 12, self.width - 6, 3))
            else:  # shield
                pygame.draw.rect(surface, self.color,
                                 (self.x, self.y, self.width, self.height))
                pygame.draw.circle(
                    surface, WHITE, (self.x + self.width//2, self.y + self.height//2), 8, 2)

    def is_off_screen(self):
        return self.y > HEIGHT or self.life <= 0

    def get_rect(self):
        return pygame.Rect(self.x, self.y, self.width, self.height)


def draw_health_bar(surface, player, x, y, width, height):
    # 背景
    pygame.draw.rect(surface, (50, 50, 50), (x, y, width, height))

    # 生命值
    health_ratio = player.health / player.max_health
    health_width = int(width * health_ratio)
    health_color = GREEN if health_ratio > 0.5 else YELLOW if health_ratio > 0.2 else RED
    pygame.draw.rect(surface, health_color, (x, y, health_width, height))

    # 边框
    pygame.draw.rect(surface, WHITE, (x, y, width, height), 1)

    # 生命值文字
    font = pygame.font.SysFont("simhei", 16)
    health_text = font.render(
        f"{player.health}/{player.max_health}", True, WHITE)
    surface.blit(health_text, (x + width//2 - health_text.get_width() //
                 2, y + height//2 - health_text.get_height()//2))


def show_game_over(surface, score, font, big_font):
    overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
    overlay.fill((0, 0, 0, 180))
    surface.blit(overlay, (0, 0))

    game_over = big_font.render("游戏结束", True, RED)
    score_text = font.render(f"最终得分: {score}", True, YELLOW)
    restart_text = font.render("按R键重新开始，ESC退出", True, WHITE)

    surface.blit(
        game_over, (WIDTH//2 - game_over.get_width()//2, HEIGHT//2 - 60))
    surface.blit(score_text, (WIDTH//2 - score_text.get_width()//2, HEIGHT//2))
    surface.blit(restart_text, (WIDTH//2 -
                 restart_text.get_width()//2, HEIGHT//2 + 60))


def show_hud(surface, player, font, small_font):
    # 分数
    score_text = font.render(f"分数: {player.score}", True, WHITE)
    surface.blit(score_text, (20, 20))

    # 武器等级
    weapon_text = small_font.render(
        f"武器等级: {player.weapon_level}", True, YELLOW)
    surface.blit(weapon_text, (20, 50))

    # 玩家生命条
    draw_health_bar(surface, player, 20, 80, 120, 20)

    # 敌人统计
    enemy_info = font.render(f"敌人: 3种类型", True, (200, 200, 200))
    surface.blit(enemy_info, (WIDTH - 150, 20))

    # 操作说明
    controls = [
        "控制: ←↑→↓ 移动, 空格/按住射击",
        "红: 普通敌人, 浅红: 快速敌人, 蓝: 坦克敌人",
        "绿: 生命回复, 紫: 武器升级, 蓝: 护盾"
    ]

    for i, text in enumerate(controls):
        control_text = small_font.render(text, True, (200, 200, 200))
        surface.blit(control_text, (20, HEIGHT - 80 + i * 20))


def main():
    player = Player()
    player_bullets = []
    enemy_bullets = []
    enemies = []
    particles = []
    powerups = []

    # 字体
    font = pygame.font.SysFont("simhei", 24)
    small_font = pygame.font.SysFont("simhei", 18)
    big_font = pygame.font.SysFont("simhei", 48)

    enemy_spawn_timer = 0
    powerup_spawn_timer = random.randint(300, 600)
    spawn_interval = 60

    game_over = False
    running = True

    # 背景星星
    stars = []
    for _ in range(100):
        stars.append({
            "x": random.randint(0, WIDTH),
            "y": random.randint(0, HEIGHT),
            "speed": random.uniform(0.1, 0.5),
            "size": random.randint(1, 3)
        })

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_SPACE:
                    player.shoot(player_bullets)

                if game_over and event.key == pygame.K_r:
                    # 重新开始游戏
                    player = Player()
                    player_bullets.clear()
                    enemy_bullets.clear()
                    enemies.clear()
                    particles.clear()
                    powerups.clear()
                    game_over = False
                    spawn_interval = 60

                if event.key == pygame.K_ESCAPE:
                    running = False

        if not game_over:
            keys = pygame.key.get_pressed()
            player.move(keys)

            # 自动射击
            if keys[pygame.K_SPACE]:
                player.shoot(player_bullets)

            player.update()

            # 更新玩家子弹
            for bullet in player_bullets[:]:
                bullet.update()
                if bullet.is_off_screen():
                    player_bullets.remove(bullet)

            # 更新敌人子弹
            for bullet in enemy_bullets[:]:
                bullet.update()

                # 检查敌人子弹是否击中玩家
                if (bullet.x >= player.x and bullet.x <= player.x + player.width and
                        bullet.y >= player.y and bullet.y <= player.y + player.height):

                    if player.take_damage(bullet.damage):
                        particles.extend(
                            [Particle(bullet.x, bullet.y, RED) for _ in range(10)])

                        if player.health <= 0:
                            game_over = True

                    if bullet in enemy_bullets:
                        enemy_bullets.remove(bullet)
                    continue

                if bullet.is_off_screen():
                    enemy_bullets.remove(bullet)

            # 生成敌人
            enemy_spawn_timer += 1
            if enemy_spawn_timer >= spawn_interval:
                # 根据分数增加敌人难度
                if player.score < 500:
                    enemy_type = "normal"
                elif player.score < 1000:
                    enemy_type = random.choice(["normal", "fast"])
                else:
                    enemy_type = random.choice(["normal", "fast", "tank"])

                enemies.append(Enemy(enemy_type))
                enemy_spawn_timer = 0

                # 随着分数增加，生成间隔减少
                if spawn_interval > 20:
                    spawn_interval = max(20, 60 - player.score // 100)

            # 更新敌人
            for enemy in enemies[:]:
                enemy.update(player.x, player.y, enemy_bullets)

                # 检查敌人是否与玩家碰撞
                if player.invincible <= 0 and player.get_rect().colliderect(enemy.get_rect()):
                    if player.take_damage(30):
                        particles.extend([Particle(
                            enemy.x + enemy.width//2, enemy.y + enemy.height//2, RED) for _ in range(20)])

                        if player.health <= 0:
                            game_over = True
                    continue

                # 检查敌人是否与玩家子弹碰撞
                for bullet in player_bullets[:]:
                    if enemy.get_rect().colliderect(pygame.Rect(bullet.x, bullet.y, bullet.width, bullet.height)):
                        enemy.health -= bullet.damage

                        particles.extend(
                            [Particle(bullet.x, bullet.y, YELLOW) for _ in range(5)])

                        if bullet in player_bullets:
                            player_bullets.remove(bullet)

                        if enemy.health <= 0:
                            if enemy.enemy_type == "normal":
                                player.score += 100
                            elif enemy.enemy_type == "fast":
                                player.score += 150
                            elif enemy.enemy_type == "tank":
                                player.score += 300

                            particles.extend([Particle(
                                enemy.x + enemy.width//2, enemy.y + enemy.height//2, ORANGE) for _ in range(20)])

                            # 掉落道具
                            drop_chance = 0.1  # 10%几率
                            if enemy.enemy_type == "tank":
                                drop_chance = 0.3  # 坦克有30%几率

                            if random.random() < drop_chance:
                                powerups.append(
                                    PowerUp(enemy.x + enemy.width//2, enemy.y + enemy.height//2))

                            enemies.remove(enemy)
                        break

                if enemy.is_off_screen():
                    enemies.remove(enemy)

            # 更新道具
            powerup_spawn_timer -= 1
            if powerup_spawn_timer <= 0:
                powerups.append(PowerUp(random.randint(0, WIDTH-20), -20))
                powerup_spawn_timer = random.randint(300, 600)

            for powerup in powerups[:]:
                powerup.update()

                # 检查玩家是否获得道具
                if player.get_rect().colliderect(powerup.get_rect()):
                    if powerup.type == "health":
                        player.health = min(
                            player.max_health, player.health + 30)
                        player.score += 50
                    elif powerup.type == "weapon" and player.weapon_level < 3:
                        player.weapon_level += 1
                        player.score += 100
                    elif powerup.type == "shield":
                        player.invincible = 180  # 3秒无敌
                        player.score += 50

                    particles.extend([Particle(powerup.x + powerup.width//2, powerup.y +
                                     powerup.height//2, powerup.color) for _ in range(20)])
                    powerups.remove(powerup)

                if powerup.is_off_screen():
                    powerups.remove(powerup)

            # 更新粒子
            for particle in particles[:]:
                particle.update()
                if particle.is_dead():
                    particles.remove(particle)

        # 绘制
        screen.fill(DARK_BLUE)

        # 绘制背景星星
        for star in stars:
            star["y"] += star["speed"]
            if star["y"] > HEIGHT:
                star["y"] = 0
                star["x"] = random.randint(0, WIDTH)
            pygame.draw.circle(
                screen, WHITE, (int(star["x"]), int(star["y"])), star["size"])

        # 绘制子弹
        for bullet in player_bullets:
            bullet.draw(screen)

        for bullet in enemy_bullets:
            bullet.draw(screen)

        # 绘制敌人
        for enemy in enemies:
            enemy.draw(screen)

        # 绘制道具
        for powerup in powerups:
            powerup.draw(screen)

        # 绘制粒子
        for particle in particles:
            particle.draw(screen)

        # 绘制玩家
        player.draw(screen)

        # 绘制HUD
        show_hud(screen, player, font, small_font)

        if game_over:
            show_game_over(screen, player.score, font, big_font)

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

    pygame.quit()
    sys.exit()


if __name__ == "__main__":
    print("二维移动射击游戏 - 增强版")
    print("游戏特性:")
    print("1. 三种敌人: 普通(红), 快速(浅红), 坦克(蓝)")
    print("2. 敌人会发射追踪子弹攻击玩家")
    print("3. 三种道具: 生命(绿), 武器升级(紫), 护盾(蓝)")
    print("4. 分数越高，敌人越强，生成越快")
    print("5. 按R键重新开始游戏")
    print("-" * 40)
    main()