import pygame
import math
import random

# --- 初始化 ---
pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("3D迷宫打怪 - Pygame Raycasting")
clock = pygame.time.Clock()

# 颜色
BG_COLOR = (20, 20, 30)
WALL_COLOR = (100, 100, 120)
WALL_DARK = (70, 70, 90)
FLOOR_COLOR = (40, 40, 50)
MONSTER_COLOR = (200, 50, 50)
WEAPON_COLOR = (200, 200, 200)
UI_COLOR = (255, 255, 255)

FONT = pygame.font.SysFont("arial", 24)
BIG_FONT = pygame.font.SysFont("arial", 60)

# 迷宫地图 (1=墙, 0=路, 2=门, 3=宝箱, 4=怪物出生点)
MAP = [
    [1,1,1,1,1,1,1,1,1,1,1,1],
    [1,0,0,0,1,0,0,0,0,0,0,1],
    [1,0,1,0,1,0,1,1,1,1,0,1],
    [1,0,1,0,0,0,0,0,0,1,0,1],
    [1,0,1,1,1,1,1,1,0,1,0,1],
    [1,0,0,0,0,0,0,1,0,0,0,1],
    [1,1,1,1,1,1,0,1,1,1,0,1],
    [1,0,0,0,0,1,0,0,0,0,0,1],
    [1,0,1,1,0,1,1,1,1,1,0,1],
    [1,0,0,1,0,0,0,0,0,1,0,1],
    [1,1,0,1,1,1,1,1,0,1,0,1],
    [1,1,1,1,1,1,1,1,1,1,1,1],
]

MAP_H = len(MAP)
MAP_W = len(MAP[0])

