"""재질별로 합쳐진 차량 glTF 에서 바퀴를 삼각형 단위로 잘라 내 시뮬레이터 관례의 .glb 로 만든다.

normalize_vehicle_gltf.py 는 바퀴 부품이 낱개 노드로 있는 모델(3ds Max 내보내기 등)을 묶는다. OBJ 에서 변환된
Sketchfab 모델(예: Genesis GV80 2021)은 노드가 'Object_N' 이고 한 메시에 네 타이어와 차체 조각이 섞여 있어 노드로는
바퀴를 고를 수 없다. 여기서는 바퀴 원통(축 = 좌우, 중심 = 차축 위치, 반지름 r, 바깥쪽 |X| > inner) 안에 무게중심이
든 삼각형을 떼어 바퀴마다 'Wheel<k>' 노드 아래 새 프리미티브로 옮긴다 (정점 버퍼는 그대로 함께 쓴다).

--repaint: 재질 이름이 실제 부위와 뒤섞인 모델(GV80 원본은 Sketchfab 에서도 후드가 노랑, 범퍼가 빨강)은 재질을 버리고
삼각형 위치·법선으로 다시 칠한다 — 타이어 · 휠, 유리(벨트라인 위 경사면), 그릴, 헤드램프 · 테일램프, 하부 트림, 나머지 차체 도장.
구역 경계는 아래 REPAINT 의 값 (바닥 · 앞뒤 끝 기준 미터).

결과 관례 (normalize_vehicle_gltf.py 와 같다): +Y 위, +Z 앞, +X 왼쪽, 바닥(타이어 아래) y = 0, 미터, 수평 중심 = 차체 경계 상자 중앙.

    python scripts/split_wheels_gltf.py assets/gltf/genesis_gv80/scene.gltf assets/gltf/GenesisGV80.glb \\
        --axes x,-z,y --scale 0.01 --axles 1.574,-1.376 --axle-height 0.395 --radius 0.41 --inner 0.60 --repaint
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))
from normalize_vehicle_gltf import pack_glb, positions  # noqa: E402

_CT = {5121: np.uint8, 5123: np.uint16, 5125: np.uint32}

# --repaint 재질 (glTF PBR) 과 구역. 높이 y 는 바닥에서, front/rear 는 차 앞끝 · 뒤끝에서 안쪽으로 잰 미터.
MATERIALS = {
    "paint": {"baseColorFactor": [0.86, 0.87, 0.88, 1], "metallicFactor": 0.35, "roughnessFactor": 0.25},  # 흰 펄
    "glass": {"baseColorFactor": [0.02, 0.025, 0.03, 1], "metallicFactor": 0.0, "roughnessFactor": 0.05},
    "trim": {"baseColorFactor": [0.03, 0.03, 0.03, 1], "metallicFactor": 0.0, "roughnessFactor": 0.6},
    "grille": {"baseColorFactor": [0.015, 0.015, 0.015, 1], "metallicFactor": 0.6, "roughnessFactor": 0.3},
    "headlight": {"baseColorFactor": [0.9, 0.92, 0.95, 1], "metallicFactor": 0.0, "roughnessFactor": 0.1,
                  "emissive": [0.6, 0.62, 0.65]},
    "taillight": {"baseColorFactor": [0.55, 0.02, 0.02, 1], "metallicFactor": 0.0, "roughnessFactor": 0.15,
                  "emissive": [0.35, 0.0, 0.0]},
    "tire": {"baseColorFactor": [0.025, 0.025, 0.025, 1], "metallicFactor": 0.0, "roughnessFactor": 0.85},
    "rim": {"baseColorFactor": [0.6, 0.6, 0.62, 1], "metallicFactor": 1.0, "roughnessFactor": 0.3},
}
REPAINT = {
    "belt_y": 1.02,         # 이 높이 위의 기운 면 = 유리, 거의 수평인 면(지붕 · 후드 끝) = 도장
    "roof_ny": 0.9,         # |법선 y| 가 이보다 크면 수평 면
    "trim_y": 0.30,         # 이 아래 = 하부 트림 (사이드실 · 범퍼 아래)
    "rim_r": 0.27,          # 바퀴 원통 안에서 축으로부터 이 반지름 안 = 휠, 밖 = 타이어
    "grille": {"front": 0.25, "half_x": 0.50, "y": [0.32, 0.85]},
    "headlight": {"front": 0.35, "min_x": 0.55, "y": [0.72, 1.00]},
    "taillight": {"rear": 0.25, "min_x": 0.45, "y": [0.86, 1.06]},
}


def classify(c: np.ndarray, n: np.ndarray, zmin: float, zmax: float) -> np.ndarray:
    """차체 삼각형(무게중심 c, 단위 법선 n, 바닥 y=0 기준)의 MATERIALS 이름."""
    R = REPAINT
    y, ax, z = c[:, 1], np.abs(c[:, 0]), c[:, 2]
    out = np.full(len(c), "paint", dtype=object)
    out[y < R["trim_y"]] = "trim"
    out[(y > R["belt_y"]) & (np.abs(n[:, 1]) < R["roof_ny"])] = "glass"
    g, h, t = R["grille"], R["headlight"], R["taillight"]
    out[(z > zmax - g["front"]) & (ax < g["half_x"]) & (y > g["y"][0]) & (y < g["y"][1])] = "grille"
    out[(z > zmax - h["front"]) & (ax > h["min_x"]) & (y > h["y"][0]) & (y < h["y"][1])] = "headlight"
    out[(z < zmin + t["rear"]) & (ax > t["min_x"]) & (y > t["y"][0]) & (y < t["y"][1])] = "taillight"
    return out


def read_indices(j: dict, blob: bytes, acc: int, nverts: int) -> np.ndarray:
    if acc is None:
        return np.arange(nverts, dtype=np.uint32)
    a = j["accessors"][acc]
    bv = j["bufferViews"][a["bufferView"]]
    dt = _CT[a["componentType"]]
    off = bv.get("byteOffset", 0) + a.get("byteOffset", 0)
    return np.frombuffer(blob, dtype=dt, count=a["count"], offset=off).astype(np.uint32)


def axes_matrix(spec: str) -> np.ndarray:
    """'x,-z,y' → 원본 좌표에서 (시뮬 X, Y, Z) 를 고르는 3x3."""
    R = np.zeros((3, 3))
    for row, s in enumerate(spec.split(",")):
        s = s.strip()
        R[row, "xyz".index(s[-1])] = -1.0 if s.startswith("-") else 1.0
    if abs(np.linalg.det(R) - 1) > 1e-6:
        raise SystemExit(f"--axes {spec} 는 회전이 아니다 (det={np.linalg.det(R):.0f}) — 좌우가 뒤집힌다")
    return R


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("src", type=Path, help="원본 .gltf (같은 폴더에 .bin)")
    ap.add_argument("out", type=Path, help="결과 .glb")
    ap.add_argument("--axes", default="x,y,z", help="시뮬 X(왼쪽),Y(위),Z(앞) 가 원본의 어느 축인지, 예: x,-z,y")
    ap.add_argument("--scale", type=float, default=1.0, help="원본 단위 → 미터")
    ap.add_argument("--axles", required=True, help="앞 차축 Z, 뒤 차축 Z (축·배율 적용 후, 중심 맞추기 전, 미터)")
    ap.add_argument("--axle-height", type=float, required=True, help="차축 높이 Y (축·배율 적용 후, 미터)")
    ap.add_argument("--radius", type=float, required=True, help="바퀴 원통 반지름 (타이어 바깥 + 여유, 미터)")
    ap.add_argument("--inner", type=float, required=True, help="바퀴로 볼 |X| 최소 (미터, 바퀴 안쪽 면)")
    ap.add_argument("--repaint", action="store_true", help="원본 재질을 버리고 위치로 다시 칠한다 (REPAINT)")
    a = ap.parse_args()

    j = json.loads(a.src.read_text())
    blob = bytearray((a.src.parent / j["buffers"][0]["uri"]).read_bytes())
    R = axes_matrix(a.axes) * a.scale
    zf, zr = (float(v) for v in a.axles.split(","))
    centers = [(+1, zf), (-1, zf), (+1, zr), (-1, zr)]  # (X 부호, Z): 앞왼, 앞오, 뒤왼, 뒤오
    names = ["WheelFrontLeft", "WheelFrontRight", "WheelRearLeft", "WheelRearRight"]

    mesh_nodes = [i for i, n in enumerate(j["nodes"]) if "mesh" in n]
    for i in mesh_nodes:
        n = j["nodes"][i]
        if any(k in n for k in ("matrix", "translation", "rotation", "scale")):
            raise SystemExit(f"메시 노드 {n.get('name')} 에 변환이 있다 — 이 스크립트는 항등 변환인 메시 노드만 다룬다")
    if len(j["buffers"]) != 1:
        raise SystemExit("버퍼가 여러 개인 glTF 는 다루지 않는다")

    def add_indices(idx: np.ndarray) -> int:
        while len(blob) % 4:
            blob.append(0)
        data = idx.astype(np.uint32).tobytes()
        j["bufferViews"].append({"buffer": 0, "byteOffset": len(blob), "byteLength": len(data), "target": 34963})
        blob.extend(data)
        j["accessors"].append({"bufferView": len(j["bufferViews"]) - 1, "componentType": 5125,
                               "count": int(idx.size), "type": "SCALAR"})
        return len(j["accessors"]) - 1

    prims = []  # (프리미티브, 시뮬 좌표 정점(중심 맞추기 전), 삼각형)
    lo, hi = np.full(3, np.inf), np.full(3, -np.inf)
    for i in mesh_nodes:
        for p in j["meshes"][j["nodes"][i]["mesh"]]["primitives"]:
            if p.get("mode", 4) != 4:
                continue
            v = positions(j, a.src.parent, p["attributes"]["POSITION"]) @ R.T
            lo, hi = np.minimum(lo, v.min(0)), np.maximum(hi, v.max(0))
            prims.append((p, v, read_indices(j, bytes(blob), p.get("indices"), len(v)).reshape(-1, 3)))

    mat_id = {}
    if a.repaint:
        for name, m in MATERIALS.items():
            pbr = {k: m[k] for k in ("baseColorFactor", "metallicFactor", "roughnessFactor")}
            mm = {"name": f"sim_{name}", "pbrMetallicRoughness": pbr, "doubleSided": True}
            if "emissive" in m:
                mm["emissiveFactor"] = m["emissive"]
            j["materials"].append(mm)
            mat_id[name] = len(j["materials"]) - 1

    def sub(p: dict, tris: np.ndarray, mat: str | None) -> dict:
        q = {kk: vv for kk, vv in p.items() if kk != "indices"}
        q["indices"] = add_indices(tris.reshape(-1))
        if mat is not None:
            q["material"] = mat_id[mat]
        return q

    wheel_prims: list[list[dict]] = [[] for _ in centers]
    n_wheel = n_all = 0
    counts: dict[str, int] = {}
    for p, v, tri in prims:
        c = v[tri].mean(1)
        owner = np.full(len(tri), -1)
        for k, (sx, zc) in enumerate(centers):
            inside = (sx * c[:, 0] > a.inner) & (np.hypot(c[:, 1] - a.axle_height, c[:, 2] - zc) < a.radius)
            owner[inside & (owner < 0)] = k
        n_all += len(tri)
        n_wheel += int((owner >= 0).sum())
        if a.repaint:
            e = np.cross(v[tri[:, 1]] - v[tri[:, 0]], v[tri[:, 2]] - v[tri[:, 0]])
            nrm = e / np.maximum(np.linalg.norm(e, axis=1, keepdims=True), 1e-12)
            cls = classify(c - [0, lo[1], 0], nrm, lo[2], hi[2])
            for k, (_, zc) in enumerate(centers):
                w = owner == k
                cls[w] = np.where(np.hypot(c[w, 1] - a.axle_height, c[w, 2] - zc) < REPAINT["rim_r"], "rim", "tire")
        new_body = []
        for k in range(len(centers)):
            sel = owner == k
            if not sel.any():
                continue
            if a.repaint:
                for m in ("tire", "rim"):
                    s2 = sel & (cls == m)
                    if s2.any():
                        wheel_prims[k].append(sub(p, tri[s2], m))
                        counts[m] = counts.get(m, 0) + int(s2.sum())
            else:
                wheel_prims[k].append(sub(p, tri[sel], None))
        body = owner < 0
        if a.repaint:
            for m in MATERIALS:
                s2 = body & (cls == m)
                if s2.any():
                    new_body.append(sub(p, tri[s2], m))
                    counts[m] = counts.get(m, 0) + int(s2.sum())
        elif body.any():
            new_body.append(sub(p, tri[body], None))
        p["_split"] = new_body
    # 차체 프리미티브를 재질별 조각으로 바꾼다
    for i in mesh_nodes:
        mesh = j["meshes"][j["nodes"][i]["mesh"]]
        out = []
        for p in mesh["primitives"]:
            out += p.pop("_split") if "_split" in p else [p]
        mesh["primitives"] = [q for q in out if j["accessors"][q["indices"]]["count"] > 0]
    if a.repaint:
        print("다시 칠한 삼각형:", ", ".join(f"{m} {n:,}" for m, n in counts.items()))

    root_children = list(mesh_nodes)
    for k, prims in enumerate(wheel_prims):
        if not prims:
            raise SystemExit(f"{names[k]} 에 삼각형이 없다 — --axles/--axle-height/--radius/--inner 확인")
        j["meshes"].append({"name": names[k], "primitives": prims})
        j["nodes"].append({"name": names[k], "children": [len(j["nodes"]) + 1]})
        j["nodes"].append({"name": names[k] + "_mesh", "mesh": len(j["meshes"]) - 1})
        root_children.append(len(j["nodes"]) - 2)
    # 메시 노드는 원래 부모에서 떼고 정렬 노드 하나 아래로 (원본 루트의 축 회전은 버리고 --axes 로 대신한다)
    for n in j["nodes"]:
        if "children" in n and not n.get("name", "").startswith("Wheel"):
            n["children"] = [c for c in n["children"] if c not in mesh_nodes]
    M = np.eye(4)
    M[:3, :3] = R
    M[:3, 3] = [-(lo[0] + hi[0]) / 2, -lo[1], -(lo[2] + hi[2]) / 2]
    j["nodes"].append({"name": "SimAlign", "matrix": M.T.reshape(-1).tolist(), "children": root_children})
    j["scenes"] = [{"nodes": [len(j["nodes"]) - 1]}]
    j["scene"] = 0
    size = hi - lo
    print(f"크기 폭 {size[0]:.3f} x 높이 {size[1]:.3f} x 길이 {size[2]:.3f} m, 축거 {zf - zr:.3f} m, "
          f"바퀴 삼각형 {n_wheel:,} / {n_all:,} ({[sum(int(j['accessors'][q['indices']]['count']) // 3 for q in w) for w in wheel_prims]})")

    tmp = a.out.with_suffix(".bin.tmp")
    tmp.write_bytes(bytes(blob))
    j["buffers"] = [{"uri": tmp.name, "byteLength": len(blob)}]
    a.out.parent.mkdir(parents=True, exist_ok=True)
    try:
        pack_glb(j, tmp.parent, a.out)
    finally:
        tmp.unlink()


if __name__ == "__main__":
    main()