import pygame
import random
import math

# ========== 全局配置 ==========
WIDTH, HEIGHT = 1200, 800
FPS = 60
BLOCK_SIZE = 24
GRID_MAX = 320
FOOD_NUM = 8

# 相机控制
rot_x = 30
rot_y = 40
cam_scale = 0.8
mouse_drag = False
last_mouse = (0, 0)

# 色彩
BG = (8, 12, 22)
GRID_LINE = (40, 60, 100)
SNAKE_GREEN_BASE = (0, 220, 90)
SNAKE_RED_BASE = (220, 40, 70)
FOOD_BASE = (255, 210, 0)
FOOD_GLOW = (255, 230, 100)

# 粒子类（死亡爆炸）
class Particle:
    def __init__(self, x, y, color):
        self.x = x
        self.y = y
        self.vx = random.uniform(-4, 4)
        self.vy = random.uniform(-6, -1)
        self.life = 100
        self.color = color

    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.vy += 0.15
        self.life -= 1

# 3D坐标工具：旋转+透视投影
def rotate_point(x, z):
    rad_y = math.radians(rot_y)
    x1 = x * math.cos(rad_y) - z * math.sin(rad_y)
    z1 = x * math.sin(rad_y) + z * math.cos(rad_y)
    rad_x = math.radians(rot_x)
    y1 = z1 * math.sin(rad_x)
    z2 = z1 * math.cos(rad_x)
    return x1, y1, z2

def project(x, y, z):
    z_offset = z + 500
    scale = 500 / z_offset * cam_scale
    sx = WIDTH // 2 + x * scale
    sy = HEIGHT // 2 + y * scale
    return int(sx), int(sy), scale

# 绘制立体方块
def draw_3d_cube(world_x, world_z, color, size):
    rx, ry, rz = rotate_point(world_x, world_z)
    half = size / 2
    points_3d = [
        (rx-half, ry-half, rz-half), (rx+half, ry-half, rz-half),
        (rx+half, ry+half, rz-half), (rx-half, ry+half, rz-half),
        (rx-half, ry-half, rz+half), (rx+half, ry-half, rz+half),
        (rx+half, ry+half, rz+half), (rx-half, ry+half, rz+half),
    ]
    pts2d = [project(p[0], p[1], p[2])[:2] for p in points_3d]
    faces = [
        [0,1,2,3], [4,5,6,7],
        [0,1,5,4], [2,3,7,6],
        [0,3,7,4], [1,2,6,5]
    ]
    face_info = []
    for f_idx in faces:
        avg_z = sum(points_3d[i][2] for i in f_idx) / 4
        face_info.append((avg_z, f_idx))
    face_info.sort(reverse=True)
    for z_val, f in face_info:
        poly = [pts2d[i] for i in f]
        pygame.draw.polygon(screen, color, poly)
        pygame.draw.polygon(screen, (255,255,255), poly, 1)

# 绘制发光食物
def draw_food(x, z):
    rx, ry, rz = rotate_point(x, z)
    sx, sy, s = project(rx, ry, rz)
    for r in range(18, 8, -3):
        alpha = int(120 - r * 6)
        glow_surf = pygame.Surface((r*2, r*2), pygame.SRCALPHA)
        pygame.draw.circle(glow_surf, (*FOOD_GLOW, alpha), (r, r), r)
        screen.blit(glow_surf, (sx - r, sy - r))
    draw_3d_cube(x, z, FOOD_BASE, BLOCK_SIZE - 4)