# --- 玩家 ---
class Player:
    def __init__(self):
        self.x = 1.5
        self.y = 1.5
        self.angle = 0  # 弧度
        self.hp = 100
        self.attack_cd = 0
        self.attack_anim = 0  # 挥砍动画进度

    def update(self, keys, dt):
        move_x, move_y = 0, 0
        speed = 3.0 * dt
        
        if keys[pygame.K_w]:
            move_x += math.cos(self.angle) * speed
            move_y += math.sin(self.angle) * speed
        if keys[pygame.K_s]:
            move_x -= math.cos(self.angle) * speed
            move_y -= math.sin(self.angle) * speed
        if keys[pygame.K_a]:
            move_x += math.cos(self.angle - math.pi/2) * speed
            move_y += math.sin(self.angle - math.pi/2) * speed
        if keys[pygame.K_d]:
            move_x += math.cos(self.angle + math.pi/2) * speed
            move_y += math.sin(self.angle + math.pi/2) * speed

        # 碰撞检测
        new_x = self.x + move_x
        new_y = self.y + move_y
        if MAP[int(self.y)][int(new_x)] == 0: self.x = new_x
        if MAP[int(new_y)][int(self.x)] == 0: self.y = new_y

        # 鼠标转向
        mx, _ = pygame.mouse.get_pos()
        self.angle += (mx - WIDTH//2) * 0.003 * dt
        pygame.mouse.set_pos(WIDTH//2, HEIGHT//2)  # 锁定鼠标

        if self.attack_cd > 0: self.attack_cd -= dt
        if self.attack_anim > 0: self.attack_anim -= dt

    def attack(self):
        if self.attack_cd <= 0:
            self.attack_cd = 0.5
            self.attack_anim = 0.3
            return True
        return False

# --- 怪物 ---
class Monster:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.hp = 3
        self.max_hp = 3
        self.alive = True
        self.attack_cd = 0
        self.flash_timer = 0

    def update(self, player, dt):
        if not self.alive: return
        
        # 追踪玩家
        dx = player.x - self.x
        dy = player.y - self.y
        dist = math.hypot(dx, dy)
        
        if dist < 8:  # 视野内才追
            if dist > 1.0:
                speed = 1.5 * dt
                self.x += (dx / dist) * speed
                self.y += (dy / dist) * speed
            else:
                # 攻击玩家
                if self.attack_cd <= 0:
                    player.hp -= 10
                    self.attack_cd = 1.0
        
        if self.attack_cd > 0: self.attack_cd -= dt
        if self.flash_timer > 0: self.flash_timer -= dt

    def take_damage(self):
        self.hp -= 1
        self.flash_timer = 0.2
        if self.hp <= 0:
            self.alive = False

    def draw(self, surface, player, z_buffer):
        if not self.alive: return
        
        dx = self.x - player.x
        dy = self.y - player.y
        dist = math.hypot(dx, dy)
        
        # 相对角度
        angle = math.atan2(dy, dx) - player.angle
        # 归一化到 -PI ~ PI
        while angle > math.pi: angle -= 2 * math.pi
        while angle < -math.pi: angle += 2 * math.pi
        
        # 是否在视野内
        if abs(angle) > math.pi / 3: return
        
        # 屏幕X
        screen_x = WIDTH // 2 + int(angle / (math.pi / 3) * WIDTH // 2)
        # 大小
        size = int(HEIGHT / dist)
        screen_y = HEIGHT // 2 - size // 2
        
        # 深度测试
        col = int(screen_x)
        if 0 <= col < WIDTH and dist < z_buffer[col]:
            color = (255, 100, 100) if self.flash_timer > 0 else MONSTER_COLOR
            pygame.draw.rect(surface, color, (screen_x - size//2, screen_y, size, size))
            # 血条
            hp_w = size * (self.hp / self.max_hp)
            pygame.draw.rect(surface, (0, 255, 0), (screen_x - size//2, screen_y - 10, hp_w, 5))

# --- Raycasting 渲染 ---
def cast_rays(player, screen):
    z_buffer = [float('inf')] * WIDTH
    fov = math.pi / 3  # 60度视野
    num_rays = WIDTH
    
    for i in range(num_rays):
        ray_angle = player.angle - fov / 2 + (i / num_rays) * fov
        
        # DDA 算法
        sin_a = math.sin(ray_angle)
        cos_a = math.cos(ray_angle)
        
        # 防止除0
        if abs(sin_a) < 1e-6: sin_a = 1e-6
        if abs(cos_a) < 1e-6: cos_a = 1e-6
        
        # 水平/垂直交点
        dist_h = float('inf')
        dist_v = float('inf')
        
        # 水平线
        y_hor = int(player.y) + (1 if sin_a > 0 else 0)
        x_hor = player.x + (y_hor - player.y) / math.tan(ray_angle)
        step_y = 1 if sin_a > 0 else -1
        step_x = step_y / math.tan(ray_angle)
        
        for _ in range(20):
            if 0 <= y_hor < MAP_H and 0 <= int(x_hor) < MAP_W:
                if MAP[y_hor][int(x_hor)] == 1:
                    dist_h = math.hypot(x_hor - player.x, y_hor - player.y)
                    break
            x_hor += step_x
            y_hor += step_y
            
        # 垂直线
        x_ver = int(player.x) + (1 if cos_a > 0 else 0)
        y_ver = player.y + (x_ver - player.x) * math.tan(ray_angle)
        step_x = 1 if cos_a > 0 else -1
        step_y = step_x * math.tan(ray_angle)
        
        for _ in range(20):
            if 0 <= int(y_ver) < MAP_H and 0 <= x_ver < MAP_W:
                if MAP[int(y_ver)][x_ver] == 1:
                    dist_v = math.hypot(x_ver - player.x, y_ver - player.y)
                    break
            y_ver += step_y
            x_ver += step_x
        
        # 取最近
        if dist_h < dist_v:
            dist = dist_h
            shade = 1.0
        else:
            dist = dist_v
            shade = 0.7
            
        # 鱼眼修正
        dist *= math.cos(ray_angle - player.angle)
        z_buffer[i] = dist
        
        # 墙高
        wall_h = int(HEIGHT / dist)
        wall_top = HEIGHT // 2 - wall_h // 2
        
        # 颜色
        c = int(100 * shade / (dist * 0.5 + 1))
        c = max(20, min(150, c))
        color = (c, c, c + 20)
        
        pygame.draw.line(screen, color, (i, wall_top), (i, wall_top + wall_h))

# --- 主程序 ---
def main():
    player = Player()
    monsters = [Monster(5.5, 5.5), Monster(9.5, 9.5)]  # 初始怪物
    game_over = False
    win = False
    
    running = True
    while running:
        dt = clock.tick(60) / 1000.0
        for event in pygame.event.get():
            if event.type == pygame.QUIT: running = False
            if event.type == pygame.MOUSEBUTTONDOWN and not game_over:
                if player.attack():
                    # 攻击判定：前方1.5格内，角度±30度
                    for m in monsters:
                        if not m.alive: continue
                        dx = m.x - player.x
                        dy = m.y - player.y
                        dist = math.hypot(dx, dy)
                        angle = math.atan2(dy, dx) - player.angle
                        while angle > math.pi: angle -= 2*math.pi
                        while angle < -math.pi: angle += 2*math.pi
                        
                        if dist < 1.5 and abs(angle) < math.pi/6:
                            m.take_damage()
                            # 击退
                            m.x += math.cos(player.angle) * 0.5
                            m.y += math.sin(player.angle) * 0.5
        
        if not game_over:
            player.update(pygame.key.get_pressed(), dt)
            for m in monsters: m.update(player, dt)
            
            if player.hp <= 0: game_over = True
            if all(not m.alive for m in monsters): win = True

        # --- 绘图 ---
        screen.fill(BG_COLOR)
        
        # 地板
        pygame.draw.rect(screen, FLOOR_COLOR, (0, HEIGHT//2, WIDTH, HEIGHT//2))
        
        # 3D墙壁
        cast_rays(player, screen)
        
        # 怪物
        z_buffer = [float('inf')] * WIDTH  # 简化版，实际应该用cast_rays返回的
        # 这里重新算一次z_buffer太慢，实际应该复用，但为了代码简洁，怪物用简单深度测试
        for m in monsters: m.draw(screen, player, [10]*WIDTH)  # 临时z_buffer
        
        # 武器
        if player.attack_anim > 0:
            # 挥砍动画
            offset = int(player.attack_anim * 200)
            pygame.draw.polygon(screen, WEAPON_COLOR, [
                (WIDTH//2 - 50, HEIGHT - 100 + offset),
                (WIDTH//2 + 50, HEIGHT - 100 + offset),
                (WIDTH//2, HEIGHT - 250 + offset)
            ])
        else:
            # 待机
            pygame.draw.rect(screen, WEAPON_COLOR, (WIDTH//2 - 20, HEIGHT - 150, 40, 100))

        # UI
        hp_txt = FONT.render(f"HP: {player.hp}", True, (255, 50, 50))
        screen.blit(hp_txt, (20, 20))
        
        # 小地图
        map_scale = 8
        map_x = WIDTH - MAP_W * map_scale - 10
        map_y = 10
        for y in range(MAP_H):
            for x in range(MAP_W):
                if MAP[y][x] == 1:
                    pygame.draw.rect(screen, (150,150,150), (map_x + x*map_scale, map_y + y*map_scale, map_scale, map_scale))
        # 玩家
        px = map_x + int(player.x * map_scale)
        py = map_y + int(player.y * map_scale)
        pygame.draw.circle(screen, (0, 255, 0), (px, py), 3)
        # 怪物
        for m in monsters:
            if m.alive:
                mx = map_x + int(m.x * map_scale)
                my = map_y + int(m.y * map_scale)
                pygame.draw.circle(screen, (255, 0, 0), (mx, my), 3)

        if game_over:
            txt = BIG_FONT.render("YOU DIED", True, (255, 0, 0))
            screen.blit(txt, txt.get_rect(center=(WIDTH//2, HEIGHT//2)))
        elif win:
            txt = BIG_FONT.render("VICTORY!", True, (0, 255, 0))
            screen.blit(txt, txt.get_rect(center=(WIDTH//2, HEIGHT//2)))

        pygame.display.flip()

    pygame.quit()

if __name__ == "__main__":
    main()