import pygame
import random
import sys
import math

# ========== 初始化 ==========
pygame.init()
WIDTH, HEIGHT = 800, 500
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("🚇 Subway Runner – Beautiful Edition")
clock = pygame.time.Clock()

# ========== 色彩方案（夜景都市） ==========
BG_TOP = (10, 10, 30)
BG_BOTTOM = (60, 50, 100)
STAR = (255, 255, 220)
GROUND = (40, 40, 50)
RAIL = (130, 130, 140)
RAIL_SHADOW = (30, 30, 35)
GOLD = (255, 215, 0)
RED = (255, 70, 70)
CYAN = (0, 255, 255)
GREEN = (80, 255, 100)
WHITE = (255, 255, 255)
BLACK = (15, 15, 25)
PLAYER_BODY = (80, 170, 255)
PLAYER_SKIN = (255, 210, 170)
HAT_RED = (240, 80, 60)
TRAIN_COLOR = (180, 50, 40)
BARRIER_COLOR = (220, 200, 40)

# ========== 字体 ==========
font_small = pygame.font.Font(None, 28)
font_med = pygame.font.Font(None, 40)
font_large = pygame.font.Font(None, 64)

# ========== 三条跑道 ==========
LANE_Y = [240, 310, 380]
PLAYER_X = 150

