import pygame
import sys
import math
import random

# ========== 初始化 ==========
pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("🏙️ 3D小镇 · 鼠标视角 + 键盘移动")
clock = pygame.time.Clock()

# ========== 字体 ==========
font_small = pygame.font.Font(None, 24)
font_med = pygame.font.Font(None, 32)

# ========== 相机参数 ==========
FOV = 90
NEAR = 0.1
cam_x = 50.0
cam_z = 50.0
cam_y = 1.8
yaw = 0.0
pitch = 0.0

MOVE_SPEED = 8.0
MOUSE_SENS = 0.002

GRID_SIZE = 100
SPACING = 1.0

# ========== 3D投影函数 ==========
def project(wx, wy, wz):
    dx = wx - cam_x
    dy = wy - cam_y
    dz = wz - cam_z

    cos_y = math.cos(-yaw)
    sin_y = math.sin(-yaw)
    rx = dx * cos_y - dz * sin_y
    rz = dx * sin_y + dz * cos_y
    ry = dy

    cos_p = math.cos(pitch)
    sin_p = math.sin(pitch)
    rz2 = rz * cos_p - ry * sin_p
    ry2 = rz * sin_p + ry * cos_p
    rx2 = rx

    if rz2 <= NEAR:
        return None

    f = (WIDTH / 2) / math.tan(math.radians(FOV) / 2)
    sx = int(WIDTH / 2 + (rx2 / rz2) * f)
    sy = int(HEIGHT / 2 - (ry2 / rz2) * f)
    return (sx, sy)

# ========== 场景元素 ==========
building_types = [
    {"w": 6, "d": 5, "h": 4,  "roof": (200, 100, 80),  "wall": (240, 220, 180), "window": (150, 200, 255)},
    {"w": 8, "d": 6, "h": 8,  "roof": (80, 80, 160),   "wall": (200, 200, 220), "window": (255, 255, 150)},
    {"w": 10,"d": 8, "h": 12, "roof": (60, 60, 140),   "wall": (180, 180, 200), "window": (255, 200, 100)},
    {"w": 5, "d": 4, "h": 3,  "roof": (180, 130, 70),  "wall": (220, 200, 160), "window": (200, 230, 255)},
    {"w": 7, "d": 5, "h": 6,  "roof": (140, 100, 100), "wall": (210, 190, 170), "window": (255, 220, 180)},
]

# 生成建筑（避开道路）
buildings = []
for _ in range(18):
    for attempt in range(50):
        x = random.uniform(10, 85)
        z = random.uniform(10, 85)
        btype = random.choice(building_types)
        w, d = btype["w"], btype["d"]
        if (42 - w/2 < x < 58 + w/2) or (42 - d/2 < z < 58 + d/2):
            continue
        overlap = False
        for b in buildings:
            if abs(x - b["x"]) < (w/2 + b["type"]["w"]/2 + 4) and abs(z - b["z"]) < (d/2 + b["type"]["d"]/2 + 4):
                overlap = True
                break
        if not overlap:
            buildings.append({"x": x, "z": z, "type": btype})
            break

# 树木
trees = []
for _ in range(60):
    for attempt in range(50):
        x = random.uniform(5, 95)
        z = random.uniform(5, 95)
        if (42 < x < 58) or (42 < z < 58):
            continue
        overlap = False
        for b in buildings:
            bw, bd = b["type"]["w"], b["type"]["d"]
            if abs(x - b["x"]) < (bw/2 + 3) and abs(z - b["z"]) < (bd/2 + 3):
                overlap = True
                break
        if not overlap:
            trees.append({"x": x, "z": z, "height": random.uniform(3, 6)})
            break

# 车辆
cars = []
car_colors = [(220, 40, 40), (40, 100, 220), (240, 200, 50), (60, 180, 100)]
for _ in range(8):
    if random.random() < 0.5:
        x = random.uniform(43, 57)
        z = random.uniform(5, 95)
    else:
        z = random.uniform(43, 57)
        x = random.uniform(5, 95)
    cars.append({"x": x, "z": z, "w": 2.0, "d": 1.2, "h": 1.0, "color": random.choice(car_colors)})

# ========== 面构建工具 ==========
def get_depth(points):
    cx = sum(p[0] for p in points) / len(points)
    cy = sum(p[1] for p in points) / len(points)
    cz = sum(p[2] for p in points) / len(points)
    dx = cx - cam_x
    dy = cy - cam_y
    dz = cz - cam_z
    cos_y = math.cos(-yaw); sin_y = math.sin(-yaw)
    rx = dx * cos_y - dz * sin_y
    rz = dx * sin_y + dz * cos_y
    ry = dy
    cos_p = math.cos(pitch); sin_p = math.sin(pitch)
    rz2 = rz * cos_p - ry * sin_p
    return rz2