# AI贪吃蛇类
class AISnake:
    def __init__(self, start_x, start_z, base_color, snake_id):
        self.id = snake_id
        self.base_r, self.base_g, self.base_b = base_color
        self.body = [(start_x, start_z),
                     (start_x - BLOCK_SIZE, start_z),
                     (start_x - BLOCK_SIZE*2, start_z)]
        self.dir = (BLOCK_SIZE, 0)
        self.next_dir = self.dir
        self.alive = True
        self.score = 0
        self.particles = []

    def head(self):
        return self.body[0]

    def get_gradient_color(self, seg_idx, total_len):
        ratio = seg_idx / max(total_len, 1)
        dark_rate = 0.4 + 0.6 * (1 - ratio)
        r = int(self.base_r * dark_rate)
        g = int(self.base_g * dark_rate)
        b = int(self.base_b * dark_rate)
        return (r, g, b)

    def ai_choose_dir(self, food_list, enemy_snake):
        hx, hz = self.head()
        target = food_list[0]
        min_dist = math.hypot(target[0]-hx, target[1]-hz)
        for f in food_list:
            d = math.hypot(f[0]-hx, f[1]-hz)
            if d < min_dist:
                min_dist = d
                target = f
        tx, tz = target
        dir_candidates = [
            (BLOCK_SIZE, 0), (-BLOCK_SIZE, 0),
            (0, BLOCK_SIZE), (0, -BLOCK_SIZE)
        ]
        safe_dirs = []
        for dx, dz in dir_candidates:
            nx = hx + dx
            nz = hz + dz
            if abs(nx) > GRID_MAX or abs(nz) > GRID_MAX:
                continue
            new_pos = (nx, nz)
            if new_pos in self.body or new_pos in enemy_snake.body:
                continue
            dist = math.hypot(tx - nx, tz - nz)
            safe_dirs.append((dist, (dx, dz)))
        if safe_dirs:
            safe_dirs.sort()
            self.next_dir = safe_dirs[0][1]
        else:
            self.next_dir = self.dir

    def move_step(self, foods, enemy):
        if not self.alive:
            return
        self.dir = self.next_dir
        hx, hz = self.head()
        nhx = hx + self.dir[0]
        nhz = hz + self.dir[1]
        new_head = (nhx, nhz)
        if abs(nhx) > GRID_MAX or abs(nhz) > GRID_MAX:
            self.die()
            return
        if new_head in self.body or new_head in enemy.body:
            self.die()
            return
        grow = False
        if new_head in foods:
            foods.remove(new_head)
            self.score += 10
            grow = True
        self.body.insert(0, new_head)
        if not grow:
            self.body.pop()

    def die(self):
        self.alive = False
        for (x, z) in self.body:
            rx, ry, rz = rotate_point(x, z)
            sx, sy, _ = project(rx, ry, rz)
            for _ in range(8):
                self.particles.append(Particle(sx, sy, (self.base_r, self.base_g, self.base_b)))

    def update_particles(self):
        for p in self.particles:
            p.update()
        self.particles = [p for p in self.particles if p.life > 0]

    def draw_snake(self):
        length = len(self.body)
        for idx, (x, z) in enumerate(self.body):
            color = self.get_gradient_color(idx, length)
            draw_3d_cube(x, z, color, BLOCK_SIZE - 2)
        for p in self.particles:
            alpha = p.life / 100
            r, g, b = p.color
            color_alpha = (r, g, b, int(alpha * 255))
            surf = pygame.Surface((6, 6), pygame.SRCALPHA)
            pygame.draw.rect(surf, color_alpha, (0,0,6,6))
            screen.blit(surf, (p.x, p.y))

# 生成食物
def spawn_food(s1, s2):
    while True:
        x = random.randrange(-GRID_MAX, GRID_MAX, BLOCK_SIZE)
        z = random.randrange(-GRID_MAX, GRID_MAX, BLOCK_SIZE)
        pos = (x, z)
        if pos not in s1.body and pos not in s2.body:
            return pos

# 地面网格
def draw_ground_grid():
    step = BLOCK_SIZE * 4
    r = GRID_MAX
    x = -r
    while x <= r:
        draw_3d_cube(x, 0, GRID_LINE, 3)
        x += step
    z = -r
    while z <= r:
        draw_3d_cube(0, z, GRID_LINE, 3)
        z += step

# ========== 主程序 ==========
pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("3D AI对战贪吃蛇")
clock = pygame.time.Clock()

# 兼容字体
try:
    font = pygame.font.SysFont(None, 32)
except Exception:
    font = pygame.font.Font(None, 32)

# 创建两条AI蛇
snake_green = AISnake(-100, 0, SNAKE_GREEN_BASE, 1)
snake_red = AISnake(100, 0, SNAKE_RED_BASE, 2)
foods = [spawn_food(snake_green, snake_red) for _ in range(FOOD_NUM)]

running = True
while running:
    clock.tick(FPS)
    screen.fill(BG)

    # 事件
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            mouse_drag = True
            last_mouse = event.pos
        if event.type == pygame.MOUSEBUTTONUP and event.button == 1:
            mouse_drag = False
        if event.type == pygame.MOUSEMOTION and mouse_drag:
            dx = event.pos[0] - last_mouse[0]
            dy = event.pos[1] - last_mouse[1]
            rot_y += dx * 0.4
            rot_x += dy * 0.4
            rot_x = max(-70, min(70, rot_x))
            last_mouse = event.pos
        if event.type == pygame.MOUSEWHEEL:
            cam_scale += event.y * 0.05
            cam_scale = max(0.4, min(1.5, cam_scale))

    # 【核心修复：对手参数传递正确】
    snake_green.ai_choose_dir(foods, snake_red)
    snake_red.ai_choose_dir(foods, snake_green)
    snake_green.move_step(foods, snake_red)
    snake_red.move_step(foods, snake_green)

    # 补充食物
    while len(foods) < FOOD_NUM:
        foods.append(spawn_food(snake_green, snake_red))

    # 更新粒子
    snake_green.update_particles()
    snake_red.update_particles()

    # 渲染
    draw_ground_grid()
    for fx, fz in foods:
        draw_food(fx, fz)
    snake_green.draw_snake()
    snake_red.draw_snake()

    # UI文字
    text_green = font.render(f"Green AI Score: {snake_green.score}", True, SNAKE_GREEN_BASE)
    text_red = font.render(f"Red AI Score: {snake_red.score}", True, SNAKE_RED_BASE)
    state_text = ""
    if not snake_green.alive and not snake_red.alive:
        state_text = "Both Snakes Dead"
    elif not snake_green.alive:
        state_text = "Red Snake Win!"
    elif not snake_red.alive:
        state_text = "Green Snake Win!"
    text_state = font.render(state_text, True, (255,255,255))

    screen.blit(text_green, (10, 10))
    screen.blit(text_red, (10, 45))
    screen.blit(text_state, (WIDTH//2 - 120, 10))

    pygame.display.flip()

pygame.quit()