# ========== 玩家类 ==========
class Player:
    def __init__(self):
        self.x = PLAYER_X
        self.y = LANE_Y[1]
        self.lane = 1
        self.target_y = self.y
        self.vy = 0
        self.gravity = 1.2
        self.jump_power = -18
        self.on_ground = True
        self.sliding = False
        self.slide_timer = 0
        self.slide_duration = 30
        self.double_jump = False
        self.shield = False
        self.magnet = False
        self.double_coin = False
        self.anim_frame = 0
        self.invincible = 0
        self.moving_lane = False
        self.trail = []

    def move_to_lane(self, lane):
        if not self.moving_lane and 0 <= lane <= 2:
            self.lane = lane
            self.target_y = LANE_Y[lane]
            self.moving_lane = True

    def jump(self):
        if self.on_ground:
            self.vy = self.jump_power
            self.on_ground = False
        elif not self.double_jump:
            self.vy = self.jump_power * 0.8
            self.double_jump = True

    def slide(self):
        if self.on_ground and not self.sliding:
            self.sliding = True
            self.slide_timer = self.slide_duration

    def update(self):
        if self.moving_lane:
            if abs(self.y - self.target_y) < 2:
                self.y = self.target_y
                self.moving_lane = False
            else:
                self.y += (self.target_y - self.y) * 0.3

        self.vy += self.gravity
        self.y += self.vy
        if self.y >= LANE_Y[self.lane]:
            self.y = LANE_Y[self.lane]
            self.vy = 0
            self.on_ground = True
            self.double_jump = False
        else:
            self.on_ground = False

        if self.sliding:
            self.slide_timer -= 1
            if self.slide_timer <= 0:
                self.sliding = False

        self.anim_frame += 0.2
        if self.invincible > 0:
            self.invincible -= 1

        if self.on_ground and not self.sliding:
            self.trail.append((self.x, self.y, self.anim_frame))
            if len(self.trail) > 5:
                self.trail.pop(0)
        else:
            self.trail.clear()

    def draw(self, surface):
        for i, (tx, ty, frame) in enumerate(self.trail):
            alpha = 50 + i * 30
            color = (PLAYER_BODY[0], PLAYER_BODY[1], PLAYER_BODY[2], alpha)
            pygame.draw.circle(surface, color, (int(tx-5), int(ty-20)), 5)

        if self.shield:
            shield_r = 34
            for i in range(shield_r, 0, -6):
                alpha = 100 - i*2
                pygame.draw.circle(surface, (*CYAN, alpha), (self.x, self.y-20), i)
            pygame.draw.circle(surface, CYAN, (self.x, self.y-20), shield_r, 4)

        if self.invincible > 0 and self.invincible % 6 < 3:
            return

        if self.sliding:
            pygame.draw.rect(surface, PLAYER_BODY, (self.x-20, self.y-5, 40, 20))
            pygame.draw.rect(surface, BLACK, (self.x-20, self.y-5, 40, 20), 2)
            pygame.draw.circle(surface, PLAYER_SKIN, (self.x, self.y-12), 12)
            pygame.draw.circle(surface, BLACK, (self.x, self.y-12), 12, 2)
            pygame.draw.rect(surface, CYAN, (self.x-7, self.y-15, 14, 6), border_radius=3)
        else:
            leg_offset = math.sin(self.anim_frame * 0.5) * 6 if self.on_ground else 0
            arm_angle = math.sin(self.anim_frame * 0.5) * 15 if self.on_ground else 0

            pygame.draw.ellipse(surface, (0,0,0,100), (self.x-15, self.y+8, 30, 8))

            pygame.draw.rect(surface, PLAYER_BODY, (self.x-15, self.y-40, 30, 40))
            pygame.draw.rect(surface, BLACK, (self.x-15, self.y-40, 30, 40), 2)

            pygame.draw.line(surface, PLAYER_BODY, (self.x-15, self.y-30), (self.x-30, self.y-10+arm_angle), 7)
            pygame.draw.line(surface, PLAYER_BODY, (self.x+15, self.y-30), (self.x+30, self.y-10-arm_angle), 7)

            pygame.draw.circle(surface, PLAYER_SKIN, (self.x, self.y-50), 16)
            pygame.draw.circle(surface, BLACK, (self.x, self.y-50), 16, 2)

            pygame.draw.circle(surface, BLACK, (self.x-5, self.y-53), 3)
            pygame.draw.circle(surface, BLACK, (self.x+5, self.y-53), 3)
            pygame.draw.circle(surface, WHITE, (self.x-6, self.y-54), 1)
            pygame.draw.circle(surface, WHITE, (self.x+4, self.y-54), 1)

            pygame.draw.rect(surface, HAT_RED, (self.x-19, self.y-67, 38, 12), border_radius=6)
            pygame.draw.rect(surface, BLACK, (self.x-19, self.y-67, 38, 12), 2, border_radius=6)
            pygame.draw.rect(surface, HAT_RED, (self.x-22, self.y-63, 44, 5), border_radius=3)

            pygame.draw.line(surface, BLACK, (self.x-8, self.y-5), (self.x-10+leg_offset, self.y+12), 5)
            pygame.draw.line(surface, BLACK, (self.x+8, self.y-5), (self.x+10-leg_offset, self.y+12), 5)

    def get_rect(self):
        if self.sliding:
            return pygame.Rect(self.x-20, self.y-5, 40, 20)
        else:
            return pygame.Rect(self.x-15, self.y-40, 30, 40)

