import pygame
import sys
import random
import math

# ========== 初始化 ==========
pygame.init()
WINDOW_WIDTH, WINDOW_HEIGHT = 1000, 700
screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
pygame.display.set_caption("🧟 僵尸病毒 – 靠近自动感染")
clock = pygame.time.Clock()

# ========== 常量 ==========
WORLD_WIDTH = 1000
WORLD_HEIGHT = 1000
TILE_SIZE = 35
GRID_COLS = WORLD_WIDTH // TILE_SIZE
GRID_ROWS = WORLD_HEIGHT // TILE_SIZE

WALL_CLUSTER_COUNT = 40
WALL_CLUSTER_MAX = 6

PLAYER_SPEED = 4.5
ZOMBIE_SPEED = 2.8
HUMAN_SPEED = 2.0
BULLET_SPEED = 12
PLAYER_MAX_HP = 2000
HUMAN_MAX_HP = 50
ZOMBIE_MAX_HP = 60
ZOMBIE_ATTACK_DAMAGE = 8
ZOMBIE_ATTACK_CD = 40
HUMAN_SHOOT_CD = 30
HUMAN_SHOOT_RANGE = 250
BULLET_DAMAGE = 25
PLAYER_REGEN = 0.05
INFECT_DISTANCE = 25              # 感染距离

# ========== 华丽色彩 ==========
COLOR_BG = (15, 15, 20)
COLOR_WALL = (55, 55, 60)
COLOR_WALL_EDGE = (90, 90, 95)
COLOR_FLOOR = (30, 30, 30)
COLOR_GRID = (45, 45, 45)
COLOR_PLAYER = (80, 180, 255)
COLOR_PLAYER_HURT = (255, 80, 80)
COLOR_ZOMBIE = (130, 255, 130)
COLOR_HUMAN = (255, 140, 140)
COLOR_BULLET = (255, 255, 120)
COLOR_BULLET_TRAIL = (255, 200, 80)
COLOR_HP_RED = (255, 50, 50)
COLOR_HP_BG = (25, 25, 25)
COLOR_UI = (240, 240, 240)

font_small = pygame.font.Font(None, 26)
font_med = pygame.font.Font(None, 34)
font_large = pygame.font.Font(None, 48)

# ========== 动态背景粒子 ==========
class Star:
    def __init__(self):
        self.x = random.randint(0, WINDOW_WIDTH)
        self.y = random.randint(0, WINDOW_HEIGHT)
        self.speed = random.uniform(0.2, 0.8)
        self.size = random.randint(1, 3)
        self.brightness = random.randint(80, 200)

    def update(self):
        self.y += self.speed
        if self.y > WINDOW_HEIGHT:
            self.y = 0
            self.x = random.randint(0, WINDOW_WIDTH)

    def draw(self, surface):
        alpha = abs(math.sin(pygame.time.get_ticks() * 0.002 + self.x)) * 100 + 100
        color = (self.brightness, self.brightness, self.brightness)
        pygame.draw.circle(surface, color, (int(self.x), int(self.y)), self.size)

