"""UE5 静态模型几何提取脚本。

对 index.json 中发现的每个 StaticMesh 提取：
- 整体包围盒（min, max, size）
- 锚点类型和 XY 偏移
- 网格复杂度（顶点数、三角形数）
- Socket 列表
- 材质槽位
- AI 占位字段（key_surfaces, style_tags, scene_tags 等）

输出: output/meshes/<SubfolderName>.json，每个子目录一个文件

用法:
  交互模式: exec(open(r"...\analyze_meshes.py").read())
  命令行:   通过 run_all.py 调用
"""

import sys
from pathlib import Path

_SCRIPT_DIR = Path(__file__).resolve().parent
if str(_SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(_SCRIPT_DIR))

from _common import (
    HAS_UNREAL, get_scan_root, get_output_dir, get_steps,
    write_json, read_json, wait_for_editor, progress_log,
    iso_now, safe_subfolder_name, classify_pivot,
)

TAG = "[MeshAnalysis]"


# ============================================================
# 自动 Tag 生成
# ============================================================

# (pattern, [style_tags])
_MESH_NAME_RULES = [
    # 家具
    (r"Table", ["furniture", "table"]),
    (r"Chair|Armchair", ["furniture", "seating"]),
    (r"Bed", ["furniture", "bed"]),
    (r"Monitor|Laptop|Keyboard", ["furniture", "electronics"]),
    (r"Broom|Radiator", ["furniture", "appliance"]),
    # 建筑
    (r"Wall", ["architecture", "wall"]),
    (r"Floor", ["architecture", "floor"]),
    (r"Door", ["architecture", "door"]),
    (r"Window", ["architecture", "window"]),
    (r"Stair", ["architecture", "stair"]),
    (r"Beam|Ladder", ["architecture", "structure"]),
    (r"Barrier|Fence", ["architecture", "barrier"]),
    # 灯光
    (r"Lamp|Light|Spot|Flash|Torch", ["lighting"]),
    # 车辆
    (r"\bCar\b|Truck|Bus|Boat|Plane|Metro|Trailer", ["vehicle"]),
    # 工业
    (r"Barrel|WoodenBox|WoodenCrate|Pallet|Container", ["industrial"]),
    # 道路
    (r"Road|Barricade|SpeedBump|Cone", ["outdoor", "road"]),
    # 基础体
    (r"\bCube\b|\bCylinder\b|\bSphere\b|\bCone\b|\bRamp\b|\bPipe\b|Quarter", ["primitive"]),
    # 浴室
    (r"Toilet|Urinal", ["bathroom"]),
    # 装饰
    (r"Book|Bottle|Mug|Plate|Can|Flask|Padlock|Bucket|Trash", ["decoration", "prop"]),
    # 自然
    (r"Tree|Foliage|Mount", ["vegetation"]),
    # 武器
    (r"Pistol|Sniper|Sword|Dagger|Target", ["weapon"]),
    # 人物
    (r"Human|Camera|Cage", ["character", "reference"]),
    # 展示
    (r"DisplayPanel", ["display"]),
]

# (subdirectory keyword, [scene_tags])
_MESH_SCENE_RULES = [
    (r"Furnitures|Bedroom|Bathroom|Decorations", ["indoor", "home"]),
    (r"Lights", ["indoor", "outdoor"]),
    (r"Vehicles|Roads", ["outdoor", "street"]),
    (r"Industrials|Pipes|Gears", ["outdoor", "industrial"]),
    (r"SmallBuildings|BigBuildings", ["outdoor", "building"]),
    (r"Architecture", ["architecture", "construction"]),
    (r"Vegetations", ["outdoor", "nature"]),
    (r"Primitives|Extra", ["utility"]),
    (r"Weapons|Interactions", ["gameplay"]),
]


def _generate_mesh_tags(name: str, subfolder: str) -> tuple:
    """根据模型名和目录生成 style_tags 和 scene_tags。

    Returns:
        (style_tags, scene_tags)
    """
    import re
    style_tags = []
    scene_tags = []

    # 基于名称的 style tags
    for pattern, tags in _MESH_NAME_RULES:
        if re.search(pattern, name):
            for t in tags:
                if t not in style_tags:
                    style_tags.append(t)

    # 基于目录的 scene tags
    for pattern, tags in _MESH_SCENE_RULES:
        if re.search(pattern, subfolder):
            for t in tags:
                if t not in scene_tags:
                    scene_tags.append(t)

    # 如果没有任何 style tag，标记为 unknown
    if not style_tags:
        style_tags = ["unknown"]
    if not scene_tags:
        scene_tags = ["unclassified"]

    return style_tags, scene_tags