# ========== 障碍物 ==========
class Obstacle:
    def __init__(self, x, lane, type_):
        self.x = x
        self.lane = lane
        self.type = type_
        self.width = 40
        self.passed = False
        self.y = LANE_Y[lane]

    def update(self, speed):
        self.x -= speed
        return self.x < -60

    def draw(self, surface):
        if self.type == "train":
            body_rect = pygame.Rect(self.x, self.y-50, self.width, 85)
            pygame.draw.rect(surface, TRAIN_COLOR, body_rect, border_radius=8)
            pygame.draw.rect(surface, BLACK, body_rect, 3, border_radius=8)
            win_rect = pygame.Rect(self.x+8, self.y-40, 24, 18)
            pygame.draw.rect(surface, CYAN, win_rect, border_radius=4)
            pygame.draw.rect(surface, BLACK, win_rect, 2, border_radius=4)
            headlight_x = self.x + self.width - 10
            pygame.draw.circle(surface, GOLD, (headlight_x, self.y-25), 6)
            pygame.draw.circle(surface, WHITE, (headlight_x, self.y-25), 3)
        elif self.type == "barrier":
            base_rect = pygame.Rect(self.x, self.y-5, self.width, 25)
            pygame.draw.rect(surface, BARRIER_COLOR, base_rect, border_radius=4)
            pygame.draw.rect(surface, BLACK, base_rect, 2, border_radius=4)
            for offset in [6, self.width-14]:
                pole_rect = pygame.Rect(self.x+offset, self.y-15, 6, 35)
                pygame.draw.rect(surface, (160, 150, 40), pole_rect)
                pygame.draw.rect(surface, BLACK, pole_rect, 2)
        elif self.type == "double":
            top_rect = pygame.Rect(self.x, self.y-70, self.width, 65)
            pygame.draw.rect(surface, TRAIN_COLOR, top_rect, border_radius=8)
            pygame.draw.rect(surface, BLACK, top_rect, 3, border_radius=8)
            pygame.draw.rect(surface, CYAN, (self.x+8, self.y-60, 24, 14), border_radius=3)
            pygame.draw.rect(surface, BLACK, (self.x+8, self.y-60, 24, 14), 2, border_radius=3)
            bot_rect = pygame.Rect(self.x, self.y-2, self.width, 22)
            pygame.draw.rect(surface, BARRIER_COLOR, bot_rect, border_radius=4)
            pygame.draw.rect(surface, BLACK, bot_rect, 2, border_radius=4)

    def get_rect(self):
        if self.type == "train":
            return pygame.Rect(self.x, self.y-50, self.width, 85)
        elif self.type == "barrier":
            return pygame.Rect(self.x, self.y-5, self.width, 25)
        elif self.type == "double":
            return pygame.Rect(self.x, self.y-70, self.width, 95)

