import pygame
import random
import sys

pygame.init()
WIDTH, HEIGHT = 480, 640
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("扶老奶奶过马路")

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GRAY = (80, 80, 80)
YELLOW = (255, 220, 0)
RED = (220, 30, 30)
GREEN = (30, 180, 30)
BLUE = (30, 140, 220)
ORANGE = (255, 130, 30)
PURPLE = (160, 60, 220)
BROWN = (110, 70, 30)

# 字体容错加载
try:
    font = pygame.font.SysFont("simhei", 32)
    small_font = pygame.font.SysFont("simhei", 24)
except:
    try:
        font = pygame.font.SysFont("msyh", 32)
        small_font = pygame.font.SysFont("msyh", 24)
    except:
        font = pygame.font.Font(None, 32)
        small_font = pygame.font.Font(None, 24)

# 游戏状态
STATE_MENU = 0
STATE_PLAY = 1
STATE_GAMEOVER = 2
game_state = STATE_MENU
select_index = 0

# 难度配置
difficulty_config = {
    "easy": {"speed_min": 2, "speed_max": 4, "spawn": 120},
    "normal": {"speed_min": 3, "speed_max": 5, "spawn": 80},
    "hard": {"speed_min": 5, "speed_max": 7, "spawn": 50}
}
diff_key = "normal"

# 障碍物类型常量
BIKE = 1
CAR = 2
TRUCK = 3

# 道具类型
ITEMS_SLOW = 1
ITEMS_SPEED = 2
ITEMS_SHIELD = 3

# 玩家
old_lady_size = 40
player_base_speed = 5
player_speed = player_base_speed
has_shield = False

# 车辆列表
cars = []
spawn_timer = 0

# 关卡、分数、生命
level = 1
score = 0
base_lane_count = 4
lives = 3

# 道具系统
items = []
item_size = 26
slow_timer = 0
speed_boost_timer = 0

# 红绿灯系统
TRAFFIC_LIGHT_GREEN = 0
TRAFFIC_LIGHT_RED = 1
traffic_light_state = TRAFFIC_LIGHT_GREEN
traffic_light_timer = 0
# 绿灯和红灯持续时间，帧数
green_duration = 180
red_duration = 0.000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000001

clock = pygame.time.Clock()

def reset_level():
    global old_lady_x, old_lady_y, cars, spawn_timer, items, player_speed
    old_lady_x = WIDTH // 2 - old_lady_size // 2
    old_lady_y = HEIGHT - 80
    cars.clear()
    items.clear()
    spawn_timer = 0
    player_speed = player_base_speed

def get_lane_positions():
    lane_num = base_lane_count + level - 1
    lanes = []
    start_y = 100
    gap = 110
    for i in range(lane_num):
        lanes.append(start_y + i * gap)
    return lanes

def spawn_random_item(lanes):
    lane_y = random.choice(lanes)
    item_x = random.randint(50, WIDTH - 50)
    item_type = random.choice([ITEMS_SLOW, ITEMS_SPEED, ITEMS_SHIELD])
    items.append({
        "x": item_x,
        "y": lane_y,
        "type": item_type
    })

def create_vehicle(lane_y, base_min, base_max):
    veh_type = random.choice([BIKE, CAR, CAR, TRUCK])
    start_x = random.choice([-80, WIDTH])
    direc = 1 if start_x < 0 else -1

    if veh_type == BIKE:
        w, h = 28, 22
        spd_mult = 1.3
        color = PURPLE
    elif veh_type == CAR:
        w, h = 60, 35
        spd_mult = 1.0
        color = random.choice([(200, 50, 50), (50, 80, 200), (40, 160, 60), (230, 190, 40)])
    else:
        w, h = 90, 42
        spd_mult = 0.65
        color = BROWN

    raw_speed = random.randint(base_min, base_max) * spd_mult

    return {
        "x": start_x,
        "y": lane_y,
        "width": w,
        "height": h,
        "speed": raw_speed * direc,
        "color": color,
        "vtype": veh_type
    }