# ============================================================
# 单个 Mesh 分析
# ============================================================

def _round3(v):
    """将浮点数或列表保留 3 位小数。"""
    if isinstance(v, (list, tuple)):
        return [round(x, 3) if isinstance(x, float) else x for x in v]
    return round(v, 3) if isinstance(v, float) else v


def _try_get_complexity(mesh) -> dict:
    """尝试获取网格复杂度信息。"""
    info = {"num_verts": None, "num_triangles": None}
    try:
        import unreal
        # 尝试获取 LOD 数量
        num_lod = mesh.get_num_lods()
        info["num_lods"] = num_lod
    except Exception:
        pass

    try:
        # 尝试获取 section 信息
        num_sections = mesh.get_num_sections(0)
        info["num_sections"] = num_sections
    except Exception:
        pass

    try:
        # 尝试通过 render data 获取顶点数
        render_data = mesh.get_render_data()
        if render_data:
            info["num_verts"] = render_data.num_vertices
            info["num_triangles"] = render_data.num_triangles
    except Exception:
        pass

    return info


def _try_get_sockets(mesh) -> list:
    """尝试获取 Socket 名称列表。"""
    sockets = []
    try:
        import unreal
        socket_names = mesh.find_socket_names() if hasattr(mesh, 'find_socket_names') else []
        for name in socket_names:
            try:
                sock = mesh.find_socket(name)
                if sock:
                    loc = sock.relative_location
                    rot = sock.relative_rotation
                    sockets.append({
                        "name": str(name),
                        "relative_location": _round3([loc.x, loc.y, loc.z]),
                        "relative_rotation": _round3([rot.pitch, rot.yaw, rot.roll]),
                    })
            except Exception:
                sockets.append({"name": str(name), "relative_location": None, "relative_rotation": None})
    except Exception:
        pass
    return sockets


def _try_get_material_slots(mesh) -> list:
    """尝试获取材质槽位。"""
    import unreal
    slots = []

    # 方法 1: get_material(i)
    try:
        num_mat = mesh.get_num_materials()
        for i in range(num_mat):
            slot_info = {"index": i, "slot_name": None, "material_name": None}
            try:
                mat = mesh.get_material(i)
                if mat:
                    slot_info["material_name"] = mat.get_name()
            except Exception:
                pass
            slots.append(slot_info)
    except Exception:
        pass

    # 方法 2: get_editor_property("static_materials")
    if not slots:
        try:
            static_mats = mesh.get_editor_property("static_materials")
            if static_mats:
                for i, sm in enumerate(static_mats):
                    slot_info = {"index": i, "slot_name": None, "material_name": None}
                    try:
                        slot_info["slot_name"] = sm.get_editor_property("material_slot_name")
                        mi = sm.get_editor_property("material_interface")
                        if mi:
                            slot_info["material_name"] = mi.get_name()
                    except Exception:
                        pass
                    slots.append(slot_info)
        except Exception:
            pass

    # 方法 3: get_material_slot_names() — 至少获取槽位名
    if not slots:
        try:
            slot_names = mesh.get_material_slot_names()
            for i, name in enumerate(slot_names):
                slots.append({"index": i, "slot_name": str(name), "material_name": None})
        except Exception:
            pass

    # 方法 4: 用 section 数作为 fallback
    if not slots:
        try:
            num_sections = mesh.get_num_sections(0)
            for i in range(num_sections):
                slots.append({"index": i, "slot_name": None, "material_name": None})
        except Exception:
            pass

    return slots