def add_quad(corners, color, faces):
    sp = [project(p[0], p[1], p[2]) for p in corners]
    if all(sp):
        faces.append((get_depth(corners), color, sp))

def add_box_faces(pts, color, faces):
    quads = [
        ([0,1,5,4], color), ([1,2,6,5], color), ([2,3,7,6], color),
        ([3,0,4,7], color), ([4,5,6,7], color), ([0,3,2,1], color)
    ]
    for idxs, col in quads:
        sp = [project(p[0], p[1], p[2]) for p in [pts[i] for i in idxs]]
        if all(sp):
            faces.append((get_depth([pts[i] for i in idxs]), col, sp))

def add_windows(origin_x, origin_z, width, height, offset, dir, color, faces):
    window_w = 0.8; window_h = 1.0
    num_x = max(1, int(width / 2))
    num_y = max(1, int(height / 2))
    for i in range(num_x):
        for j in range(num_y):
            if dir == 'z':
                wx = origin_x + (i+0.5)*(width/num_x) - window_w/2
                wy = (j+0.5)*(height/num_y)
                wz = origin_z
                pts = [(wx, wy, wz), (wx+window_w, wy, wz), (wx+window_w, wy+window_h, wz), (wx, wy+window_h, wz)]
            else:
                wz = origin_x + (i+0.5)*(width/num_x) - window_w/2
                wy = (j+0.5)*(height/num_y)
                wx = origin_z
                pts = [(wx, wy, wz), (wx, wy, wz+window_w), (wx, wy+window_h, wz+window_w), (wx, wy+window_h, wz)]
            sp = [project(p[0], p[1], p[2]) for p in pts]
            if all(sp):
                faces.append((get_depth(pts), color, sp))

def collect_all_faces():
    faces = []
    # 道路
    road_ns = [(42, 0.01, 0), (58, 0.01, 0), (58, 0.01, 100), (42, 0.01, 100)]
    add_quad(road_ns, (90, 90, 90), faces)
    road_ew = [(0, 0.01, 42), (100, 0.01, 42), (100, 0.01, 58), (0, 0.01, 58)]
    add_quad(road_ew, (90, 90, 90), faces)
    # 道路中线
    for i in range(0, 100, 10):
        if i % 20 == 0:
            add_quad([(49.5, 0.02, i), (50.5, 0.02, i), (50.5, 0.02, i+5), (49.5, 0.02, i+5)], (240, 230, 50), faces)
            add_quad([(i, 0.02, 49.5), (i+5, 0.02, 49.5), (i+5, 0.02, 50.5), (i, 0.02, 50.5)], (240, 230, 50), faces)

    # 建筑
    for b in buildings:
        x0, z0 = b["x"], b["z"]
        t = b["type"]; w, d, h = t["w"], t["d"], t["h"]
        roof_h = 1.2
        pts = [
            (x0,0,z0), (x0+w,0,z0), (x0+w,0,z0+d), (x0,0,z0+d),
            (x0,h,z0), (x0+w,h,z0), (x0+w,h,z0+d), (x0,h,z0+d)
        ]
        mid_x = x0 + w/2
        pts.append((mid_x, h+roof_h, z0))
        pts.append((mid_x, h+roof_h, z0+d))

        wall_faces = [
            ([0,1,5,4], t["wall"], 0), ([3,2,6,7], t["wall"], 1),
            ([0,3,7,4], t["wall"], 2), ([1,2,6,5], t["wall"], 3)
        ]
        for idxs, col, face_id in wall_faces:
            sp = [project(p[0], p[1], p[2]) for p in [pts[i] for i in idxs]]
            if all(sp):
                faces.append((get_depth([pts[i] for i in idxs]), col, sp))
                if face_id == 0: add_windows(x0, z0, w, h, 0, 'z', t["window"], faces)
                elif face_id == 1: add_windows(x0, z0+d, w, h, d, 'z', t["window"], faces)
                elif face_id == 2: add_windows(z0, x0, d, h, 0, 'x', t["window"], faces)
                elif face_id == 3: add_windows(z0, x0+w, d, h, w, 'x', t["window"], faces)

        roof_faces = [([4,5,8], t["roof"]), ([7,6,9], t["roof"]), ([5,6,9,8], t["roof"]), ([4,7,9,8], t["roof"])]
        for idxs, col in roof_faces:
            sp = [project(p[0], p[1], p[2]) for p in [pts[i] for i in idxs]]
            if all(sp):
                faces.append((get_depth([pts[i] for i in idxs]), col, sp))

    # 树
    for t in trees:
        x, z = t["x"], t["z"]; h = t["height"]; tr = 0.3
        trunk_pts = [
            (x-tr,0,z-tr), (x+tr,0,z-tr), (x+tr,0,z+tr), (x-tr,0,z+tr),
            (x-tr,h*0.5,z-tr), (x+tr,h*0.5,z-tr), (x+tr,h*0.5,z+tr), (x-tr,h*0.5,z+tr)
        ]
        add_box_faces(trunk_pts, (120, 80, 40), faces)
        crown_r = 2.0; crown_y = h*0.5
        for angle in [0, math.pi/3, 2*math.pi/3]:
            p1 = (x, crown_y + crown_r*0.6, z)
            p2 = (x + crown_r*math.cos(angle), crown_y, z + crown_r*math.sin(angle))
            p3 = (x + crown_r*math.cos(angle+0.5), crown_y, z + crown_r*math.sin(angle+0.5))
            sp = [project(p[0], p[1], p[2]) for p in [p1,p2,p3]]
            if all(sp): faces.append((get_depth([p1,p2,p3]), (50, 150, 50), sp))

    # 车
    for car in cars:
        x, z = car["x"], car["z"]; w, d, h = car["w"], car["d"], car["h"]
        pts = [
            (x-w/2,0,z-d/2), (x+w/2,0,z-d/2), (x+w/2,0,z+d/2), (x-w/2,0,z+d/2),
            (x-w/2,h,z-d/2), (x+w/2,h,z-d/2), (x+w/2,h,z+d/2), (x-w/2,h,z+d/2)
        ]
        add_box_faces(pts, car["color"], faces)

    return faces

