from ursina import *
import random
import math
import sys

app = Ursina()

# ========== 游戏设置 ==========
GRID_SIZE = 20
CELL_SIZE = 1
MOVE_INTERVAL = 0.15
SPEEDUP_FACTOR = 1.0
score = 0
game_over = False
move_timer = 0

# 存储延迟任务以便取消
pending_invokes = []

# ========== 场景 ==========
window.title = "3D Neon Snake"
window.borderless = False
window.size = (800, 600)
window.color = color.rgb(5, 5, 15)

# 星空
for _ in range(80):
    star = Entity(model='quad', color=color.rgb(random.randint(150,255), random.randint(150,255), random.randint(200,255)),
                  scale=0.05, position=(random.uniform(-15,15), random.uniform(5,15), random.uniform(-15,15)), billboard=True)
    star.animate_color(color.rgb(random.randint(50,150), random.randint(50,150), random.randint(100,200)), duration=random.uniform(1,3), loop=True)

DirectionalLight(parent=camera, rotation=(45, -30, 0), color=color.rgba(200,200,255,150))
AmbientLight(color=color.rgba(50,50,80,80))

# 平台
platform = Entity(model='plane', scale=(GRID_SIZE+1, 1, GRID_SIZE+1), color=color.rgba(10, 10, 30, 200),
                  position=(0, -0.05, 0), texture='white_cube', texture_scale=(GRID_SIZE+1, GRID_SIZE+1))

# 网格线
for i in range(GRID_SIZE+1):
    Entity(model='cube', scale=(GRID_SIZE+1, 0.02, 0.05), color=color.rgba(0, 255, 255, 100), position=(0, 0.01, i - GRID_SIZE/2))
    Entity(model='cube', scale=(0.05, 0.02, GRID_SIZE+1), color=color.rgba(0, 255, 255, 100), position=(i - GRID_SIZE/2, 0.01, 0))

# 边框灯
for side in range(4):
    for j in range(GRID_SIZE+1):
        if side == 0: pos = (-GRID_SIZE/2, 0.02, j - GRID_SIZE/2)
        elif side == 1: pos = (GRID_SIZE/2, 0.02, j - GRID_SIZE/2)
        elif side == 2: pos = (j - GRID_SIZE/2, 0.02, -GRID_SIZE/2)
        else: pos = (j - GRID_SIZE/2, 0.02, GRID_SIZE/2)
        Entity(model='sphere', scale=0.15, color=color.cyan, position=pos, unlit=True)