def _analyze_single_mesh(asset_info: dict, subfolder: str = "") -> dict:
    """分析单个 StaticMesh，返回完整信息。"""
    import unreal

    ue_path = asset_info["ue_path"]
    name = asset_info["name"]

    result = {
        "ue_path": ue_path,
        "bounding_box": None,
        "pivot": None,
        "mesh_complexity": {},
        "sockets": [],
        "material_slots": [],
        "analysis_notes": [],
        # AI Pass 2 占位
        "key_surfaces": [],
        "ground_sink_estimate": None,
        "style_tags": [],
        "scene_tags": [],
    }

    # 加载 mesh
    try:
        mesh = unreal.load_asset(ue_path)
    except Exception as e:
        result["analysis_notes"].append(f"加载失败: {e}")
        return result

    if mesh is None:
        result["analysis_notes"].append("加载返回 None")
        return result

    # ---- 1. 包围盒 ----
    try:
        bb = mesh.get_bounding_box()
        bmin = (bb.min.x, bb.min.y, bb.min.z)
        bmax = (bb.max.x, bb.max.y, bb.max.z)
        size = (bmax[0] - bmin[0], bmax[1] - bmin[1], bmax[2] - bmin[2])

        result["bounding_box"] = {
            "min": _round3(list(bmin)),
            "max": _round3(list(bmax)),
            "size": _round3(list(size)),
        }

        # ---- 2. 锚点 ----
        result["pivot"] = classify_pivot(bmin, bmax)

    except Exception as e:
        result["analysis_notes"].append(f"包围盒提取失败: {e}")
        return result

    # ---- 3. 网格复杂度 ----
    result["mesh_complexity"] = _try_get_complexity(mesh)

    # ---- 4. Sockets ----
    result["sockets"] = _try_get_sockets(mesh)

    # ---- 5. 材质槽位 ----
    result["material_slots"] = _try_get_material_slots(mesh)

    # ---- 6. 自动 Tags ----
    style_tags, scene_tags = _generate_mesh_tags(name, subfolder)
    result["style_tags"] = style_tags
    result["scene_tags"] = scene_tags

    return result


# ============================================================
# 主流程
# ============================================================

def run(index_data: dict = None) -> dict:
    """执行模型分析。

    Args:
        index_data: 已加载的 index.json 数据，若为 None 则从文件读取

    Returns:
        统计信息 dict
    """
    if not HAS_UNREAL:
        print(f"{TAG} 错误: unreal 模块不可用")
        return {"status": "error", "reason": "no unreal"}

    output_dir = get_output_dir()

    # 加载 index.json
    if index_data is None:
        index_path = output_dir / "index.json"
        if not index_path.exists():
            print(f"{TAG} 错误: index.json 不存在，请先运行 discover")
            return {"status": "error", "reason": "no index.json"}
        index_data = read_json(index_path)

    meshes_by_subfolder = index_data.get("meshes", {})
    if not meshes_by_subfolder:
        print(f"{TAG} 无 StaticMesh 资产需要分析")
        return {"status": "ok", "analyzed": 0}

    # 创建输出目录
    meshes_dir = output_dir / "meshes"
    meshes_dir.mkdir(parents=True, exist_ok=True)

    total_all = sum(len(assets) for assets in meshes_by_subfolder.values())
    print(f"{TAG} 共 {total_all} 个 StaticMesh，分布在 {len(meshes_by_subfolder)} 个子目录")

    analyzed = 0
    failed = 0
    global_idx = 0

    for subfolder, assets in meshes_by_subfolder.items():
        safe_name = safe_subfolder_name(subfolder)
        output_file = meshes_dir / f"{safe_name}.json"

        # 断点续跑: 跳过已完成的文件
        if output_file.exists():
            try:
                existing = read_json(output_file)
                if existing.get("completed"):
                    print(f"{TAG} 跳过已完成的: {safe_name}")
                    analyzed += existing.get("count", len(assets))
                    global_idx += len(assets)
                    continue
            except Exception:
                pass

        print(f"{TAG} 分析子目录: {subfolder} ({len(assets)} 个模型)")
        mesh_results = {}
        subfolder_analyzed = 0
        subfolder_failed = 0

        for i, asset_info in enumerate(assets):
            progress_log(i, len(assets), f"{asset_info['name']}")

            result = _analyze_single_mesh(asset_info, subfolder)
            mesh_results[asset_info["name"]] = result

            if result.get("bounding_box") is not None:
                subfolder_analyzed += 1
            else:
                subfolder_failed += 1

            global_idx += 1

        # 写入分组文件
        group_data = {
            "subfolder": subfolder,
            "source_path": f"{get_scan_root()}/Meshes/{subfolder}" if subfolder else f"{get_scan_root()}/Meshes",
            "count": len(assets),
            "completed": True,
            "meshes": mesh_results,
        }
        write_json(group_data, output_file)

        analyzed += subfolder_analyzed
        failed += subfolder_failed
        print(f"{TAG} ✅ {safe_name}: {subfolder_analyzed} 成功, {subfolder_failed} 失败 → {output_file}")

    print(f"{TAG} 全部完成: {analyzed} 成功, {failed} 失败 / {total_all} 总计")

    return {
        "status": "ok",
        "total": total_all,
        "analyzed": analyzed,
        "failed": failed,
    }


# ---- 入口 ----
if __name__ == "__main__":
    if not wait_for_editor():
        sys.exit(1)
    result = run()
    if result.get("status") != "ok":
        sys.exit(1)