# ========== 金币 ==========
class Coin:
    def __init__(self, x, lane, value=1):
        self.x = x
        self.y = LANE_Y[lane] - 30
        self.value = value
        self.radius = 10
        self.collected = False
        self.angle = random.uniform(0, 2*math.pi)

    def update(self, speed):
        self.x -= speed
        if self.x < -30:
            return True
        self.angle += 0.15
        return False

    def draw(self, surface):
        if self.collected:
            return
        glow_r = 4 + int(math.sin(self.angle) * 2)
        for r in range(self.radius+glow_r, 0, -3):
            alpha = 50
            pygame.draw.circle(surface, (*GOLD, alpha), (int(self.x), int(self.y)), r)
        pygame.draw.circle(surface, GOLD, (int(self.x), int(self.y)), self.radius)
        pygame.draw.circle(surface, WHITE, (int(self.x-2), int(self.y-3)), self.radius//2)
        coin_text = font_small.render("$", True, BLACK)
        surface.blit(coin_text, (self.x-6, self.y-12))

    def get_rect(self):
        return pygame.Rect(self.x-10, self.y-10, 20, 20)

# ========== 道具 ==========
class PowerUp:
    def __init__(self, x, lane, type_):
        self.x = x
        self.y = LANE_Y[lane] - 35
        self.type = type_
        self.radius = 14
        self.collected = False
        self.float_offset = 0

    def update(self, speed):
        self.x -= speed
        if self.x < -30:
            return True
        self.float_offset = math.sin(pygame.time.get_ticks() * 0.01) * 3
        return False

    def draw(self, surface):
        colors = {"shield": CYAN, "magnet": RED, "double": GOLD}
        letters = {"shield": "S", "magnet": "M", "double": "x2"}
        color = colors[self.type]
        y_draw = self.y + self.float_offset

        for i in range(3):
            angle = pygame.time.get_ticks() * 0.08 + i * 2.1
            dx = math.cos(angle) * 8
            dy = math.sin(angle) * 8
            pygame.draw.circle(surface, color, (int(self.x+dx), int(y_draw+dy)), self.radius-4)

        pygame.draw.circle(surface, color, (int(self.x), int(y_draw)), self.radius)
        pygame.draw.circle(surface, WHITE, (int(self.x), int(y_draw)), self.radius, 2)
        txt = font_small.render(letters[self.type], True, BLACK)
        surface.blit(txt, (self.x-10, y_draw-12))

    def get_rect(self):
        return pygame.Rect(self.x-14, self.y-14, 28, 28)

# ========== 粒子系统 ==========
class Particle:
    def __init__(self, x, y, color, vel, life=25):
        self.x = x
        self.y = y
        self.vx = vel[0]
        self.vy = vel[1]
        self.color = color
        self.life = life
        self.max_life = life

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

    def draw(self, surface):
        alpha = int(255 * (self.life / self.max_life))
        r = int(3 * (self.life / self.max_life))
        pygame.draw.circle(surface, (*self.color, alpha), (int(self.x), int(self.y)), r)

# ========== 游戏主类 ==========
class Game:
    def __init__(self):
        self.reset()

    def reset(self):
        self.player = Player()
        self.obstacles = []
        self.coins = []
        self.powerups = []
        self.particles = []
        self.speed = 6
        self.score = 0
        self.distance = 0
        self.spawn_timer = 0
        self.coin_spawn_timer = 0
        self.powerup_timer = 0
        self.magnet_timer = 0
        self.shield_timer = 0
        self.double_timer = 0
        self.bg_offset = 0
        self.city_buildings = self.generate_buildings()

    def generate_buildings(self):
        buildings = []
        for i in range(0, WIDTH+300, 100):
            h = random.randint(80, 150)
            buildings.append([i, h])
        return buildings

    def spawn_obstacle(self):
        lane = random.randint(0, 2)
        types = ["train", "barrier", "train", "barrier", "double"]
        self.obstacles.append(Obstacle(WIDTH+60, lane, random.choice(types)))

    def spawn_coin_line(self):
        lane = random.randint(0, 2)
        for i in range(5):
            self.coins.append(Coin(WIDTH+120 + i*30, lane))

    def spawn_powerup(self):
        lane = random.randint(0, 2)
        type_ = random.choice(["magnet", "double", "shield"])
        self.powerups.append(PowerUp(WIDTH+60, lane, type_))

    def handle_input(self):
        keys = pygame.key.get_pressed()
        if not self.player.moving_lane:
            if keys[pygame.K_LEFT] or keys[pygame.K_a]:
                self.player.move_to_lane(max(0, self.player.lane - 1))
            if keys[pygame.K_RIGHT] or keys[pygame.K_d]:
                self.player.move_to_lane(min(2, self.player.lane + 1))
        if keys[pygame.K_SPACE] or keys[pygame.K_UP]:
            self.player.jump()
        if keys[pygame.K_DOWN]:
            self.player.slide()

    def update(self):
        if self.state != "playing":
            return

        self.player.update()
        self.speed = 6 + self.distance // 400
        self.bg_offset = (self.bg_offset + self.speed) % 60

        self.spawn_timer += 1
        spawn_rate = max(30, 70 - self.distance // 200)
        if self.spawn_timer > spawn_rate:
            self.spawn_timer = 0
            self.spawn_obstacle()

        self.coin_spawn_timer += 1
        if self.coin_spawn_timer > 50:
            self.coin_spawn_timer = 0
            self.spawn_coin_line()

        self.powerup_timer += 1
        if self.powerup_timer > 250:
            self.powerup_timer = 0
            self.spawn_powerup()

        if self.player.magnet:
            self.magnet_timer -= 1
            if self.magnet_timer <= 0:
                self.player.magnet = False
        if self.player.shield:
            self.shield_timer -= 1
            if self.shield_timer <= 0:
                self.player.shield = False
        if self.player.double_coin:
            self.double_timer -= 1
            if self.double_timer <= 0:
                self.player.double_coin = False

        for obs in self.obstacles[:]:
            if obs.update(self.speed):
                self.obstacles.remove(obs)
                continue
            if not obs.passed and obs.x + obs.width < self.player.x:
                obs.passed = True
                self.distance += 50
            if obs.lane == self.player.lane and not self.player.moving_lane:
                if self.player.get_rect().colliderect(obs.get_rect()):
                    if self.player.shield:
                        self.player.shield = False
                        self.player.invincible = 30
                        for _ in range(30):
                            self.particles.append(Particle(self.player.x, self.player.y-20, CYAN,
                                                           (random.uniform(-5,5), random.uniform(-8,-2))))
                        self.obstacles.remove(obs)
                    else:
                        self.game_over()
                        return

        for coin in self.coins[:]:
            if coin.update(self.speed):
                self.coins.remove(coin)
                continue
            if self.player.magnet and abs(coin.x - self.player.x) < 140:
                coin.x += (self.player.x - coin.x) * 0.15
            if not coin.collected and coin.get_rect().colliderect(self.player.get_rect()):
                coin.collected = True
                value = coin.value * (2 if self.player.double_coin else 1)
                self.score += value
                for _ in range(8):
                    self.particles.append(Particle(coin.x, coin.y, GOLD,
                                                   (random.uniform(-3,3), random.uniform(-5,-1))))
                self.coins.remove(coin)

        for pu in self.powerups[:]:
            if pu.update(self.speed):
                self.powerups.remove(pu)
                continue
            if pu.get_rect().colliderect(self.player.get_rect()):
                if pu.type == "shield":
                    self.player.shield = True
                    self.shield_timer = 360
                elif pu.type == "magnet":
                    self.player.magnet = True
                    self.magnet_timer = 360
                elif pu.type == "double":
                    self.player.double_coin = True
                    self.double_timer = 360
                for _ in range(15):
                    self.particles.append(Particle(pu.x, pu.y, CYAN,
                                                   (random.uniform(-4,4), random.uniform(-6,-2))))
                self.powerups.remove(pu)

        for p in self.particles[:]:
            p.update()
            if p.life <= 0:
                self.particles.remove(p)

        for b in self.city_buildings:
            b[0] -= self.speed * 0.3
            if b[0] < -100:
                b[0] += WIDTH + 300
                b[1] = random.randint(80, 150)

        self.distance += 1

    def game_over(self):
        self.state = "gameover"

    def draw_background(self):
        for y in range(0, 300):
            ratio = y / 300
            r = int(10 + 50 * ratio)
            g = int(10 + 40 * ratio)
            b = int(30 + 70 * ratio)
            pygame.draw.line(screen, (r, g, b), (0, y), (WIDTH, y))

        for i in range(0, WIDTH, 40):
            x = (i + self.bg_offset * 0.5) % WIDTH
            pygame.draw.circle(screen, STAR, (x, 60), 2)
            pygame.draw.circle(screen, STAR, (x + 15, 100), 1)

        for b in self.city_buildings:
            x = int(b[0])           # ★ 转成整数
            h = b[1]
            pygame.draw.rect(screen, (25, 20, 40), (x, 300 - h, 80, h))
            for wy in range(300 - h + 10, 295, 20):
                for wx in range(x + 10, x + 70, 25):
                    if random.random() < 0.6:
                        light_color = (255, 240, 180) if random.random() < 0.7 else (100, 100, 120)
                        pygame.draw.rect(screen, light_color, (wx, wy, 12, 10))

        pygame.draw.rect(screen, GROUND, (0, 320, WIDTH, 180))

        for i in range(0, WIDTH + 30, 40):
            x = (i - self.bg_offset) % (WIDTH + 30) - 15
            for ly in LANE_Y:
                pygame.draw.line(screen, RAIL_SHADOW, (x, ly + 14), (x + 20, ly + 14), 2)
                pygame.draw.line(screen, RAIL, (x, ly + 12), (x + 20, ly + 12), 2)

        pygame.draw.rect(screen, (25, 25, 35), (0, 0, WIDTH, 30))
        pygame.draw.rect(screen, (80, 80, 90), (0, 0, WIDTH, 30), 3)

    def draw_ui(self):
        panel = pygame.Surface((190, 80), pygame.SRCALPHA)
        panel.fill((0, 0, 0, 140))
        screen.blit(panel, (15, 40))
        score_txt = font_med.render(f"Score: {self.score}", True, WHITE)
        screen.blit(score_txt, (30, 48))
        dist_txt = font_small.render(f"Dist: {self.distance}m", True, WHITE)
        screen.blit(dist_txt, (30, 78))

        icon_x = 230
        if self.player.shield:
            pygame.draw.circle(screen, CYAN, (icon_x, 50), 13)
            screen.blit(font_small.render("S", True, BLACK), (icon_x-6, 40))
            icon_x += 40
        if self.player.magnet:
            pygame.draw.circle(screen, RED, (icon_x, 50), 13)
            screen.blit(font_small.render("M", True, BLACK), (icon_x-6, 40))
            icon_x += 40
        if self.player.double_coin:
            pygame.draw.circle(screen, GOLD, (icon_x, 50), 13)
            screen.blit(font_small.render("x2", True, BLACK), (icon_x-11, 40))

    def draw_menu(self):
        screen.fill((10, 10, 30))
        title = font_large.render("SUBWAY RUNNER", True, WHITE)
        screen.blit(title, (WIDTH//2 - title.get_width()//2, 100))
        start = font_med.render("Press SPACE to Start", True, GOLD)
        screen.blit(start, (WIDTH//2 - start.get_width()//2, 330))
        hints = font_small.render("← → lanes | SPACE/↑ jump | ↓ slide", True, WHITE)
        screen.blit(hints, (WIDTH//2 - hints.get_width()//2, 390))

    def draw_gameover(self):
        overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
        overlay.fill((0, 0, 0, 190))
        screen.blit(overlay, (0, 0))
        go = font_large.render("GAME OVER", True, RED)
        screen.blit(go, (WIDTH//2 - go.get_width()//2, 150))
        score_txt = font_med.render(f"Score: {self.score}", True, WHITE)
        screen.blit(score_txt, (WIDTH//2 - score_txt.get_width()//2, 260))
        restart = font_med.render("Press R to Restart", True, GOLD)
        screen.blit(restart, (WIDTH//2 - restart.get_width()//2, 330))

    def draw(self):
        if self.state == "menu":
            self.draw_menu()
        elif self.state == "playing":
            self.draw_background()
            for obs in self.obstacles:
                obs.draw(screen)
            for coin in self.coins:
                coin.draw(screen)
            for pu in self.powerups:
                pu.draw(screen)
            self.player.draw(screen)
            for p in self.particles:
                p.draw(screen)
            self.draw_ui()
        elif self.state == "gameover":
            self.draw_background()
            for obs in self.obstacles:
                obs.draw(screen)
            for coin in self.coins:
                coin.draw(screen)
            for pu in self.powerups:
                pu.draw(screen)
            self.player.draw(screen)
            for p in self.particles:
                p.draw(screen)
            self.draw_gameover()

    def run(self):
        self.state = "menu"
        running = True
        while running:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    running = False
                if event.type == pygame.KEYDOWN:
                    if self.state == "menu" and event.key == pygame.K_SPACE:
                        self.reset()
                        self.state = "playing"
                    elif self.state == "gameover" and event.key == pygame.K_r:
                        self.reset()
                        self.state = "playing"
            if self.state == "playing":
                self.handle_input()
                self.update()
            self.draw()
            pygame.display.flip()
            clock.tick(60)

        pygame.quit()
        sys.exit()

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