def draw_menu():
    screen.fill(GRAY)
    title = font.render("Choose Difficulty", True, WHITE)
    screen.blit(title, (WIDTH // 2 - title.get_width() // 2, 120))
    options = ["Easy", "Normal", "Hard"]
    y_start = 220
    for i, text in enumerate(options):
        color = YELLOW if i == select_index else WHITE
        txt = font.render(text, True, color)
        screen.blit(txt, (WIDTH // 2 - txt.get_width() // 2, y_start + i * 60))
    hint = small_font.render("UP/DOWN select, ENTER confirm", True, WHITE)
    screen.blit(hint, (WIDTH // 2 - hint.get_width() // 2, 420))

def draw_traffic_light():
    """绘制红绿灯小图标"""
    box_x = WIDTH - 50
    box_y = 10
    pygame.draw.rect(screen, BLACK, (box_x, box_y, 30, 55))

    if traffic_light_state == TRAFFIC_LIGHT_GREEN:
        pygame.draw.circle(screen, GREEN, (box_x + 15, box_y + 15), 10)
        pygame.draw.circle(screen, (100, 100, 100), (box_x + 15, box_y + 40), 10)
    else:
        pygame.draw.circle(screen, (100, 100, 100), (box_x + 15, box_y + 15), 10)
        pygame.draw.circle(screen, RED, (box_x + 15, box_y + 40), 10)

# 主循环
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()

        if event.type == pygame.KEYDOWN:
            if game_state == STATE_MENU:
                if event.key == pygame.K_UP:
                    select_index = max(0, select_index - 1)
                if event.key == pygame.K_DOWN:
                    select_index = min(2, select_index + 1)
                if event.key == pygame.K_RETURN:
                    diff_list = ["easy", "normal", "hard"]
                    diff_key = diff_list[select_index]
                    level = 1
                    score = 0
                    lives = 3
                    has_shield = False
                    slow_timer = 0
                    speed_boost_timer = 0
                    traffic_light_state = TRAFFIC_LIGHT_GREEN
                    traffic_light_timer = 0
                    reset_level()
                    game_state = STATE_PLAY

            elif game_state == STATE_GAMEOVER:
                if event.key == pygame.K_SPACE:
                    game_state = STATE_MENU

    if game_state == STATE_MENU:
        draw_menu()

    elif game_state == STATE_PLAY:
        cfg = difficulty_config[diff_key]
        lanes = get_lane_positions()
        screen.fill(GRAY)

        # 车道线
        for ly in lanes:
            pygame.draw.line(screen, YELLOW, (0, ly), (WIDTH, ly), 3)

        # 红绿灯计时
        traffic_light_timer += 1
        if traffic_light_state == TRAFFIC_LIGHT_GREEN:
            if traffic_light_timer >= green_duration:
                traffic_light_state = TRAFFIC_LIGHT_RED
                traffic_light_timer = 0
        else:
            if traffic_light_timer >= red_duration:
                traffic_light_state = TRAFFIC_LIGHT_GREEN
                traffic_light_timer = 0

        # 道具计时
        if slow_timer > 0:
            slow_timer -= 1
        if speed_boost_timer > 0:
            speed_boost_timer -= 1
            player_speed = player_base_speed + 2
        else:
            player_speed = player_base_speed

        # 老奶奶移动
        keys = pygame.key.get_pressed()
        if keys[pygame.K_LEFT] and old_lady_x > 0:
            old_lady_x -= player_speed
        if keys[pygame.K_RIGHT] and old_lady_x < WIDTH - old_lady_size:
            old_lady_x += player_speed
        if keys[pygame.K_UP] and old_lady_y > 0:
            old_lady_y -= player_speed
        if keys[pygame.K_DOWN] and old_lady_y < HEIGHT - old_lady_size:
            old_lady_y += player_speed

        # 生成车辆
        spawn_timer += 1
        if spawn_timer >= cfg["spawn"]:
            lane_y = random.choice(lanes)
            cars.append(create_vehicle(lane_y, cfg["speed_min"], cfg["speed_max"]))
            spawn_timer = 0
            if random.random() < 0.25:
                spawn_random_item(lanes)

        # 更新车辆位置：红灯时车辆停止
        for car in cars:
            if traffic_light_state == TRAFFIC_LIGHT_GREEN:
                car["x"] += car["speed"]
        cars = [c for c in cars if -100 < c["x"] < WIDTH + 20]

        # 碰撞检测
        player_rect = pygame.Rect(old_lady_x, old_lady_y, old_lady_size, old_lady_size)
        crash = False
        for car in cars:
            car_rect = pygame.Rect(car["x"], car["y"], car["width"], car["height"])
            if player_rect.colliderect(car_rect):
                crash = True
                break

        if crash:
            if has_shield:
                has_shield = False
                reset_level()
            else:
                lives -= 1
                reset_level()
                if lives <= 0:
                    game_state = STATE_GAMEOVER

        # 拾取道具
        new_items = []
        for item in items:
            item_rect = pygame.Rect(item["x"], item["y"], item_size, item_size)
            if player_rect.colliderect(item_rect):
                if item["type"] == ITEMS_SLOW:
                    slow_timer = 240
                elif item["type"] == ITEMS_SPEED:
                    speed_boost_timer = 240
                elif item["type"] == ITEMS_SHIELD:
                    has_shield = True
            else:
                new_items.append(item)
        items = new_items

        # 通关进入下一关
        if old_lady_y <= 10:
            score += 100 * level
            level += 1
            reset_level()

        # 绘制老奶奶护盾外圈
        if has_shield:
            pygame.draw.circle(screen, GREEN,
                               (old_lady_x + old_lady_size // 2, old_lady_y + old_lady_size // 2),
                               old_lady_size // 2 + 6, 3)

        # 绘制老奶奶本体
        pygame.draw.rect(screen, (140, 90, 30), (old_lady_x, old_lady_y, old_lady_size, old_lady_size))
        pygame.draw.circle(screen, (255, 220, 200), (old_lady_x + 20, old_lady_y + 12), 11)

        # 绘制全部车辆障碍物
        for car in cars:
            pygame.draw.rect(screen, car["color"], (car["x"], car["y"], car["width"], car["height"]))

        # 绘制道具
        for item in items:
            cx = item["x"] + item_size // 2
            cy = item["y"] + item_size // 2
            if item["type"] == ITEMS_SLOW:
                pygame.draw.circle(screen, BLUE, (cx, cy), item_size // 2)
            elif item["type"] == ITEMS_SPEED:
                pygame.draw.circle(screen, ORANGE, (cx, cy), item_size // 2)
            elif item["type"] == ITEMS_SHIELD:
                pygame.draw.circle(screen, GREEN, (cx, cy), item_size // 2)

        # 绘制红绿灯
        draw_traffic_light()

        # 顶部信息
        shield_text = " Shield:ON" if has_shield else ""
        light_text = " GREEN" if traffic_light_state == TRAFFIC_LIGHT_GREEN else " RED"
        info_text = small_font.render(f"Level:{level}  Score:{score}  Lives:{lives}{shield_text}{light_text}", True, WHITE)
        screen.blit(info_text, (10, 8))

    elif game_state == STATE_GAMEOVER:
        screen.fill(GRAY)
        t1 = font.render("Game Over!", True, RED)
        t2 = font.render(f"Final Score: {score}", True, WHITE)
        t3 = small_font.render("Press SPACE back to menu", True, WHITE)
        screen.blit(t1, (WIDTH // 2 - t1.get_width() // 2, 200))
        screen.blit(t2, (WIDTH // 2 - t2.get_width() // 2, 260))
        screen.blit(t3, (WIDTH // 2 - t3.get_width() // 2, 340))

    pygame.display.flip()
    clock.tick(60)