# ========== 蛇 ==========
class Snake:
    def __init__(self):
        self.segments = []
        self.direction = Vec3(0, 0, 1)
        self.next_direction = Vec3(0, 0, 1)
        self.grow_count = 0
        self.invincible = False
        self.head = None
        self.eye_l = None
        self.eye_r = None
        self.pupil_l = None
        self.pupil_r = None
        # 初始身体
        for i in range(3):
            seg = Entity(model='sphere', scale=0.45,
                         color=color.rgb(0, 255 - i*30, 255 - i*30),
                         position=(0, 0.3, -3 + i), unlit=True)
            self.segments.append(seg)
        self.head = self.segments[0]
        self.head.scale = 0.5
        self.eye_l = Entity(model='sphere', scale=0.1, color=color.white, parent=self.head, position=(-0.15, 0.1, 0.2))
        self.eye_r = Entity(model='sphere', scale=0.1, color=color.white, parent=self.head, position=(0.15, 0.1, 0.2))
        self.pupil_l = Entity(model='sphere', scale=0.06, color=color.black, parent=self.eye_l, position=(0,0,0.05))
        self.pupil_r = Entity(model='sphere', scale=0.06, color=color.black, parent=self.eye_r, position=(0,0,0.05))

    def move(self):
        global game_over
        if game_over: return

        if self.next_direction + self.direction != Vec3(0,0,0):
            self.direction = self.next_direction

        head_pos = self.head.position + self.direction * CELL_SIZE * SPEEDUP_FACTOR

        # 边界检测
        if not self.invincible:
            if abs(head_pos.x) > GRID_SIZE/2 - 0.5 or abs(head_pos.z) > GRID_SIZE/2 - 0.5:
                self.die()
                return

        # 自身碰撞（跳过头部自身，并排除刚移走的位置）
        for seg in self.segments[2:]:  # 跳过第一个（头部）和紧挨着的第二节，因为第二节还没移动
            if distance(head_pos, seg.position) < 0.5:
                self.die()
                return

        # 食物碰撞
        for food in foods[:]:
            if distance(head_pos, food.position) < 0.6:
                global score
                score += 10
                self.grow_count += 1
                for _ in range(12):
                    Particle(food.position, food.color)
                foods.remove(food)
                destroy(food)
                spawn_food()
                if random.random() < 0.3:
                    spawn_powerup()
                break

        # 道具碰撞
        for pu in powerups[:]:
            if distance(head_pos, pu.position) < 0.6:
                apply_powerup(pu.power_type)
                powerups.remove(pu)
                destroy(pu)
                break

        # 创建新头部
        new_head = Entity(model='sphere', scale=0.5, color=color.rgb(0,255,255), position=head_pos, unlit=True)
        self.eye_l.parent = new_head
        self.eye_r.parent = new_head
        self.eye_l.position = (-0.15, 0.1, 0.2)
        self.eye_r.position = (0.15, 0.1, 0.2)
        self.head = new_head
        self.segments.insert(0, new_head)

        if self.grow_count > 0:
            self.grow_count -= 1
        else:
            tail = self.segments.pop()
            destroy(tail)
        self.update_colors()

    def update_colors(self):
        total = len(self.segments)
        for i, seg in enumerate(self.segments):
            t = i / max(1, total-1)
            r = 0
            g = int(255 * (1 - t))
            b = int(255 * (1 - t*0.5) + 100 * t)
            seg.color = color.rgb(r, g, b)

    def set_direction(self, key):
        if key in ('w', 'up arrow') and self.direction != Vec3(0,0,-1):
            self.next_direction = Vec3(0,0,1)
        elif key in ('s', 'down arrow') and self.direction != Vec3(0,0,1):
            self.next_direction = Vec3(0,0,-1)
        elif key in ('a', 'left arrow') and self.direction != Vec3(1,0,0):
            self.next_direction = Vec3(-1,0,0)
        elif key in ('d', 'right arrow') and self.direction != Vec3(-1,0,0):
            self.next_direction = Vec3(1,0,0)

    def die(self):
        global game_over
        if game_over:  # 防止重复触发
            return
        game_over = True
        for seg in self.segments:
            seg.color = color.red
        # 延迟显示结束文字，避免多次创建
        invoke(lambda: show_game_over_text(), delay=0.1)

# ========== 食物 ==========
class Food(Entity):
    def __init__(self, position):
        super().__init__(model='sphere', scale=0.35,
                         color=random.choice([color.rgb(255,100,100), color.gold, color.magenta]),
                         position=position, unlit=True)
        Entity(model='sphere', scale=0.5, color=color.rgba(255,255,255,80), parent=self, unlit=True)

# ========== 道具 ==========
class PowerUp(Entity):
    def __init__(self, position, power_type):
        colors = {'speed': color.green, 'invincible': color.orange}
        super().__init__(model='diamond', scale=0.4, color=colors.get(power_type, color.white),
                         position=position, unlit=True)
        self.power_type = power_type
        self.animate_rotation_y(360, duration=2, loop=True)

# ========== 粒子 ==========
class Particle(Entity):
    def __init__(self, pos, col):
        super().__init__(model='quad', color=col, scale=0.08, position=pos, billboard=True, unlit=True)
        self.velocity = Vec3(random.uniform(-0.5,0.5), random.uniform(0.5,1.5), random.uniform(-0.5,0.5))
        self.life = 0.4
    def update(self):
        self.position += self.velocity * time.dt
        self.velocity.y -= 4 * time.dt
        self.life -= time.dt
        if self.life <= 0: destroy(self)