# ========== 地图生成 ==========
def generate_walls():
    walls = set()
    for col in range(GRID_COLS):
        walls.add((col, 0))
        walls.add((col, GRID_ROWS - 1))
    for row in range(GRID_ROWS):
        walls.add((0, row))
        walls.add((GRID_COLS - 1, row))

    center_c = WORLD_WIDTH // 2 // TILE_SIZE
    center_r = WORLD_HEIGHT // 2 // TILE_SIZE

    for _ in range(WALL_CLUSTER_COUNT):
        cw = random.randint(2, WALL_CLUSTER_MAX)
        ch = random.randint(2, WALL_CLUSTER_MAX)
        cx = random.randint(2, GRID_COLS - cw - 2)
        cy = random.randint(2, GRID_ROWS - ch - 2)
        if abs(cx + cw//2 - center_c) < 3 and abs(cy + ch//2 - center_r) < 3:
            continue
        for dx in range(cw):
            for dy in range(ch):
                walls.add((cx + dx, cy + dy))
    return walls

def is_wall(wx, wy, walls):
    col = int(wx // TILE_SIZE)
    row = int(wy // TILE_SIZE)
    return (col, row) in walls

def world_to_screen(wx, wy, cam_x, cam_y):
    return (wx - cam_x, wy - cam_y)

# ========== 实体类 ==========
class Entity:
    def __init__(self, x, y, radius, color, speed, max_hp):
        self.x = x
        self.y = y
        self.radius = radius
        self.color = color
        self.speed = speed
        self.max_hp = max_hp
        self.hp = max_hp
        self.flash_timer = 0

    def move(self, dx, dy, walls):
        new_x = self.x + dx
        new_y = self.y + dy
        if not is_wall(new_x, self.y, walls):
            self.x = new_x
        if not is_wall(self.x, new_y, walls):
            self.y = new_y
        self.x = max(self.radius, min(WORLD_WIDTH - self.radius, self.x))
        self.y = max(self.radius, min(WORLD_HEIGHT - self.radius, self.y))

    def take_damage(self, dmg):
        self.hp -= dmg
        self.flash_timer = 15

    def draw(self, surface, cam_x, cam_y):
        sx, sy = world_to_screen(self.x, self.y, cam_x, cam_y)
        if sx < -self.radius or sx > WINDOW_WIDTH + self.radius or sy < -self.radius or sy > WINDOW_HEIGHT + self.radius:
            return

        if self.flash_timer > 0:
            draw_color = COLOR_PLAYER_HURT if isinstance(self, Player) else (255, 100, 100)
            self.flash_timer -= 1
        else:
            draw_color = self.color

        pygame.draw.circle(surface, (0,0,0, 180), (int(sx+3), int(sy+3)), self.radius)
        for r in range(self.radius + 4, self.radius, -2):
            alpha = 60 - (self.radius + 4 - r) * 10
            if alpha > 0:
                glow_surf = pygame.Surface((r*2, r*2), pygame.SRCALPHA)
                pygame.draw.circle(glow_surf, (*draw_color, alpha), (r, r), r)
                surface.blit(glow_surf, (sx - r, sy - r))
        pygame.draw.circle(surface, draw_color, (int(sx), int(sy)), self.radius)
        light = tuple(min(255, c + 100) for c in draw_color)
        pygame.draw.circle(surface, light, (int(sx - self.radius*0.3), int(sy - self.radius*0.3)), self.radius//3)

        if self.hp < self.max_hp:
            bw = self.radius * 2 + 4
            bh = 5
            bx = sx - bw//2
            by = sy - self.radius - 10
            fill = (self.hp / self.max_hp) * bw
            pygame.draw.rect(surface, COLOR_HP_BG, (bx, by, bw, bh), border_radius=2)
            r_val = max(0, min(255, 255 - int(200 * fill / bw)))
            g_val = max(0, min(255, int(200 * fill / bw)))
            hp_color = (r_val, g_val, 30)
            pygame.draw.rect(surface, hp_color, (bx, by, fill, bh), border_radius=2)

class Player(Entity):
    def __init__(self, x, y):
        super().__init__(x, y, 13, COLOR_PLAYER, PLAYER_SPEED, PLAYER_MAX_HP)
        self.shoot_cd = 0

    def update(self, keys, walls, bullets):
        dx, dy = 0, 0
        if keys[pygame.K_LEFT] or keys[pygame.K_a]: dx -= 1
        if keys[pygame.K_RIGHT] or keys[pygame.K_d]: dx += 1
        if keys[pygame.K_UP] or keys[pygame.K_w]: dy -= 1
        if keys[pygame.K_DOWN] or keys[pygame.K_s]: dy += 1
        if dx != 0 or dy != 0:
            length = math.hypot(dx, dy)
            dx = dx / length * self.speed
            dy = dy / length * self.speed
        self.move(dx, dy, walls)
        if self.hp < self.max_hp:
            self.hp = min(self.max_hp, self.hp + PLAYER_REGEN)
        if self.shoot_cd > 0:
            self.shoot_cd -= 1

    def shoot(self, target_wx, target_wy, bullets):
        if self.shoot_cd > 0: return
        dx = target_wx - self.x
        dy = target_wy - self.y
        dist = math.hypot(dx, dy)
        if dist == 0: return
        vx = (dx / dist) * BULLET_SPEED
        vy = (dy / dist) * BULLET_SPEED
        bullets.append(Bullet(self.x, self.y, vx, vy, BULLET_DAMAGE, True))
        self.shoot_cd = 10

class Human(Entity):
    def __init__(self, x, y):
        super().__init__(x, y, 10, COLOR_HUMAN, HUMAN_SPEED, HUMAN_MAX_HP)
        self.shoot_cd = 0

    def update(self, zombies, player, walls, bullets):
        threat = zombies + ([player] if player.hp > 0 else [])
        if not threat: return
        closest = min(threat, key=lambda z: math.hypot(self.x - z.x, self.y - z.y))
        dist = math.hypot(self.x - closest.x, self.y - closest.y)
        if dist < HUMAN_SHOOT_RANGE * 0.7:
            dx = self.x - closest.x
            dy = self.y - closest.y
            if dist > 0:
                dx = dx / dist * self.speed
                dy = dy / dist * self.speed
            self.move(dx, dy, walls)
            if self.shoot_cd <= 0 and dist < HUMAN_SHOOT_RANGE:
                self.shoot_at(closest, bullets)
                self.shoot_cd = HUMAN_SHOOT_CD
        if self.shoot_cd > 0:
            self.shoot_cd -= 1

    def shoot_at(self, target, bullets):
        dx = target.x - self.x
        dy = target.y - self.y
        dist = math.hypot(dx, dy)
        if dist == 0: return
        vx = (dx / dist) * BULLET_SPEED
        vy = (dy / dist) * BULLET_SPEED
        bullets.append(Bullet(self.x, self.y, vx, vy, BULLET_DAMAGE, False))

class Zombie(Entity):
    def __init__(self, x, y):
        super().__init__(x, y, 12, COLOR_ZOMBIE, ZOMBIE_SPEED, ZOMBIE_MAX_HP)
        self.attack_cd = 0

    def update(self, targets, walls):
        if not targets: return
        closest = min(targets, key=lambda t: math.hypot(self.x - t.x, self.y - t.y))
        dx = closest.x - self.x
        dy = closest.y - self.y
        dist = math.hypot(dx, dy)
        if dist > 0:
            dx = dx / dist * self.speed
            dy = dy / dist * self.speed
        self.move(dx, dy, walls)
        if dist < self.radius + closest.radius + 2 and self.attack_cd <= 0:
            closest.take_damage(ZOMBIE_ATTACK_DAMAGE)
            self.attack_cd = ZOMBIE_ATTACK_CD
        if self.attack_cd > 0:
            self.attack_cd -= 1

class Bullet:
    def __init__(self, x, y, vx, vy, damage, is_player):
        self.x = x
        self.y = y
        self.vx = vx
        self.vy = vy
        self.damage = damage
        self.is_player = is_player
        self.alive = True
        self.trail = []

    def update(self, walls):
        self.trail.append((self.x, self.y))
        if len(self.trail) > 8:
            self.trail.pop(0)
        self.x += self.vx
        self.y += self.vy
        if is_wall(self.x, self.y, walls) or not (0 <= self.x <= WORLD_WIDTH and 0 <= self.y <= WORLD_HEIGHT):
            self.alive = False

    def draw(self, surface, cam_x, cam_y):
        if not self.alive: return
        for i, (tx, ty) in enumerate(self.trail):
            alpha = int(180 * (i / len(self.trail)))
            sx, sy = world_to_screen(tx, ty, cam_x, cam_y)
            if 0 <= sx <= WINDOW_WIDTH and 0 <= sy <= WINDOW_HEIGHT:
                pygame.draw.circle(surface, (*COLOR_BULLET_TRAIL, alpha), (int(sx), int(sy)), 2 + i//2)
        sx, sy = world_to_screen(self.x, self.y, cam_x, cam_y)
        if 0 <= sx <= WINDOW_WIDTH and 0 <= sy <= WINDOW_HEIGHT:
            pygame.draw.circle(surface, COLOR_BULLET, (int(sx), int(sy)), 4)
            pygame.draw.circle(surface, (255, 255, 255), (int(sx-1), int(sy-1)), 2)

# ========== 游戏主类 ==========
class Game:
    def __init__(self):
        self.walls = generate_walls()
        cx, cy = WORLD_WIDTH // 2, WORLD_HEIGHT // 2
        if is_wall(cx, cy, self.walls):
            for offset in range(1, 20):
                if not is_wall(cx + offset * 10, cy, self.walls):
                    cx += offset * 10
                    break
        self.player = Player(cx, cy)
        self.humans = []
        self.zombies = []
        self.bullets = []
        self.stars = [Star() for _ in range(60)]
        self.camera_x = self.player.x - WINDOW_WIDTH // 2
        self.camera_y = self.player.y - WINDOW_HEIGHT // 2
        self.spawn_entities()

    def spawn_entities(self):
        for _ in range(20):  # 增加人类数量，更容易遇到
            for _ in range(50):
                x = random.uniform(100, WORLD_WIDTH - 100)
                y = random.uniform(100, WORLD_HEIGHT - 100)
                if not is_wall(x, y, self.walls) and math.hypot(x - self.player.x, y - self.player.y) > 100:
                    self.humans.append(Human(x, y))
                    break
        for _ in range(8):
            for _ in range(50):
                x = random.uniform(100, WORLD_WIDTH - 100)
                y = random.uniform(100, WORLD_HEIGHT - 100)
                if not is_wall(x, y, self.walls) and math.hypot(x - self.player.x, y - self.player.y) > 100:
                    self.zombies.append(Zombie(x, y))
                    break

    def update(self):
        keys = pygame.key.get_pressed()
        if self.player.hp > 0:
            self.player.update(keys, self.walls, self.bullets)
            if pygame.mouse.get_pressed()[0]:
                mx, my = pygame.mouse.get_pos()
                wx = self.player.x + (mx - WINDOW_WIDTH // 2)
                wy = self.player.y + (my - WINDOW_HEIGHT // 2)
                self.player.shoot(wx, wy, self.bullets)

            # ★ 核心感染逻辑：玩家靠近人类时自动转化
            for human in self.humans[:]:
                if math.hypot(self.player.x - human.x, self.player.y - human.y) < INFECT_DISTANCE:
                    self.humans.remove(human)
                    self.zombies.append(Zombie(human.x, human.y))

        # 人类AI
        for human in self.humans[:]:
            if human.hp <= 0:
                self.zombies.append(Zombie(human.x, human.y))
                self.humans.remove(human)
                continue
            human.update(self.zombies, self.player, self.walls, self.bullets)

        # 僵尸AI
        targets = self.humans + ([self.player] if self.player.hp > 0 else [])
        for zombie in self.zombies[:]:
            if zombie.hp <= 0:
                self.zombies.remove(zombie)
                continue
            zombie.update(targets, self.walls)

        # 子弹碰撞
        for bullet in self.bullets[:]:
            bullet.update(self.walls)
            if not bullet.alive:
                self.bullets.remove(bullet)
                continue
            if bullet.is_player:
                for z in self.zombies:
                    if math.hypot(bullet.x - z.x, bullet.y - z.y) < z.radius:
                        z.take_damage(bullet.damage)
                        bullet.alive = False
                        break
            else:
                for z in self.zombies:
                    if math.hypot(bullet.x - z.x, bullet.y - z.y) < z.radius:
                        z.take_damage(bullet.damage)
                        bullet.alive = False
                        break
                if bullet.alive and self.player.hp > 0:
                    if math.hypot(bullet.x - self.player.x, bullet.y - self.player.y) < self.player.radius:
                        self.player.take_damage(bullet.damage)
                        bullet.alive = False

        self.camera_x = self.player.x - WINDOW_WIDTH // 2
        self.camera_y = self.player.y - WINDOW_HEIGHT // 2
        for star in self.stars:
            star.update()

    def draw_walls(self):
        start_col = max(0, int(self.camera_x // TILE_SIZE) - 1)
        end_col = min(GRID_COLS, int((self.camera_x + WINDOW_WIDTH) // TILE_SIZE) + 2)
        start_row = max(0, int(self.camera_y // TILE_SIZE) - 1)
        end_row = min(GRID_ROWS, int((self.camera_y + WINDOW_HEIGHT) // TILE_SIZE) + 2)
        for col in range(start_col, end_col):
            for row in range(start_row, end_row):
                wx = col * TILE_SIZE
                wy = row * TILE_SIZE
                sx, sy = world_to_screen(wx, wy, self.camera_x, self.camera_y)
                if (col, row) in self.walls:
                    pygame.draw.rect(screen, COLOR_WALL, (sx, sy, TILE_SIZE, TILE_SIZE))
                    pygame.draw.rect(screen, COLOR_WALL_EDGE, (sx, sy, TILE_SIZE, TILE_SIZE), 1)
                else:
                    pygame.draw.rect(screen, COLOR_FLOOR, (sx, sy, TILE_SIZE, TILE_SIZE))
        for col in range(start_col, end_col):
            x = col * TILE_SIZE - self.camera_x
            if x < 0 or x > WINDOW_WIDTH: continue
            pygame.draw.line(screen, COLOR_GRID, (x, 0), (x, WINDOW_HEIGHT))
        for row in range(start_row, end_row):
            y = row * TILE_SIZE - self.camera_y
            if y < 0 or y > WINDOW_HEIGHT: continue
            pygame.draw.line(screen, COLOR_GRID, (0, y), (WINDOW_WIDTH, y))

    def draw_ui(self):
        panel_surf = pygame.Surface((240, 80), pygame.SRCALPHA)
        panel_surf.fill((0,0,0, 140))
        screen.blit(panel_surf, (15, 15))

        bw, bh = 210, 22
        bx, by = 20, 20
        fill = (self.player.hp / PLAYER_MAX_HP) * bw
        pygame.draw.rect(screen, COLOR_HP_BG, (bx, by, bw, bh), border_radius=4)
        for i in range(int(fill)):
            ratio = i / bw
            r = max(0, min(255, 255 - int(200 * ratio)))
            g = max(0, min(255, int(200 * ratio)))
            pygame.draw.line(screen, (r, g, 30), (bx + i, by), (bx + i, by + bh))
        hp_txt = font_small.render(f"HP: {int(self.player.hp)}/{PLAYER_MAX_HP}", True, COLOR_UI)
        screen.blit(hp_txt, (bx + 5, by + 2))
        stats = f"Humans: {len(self.humans)}  Zombies: {len(self.zombies)}"
        stats_txt = font_small.render(stats, True, COLOR_UI)
        screen.blit(stats_txt, (bx, by + bh + 5))
        hint = font_small.render("WASD move | Mouse shoot | Walk to humans to infect", True, COLOR_UI)
        screen.blit(hint, (WINDOW_WIDTH - hint.get_width() - 20, 20))

        cross_size = 10
        cx, cy = WINDOW_WIDTH // 2, WINDOW_HEIGHT // 2
        cross_surf = pygame.Surface((cross_size*2, 2), pygame.SRCALPHA)
        cross_surf.fill((255, 255, 255, 120))
        screen.blit(cross_surf, (cx - cross_size, cy - 1))
        cross_surf = pygame.Surface((2, cross_size*2), pygame.SRCALPHA)
        cross_surf.fill((255, 255, 255, 120))
        screen.blit(cross_surf, (cx - 1, cy - cross_size))

        if self.player.hp <= 0:
            over = font_large.render("YOU DIED", True, (255, 50, 50))
            screen.blit(over, (WINDOW_WIDTH//2 - over.get_width()//2, WINDOW_HEIGHT//2 - 30))

    def draw(self):
        screen.fill(COLOR_BG)
        for star in self.stars:
            star.draw(screen)
        self.draw_walls()
        for human in self.humans:
            human.draw(screen, self.camera_x, self.camera_y)
        for zombie in self.zombies:
            zombie.draw(screen, self.camera_x, self.camera_y)
        if self.player.hp > 0:
            self.player.draw(screen, self.camera_x, self.camera_y)
        for bullet in self.bullets:
            bullet.draw(screen, self.camera_x, self.camera_y)
        self.draw_ui()

    def run(self):
        running = True
        while running:
            dt = clock.tick(60)
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                if event.type == pygame.KEYDOWN:
                    if event.key == pygame.K_r:
                        self.__init__()
            self.update()
            self.draw()
            pygame.display.flip()
        pygame.quit()
        sys.exit()

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