import pygame
import random
import sys

# --- 初始化 Pygame ---
pygame.init()

# --- 常量定义 ---
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
FPS = 60

# 颜色定义 (R, G, B)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GREEN = (0, 255, 0)   # 玩家颜色
RED = (255, 0, 0)     # 敌人颜色
YELLOW = (255, 255, 0)  # 子弹颜色

# 游戏设置
PLAYER_SPEED = 5
BULLET_SPEED = 7
ENEMY_SPEED = 3
ENEMY_SPAWN_RATE = 30  # 帧数间隔，数值越小生成越快

# 设置屏幕
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Pygame 二维移动射击演示")
clock = pygame.time.Clock()
font = pygame.font.SysFont("Arial", 24)


class Player(pygame.sprite.Sprite):
    def __init__(self):
        super().__init__()
        self.image = pygame.Surface((40, 40))
        self.image.fill(GREEN)
        self.rect = self.image.get_rect()
        self.rect.centerx = SCREEN_WIDTH // 2
        self.rect.bottom = SCREEN_HEIGHT - 10
        self.speed_x = 0
        self.speed_y = 0

    def update(self):
        # 键盘控制
        keys = pygame.key.get_pressed()
        self.speed_x = 0
        self.speed_y = 0

        if keys[pygame.K_LEFT] or keys[pygame.K_a]:
            self.speed_x = -PLAYER_SPEED
        if keys[pygame.K_RIGHT] or keys[pygame.K_d]:
            self.speed_x = PLAYER_SPEED
        if keys[pygame.K_UP] or keys[pygame.K_w]:
            self.speed_y = -PLAYER_SPEED
        if keys[pygame.K_DOWN] or keys[pygame.K_s]:
            self.speed_y = PLAYER_SPEED

        self.rect.x += self.speed_x
        self.rect.y += self.speed_y

        # 边界检查
        if self.rect.left < 0:
            self.rect.left = 0
        if self.rect.right > SCREEN_WIDTH:
            self.rect.right = SCREEN_WIDTH
        if self.rect.top < 0:
            self.rect.top = 0
        if self.rect.bottom > SCREEN_HEIGHT:
            self.rect.bottom = SCREEN_HEIGHT

    def shoot(self):
        bullet = Bullet(self.rect.centerx, self.rect.top)
        all_sprites.add(bullet)
        bullets.add(bullet)


class Bullet(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.image = pygame.Surface((6, 15))
        self.image.fill(YELLOW)
        self.rect = self.image.get_rect()
        self.rect.centerx = x
        self.rect.bottom = y

    def update(self):
        self.rect.y -= BULLET_SPEED
        # 如果子弹飞出屏幕顶部，则移除
        if self.rect.bottom < 0:
            self.kill()


class Enemy(pygame.sprite.Sprite):
    def __init__(self):
        super().__init__()
        self.image = pygame.Surface((30, 30))
        self.image.fill(RED)
        self.rect = self.image.get_rect()
        self.rect.x = random.randrange(0, SCREEN_WIDTH - self.rect.width)
        self.rect.y = random.randrange(-100, -40)
        self.speed_y = random.randrange(2, ENEMY_SPEED + 2)

    def update(self):
        self.rect.y += self.speed_y
        # 如果敌人飞出屏幕底部，则移除
        if self.rect.top > SCREEN_HEIGHT:
            self.kill()


# --- 精灵组管理 ---
all_sprites = pygame.sprite.Group()
bullets = pygame.sprite.Group()
enemies = pygame.sprite.Group()

player = Player()
all_sprites.add(player)

# --- 游戏变量 ---
score = 0
game_over = False
frame_count = 0


def draw_text(surf, text, size, x, y):
    font_obj = pygame.font.SysFont("Arial", size)
    text_surface = font_obj.render(text, True, WHITE)
    text_rect = text_surface.get_rect()
    text_rect.midtop = (x, y)
    surf.blit(text_surface, text_rect)


def show_go_screen():
    screen.fill(BLACK)
    draw_text(screen, "GAME OVER", 64, SCREEN_WIDTH / 2, SCREEN_HEIGHT / 4)
    draw_text(screen, f"Final Score: {score}",
              22, SCREEN_WIDTH / 2, SCREEN_HEIGHT / 2)
    draw_text(screen, "Press SPACE to Restart or Q to Quit",
              18, SCREEN_WIDTH / 2, SCREEN_HEIGHT * 3 / 4)
    pygame.display.flip()

    waiting = True
    while waiting:
        clock.tick(FPS)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
            if event.type == pygame.KEYUP:
                if event.key == pygame.K_q:
                    return False
                if event.key == pygame.K_SPACE:
                    return True
    return False


# --- 主游戏循环 ---
running = True
while running:
    if game_over:
        if not show_go_screen():
            running = False
            continue
        # 重置游戏
        game_over = False
        all_sprites.empty()
        bullets.empty()
        enemies.empty()
        player = Player()
        all_sprites.add(player)
        score = 0
        frame_count = 0

    # 1. 处理事件
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE and not game_over:
                player.shoot()

    # 2. 更新逻辑
    if not game_over:
        all_sprites.update()

        # 生成敌人
        frame_count += 1
        if frame_count % ENEMY_SPAWN_RATE == 0:
            enemy = Enemy()
            all_sprites.add(enemy)
            enemies.add(enemy)

        # 检测子弹击中敌人
        hits = pygame.sprite.groupcollide(enemies, bullets, True, True)
        for hit in hits:
            score += 10
            # 可选：在这里添加爆炸效果或声音

        # 检测敌人撞击玩家
        hits = pygame.sprite.spritecollide(player, enemies, False)
        if hits:
            game_over = True

    # 3. 绘制渲染
    screen.fill(BLACK)
    all_sprites.draw(screen)

    # 绘制分数
    draw_text(screen, f"Score: {score}", 18, SCREEN_WIDTH - 60, 10)

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

pygame.quit()
sys.exit()