# ========== 全局容器 ==========
foods = []
powerups = []
snake = Snake()
game_over_text_entity = None  # 确保只有一个结束文字

def show_game_over_text():
    global game_over_text_entity
    if game_over_text_entity:
        destroy(game_over_text_entity)
    game_over_text_entity = Text("GAME OVER\nPress R to restart", position=(0,0.1), scale=3,
                                 color=color.red, origin=(0,0), font='VeraMono.ttf')

# ========== 生成函数 ==========
def spawn_food():
    while True:
        x = random.randint(-GRID_SIZE//2+1, GRID_SIZE//2-1)
        z = random.randint(-GRID_SIZE//2+1, GRID_SIZE//2-1)
        pos = Vec3(x, 0.25, z)
        if not any(distance(pos, seg.position) < 0.8 for seg in snake.segments):
            break
    f = Food(pos)
    foods.append(f)

def spawn_powerup():
    if len(powerups) >= 2: return
    while True:
        x = random.randint(-GRID_SIZE//2+1, GRID_SIZE//2-1)
        z = random.randint(-GRID_SIZE//2+1, GRID_SIZE//2-1)
        pos = Vec3(x, 0.3, z)
        if not any(distance(pos, seg.position) < 0.8 for seg in snake.segments) and \
           not any(distance(pos, f.position) < 0.8 for f in foods):
            break
    ptype = random.choice(['speed', 'invincible'])
    pu = PowerUp(pos, ptype)
    powerups.append(pu)

def apply_powerup(ptype):
    global SPEEDUP_FACTOR
    if ptype == 'speed':
        SPEEDUP_FACTOR = 1.5
        # 取消之前可能遗留的invoke
        for inv in pending_invokes:
            try: inv.cancel()
            except: pass
        inv = invoke(lambda: reset_speed(), delay=4)
        pending_invokes.append(inv)
    elif ptype == 'invincible':
        snake.invincible = True
        for seg in snake.segments:
            seg.color = color.white
        inv = invoke(lambda: disable_invincible(), delay=4)
        pending_invokes.append(inv)

def reset_speed():
    global SPEEDUP_FACTOR
    SPEEDUP_FACTOR = 1.0

def disable_invincible():
    snake.invincible = False
    if not game_over:
        snake.update_colors()

# ========== 重置游戏 ==========
def reset_game():
    global game_over, score, SPEEDUP_FACTOR, move_timer, game_over_text_entity
    # 取消所有延迟调用
    for inv in pending_invokes:
        try: inv.cancel()
        except: pass
    pending_invokes.clear()

    game_over = False
    score = 0
    SPEEDUP_FACTOR = 1.0
    move_timer = 0

    # 销毁蛇身
    for seg in snake.segments:
        destroy(seg)
    snake.segments.clear()

    # 销毁食物道具
    for f in foods: destroy(f)
    foods.clear()
    for pu in powerups: destroy(pu)
    powerups.clear()

    # 销毁结束文字
    if game_over_text_entity:
        destroy(game_over_text_entity)
        game_over_text_entity = None

    # 重新创建蛇
    snake.__init__()
    # 确保蛇的初始位置不与新生成的食物重叠
    for _ in range(3): spawn_food()

# ========== 输入 ==========
def input(key):
    if key == 'r' and game_over:
        reset_game()
        return
    if game_over: return
    snake.set_direction(key)

# ========== 主循环 ==========
def update():
    global move_timer, game_over
    if game_over: return

    move_timer += time.dt
    if move_timer >= MOVE_INTERVAL / SPEEDUP_FACTOR:
        snake.move()
        move_timer = 0

    if not hasattr(update, 'hud'):
        update.hud = Text(f"Score: {score}", position=(-0.85, 0.45), scale=1.5,
                          background=True, color=color.white)
    else:
        update.hud.text = f"Score: {score}"

# ========== 初始生成 ==========
for _ in range(3): spawn_food()

# 摄像机
camera.position = (0, 18, -18)
camera.rotation_x = 45
camera.look_at((0, 0, 0))

app.run()