# ========== 主循环 ==========
def main():
    global cam_x, cam_z, yaw, pitch
    running = True
    # 启用鼠标抓取以获取相对移动（隐藏鼠标）
    pygame.mouse.set_visible(False)
    pygame.event.set_grab(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.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    running = False

        # 鼠标视角
        dx, dy = pygame.mouse.get_rel()
        yaw += dx * MOUSE_SENS
        pitch -= dy * MOUSE_SENS
        pitch = max(-1.5, min(1.5, pitch))

        # 键盘移动（上下左右 / WASD）
        keys = pygame.key.get_pressed()
        forward = 0.0
        strafe = 0.0

        if keys[pygame.K_UP] or keys[pygame.K_w]:
            forward = 1.0
        if keys[pygame.K_DOWN] or keys[pygame.K_s]:
            forward = -1.0
        if keys[pygame.K_LEFT] or keys[pygame.K_a]:
            strafe = -1.0
        if keys[pygame.K_RIGHT] or keys[pygame.K_d]:
            strafe = 1.0

        # 移动
        move_dir_x = math.cos(yaw) * forward + math.cos(yaw + math.pi/2) * strafe
        move_dir_z = math.sin(yaw) * forward + math.sin(yaw + math.pi/2) * strafe
        length = math.hypot(move_dir_x, move_dir_z)
        if length > 0:
            move_dir_x /= length
            move_dir_z /= length
            cam_x += move_dir_x * MOVE_SPEED * dt
            cam_z += move_dir_z * MOVE_SPEED * dt

        cam_x = max(1, min(GRID_SIZE - 1, cam_x))
        cam_z = max(1, min(GRID_SIZE - 1, cam_z))

        # 绘制
        for y in range(HEIGHT):
            ratio = y / HEIGHT
            r = int(100 + 60 * ratio); g = int(150 + 60 * ratio); b = int(220 + 30 * ratio)
            pygame.draw.line(screen, (r, g, b), (0, y), (WIDTH, y))

        # 网格线
        step = 5
        for i in range(0, GRID_SIZE+1, step):
            for j in range(0, GRID_SIZE, step):
                p1 = project(i, 0, j); p2 = project(i, 0, j+step)
                if p1 and p2: pygame.draw.line(screen, (180,180,180), p1, p2, 1)
        for j in range(0, GRID_SIZE+1, step):
            for i in range(0, GRID_SIZE, step):
                p1 = project(i, 0, j); p2 = project(i+step, 0, j)
                if p1 and p2: pygame.draw.line(screen, (180,180,180), p1, p2, 1)

        # 3D面
        faces = collect_all_faces()
        faces.sort(key=lambda x: x[0], reverse=True)
        for depth, color, sp in faces:
            if len(sp) == 3: pygame.draw.polygon(screen, color, sp)
            else: pygame.draw.polygon(screen, color, sp)
            pygame.draw.polygon(screen, (0,0,0), sp, 1)

        # 信息
        info = f"Pos: ({cam_x:.1f}, {cam_z:.1f})  Yaw: {math.degrees(yaw):.0f}°"
        text_surf = font_med.render(info, True, (255,255,255))
        screen.blit(text_surf, (20, 20))
        help_surf = font_small.render("↑↓前进后退  ←→平移  鼠标旋转视角  ESC退出", True, (255,255,255))
        screen.blit(help_surf, (20, HEIGHT - 30))

        pygame.display.flip()

    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()