#!/usr/bin/env python3
"""按空间平面把一个 GLB 拆成多个模块（默认沿 Y 轴切成 3 份）。

针对"多个相同模块 + 共享地板"的场景：
  - 完全落在某个模块范围内的零件，整体归入该模块（保留零件名/材质/世界变换）；
  - 跨越分割面的共享零件（如地板），在分割面处做几何裁剪，每个模块各留一段。

用法示例：
  python3 viewer/tools/split_glb.py \
      --input viewer/压缩机2326-3000.glb \
      --axis y --planes 16 32 \
      --outdir viewer/models --prefix 压缩机2326-3000
"""
import argparse, json, os, struct
import numpy as np

GLTF_JSON = 0x4E4F534A
GLTF_BIN = 0x004E4942

COMP_DTYPE = {5120: np.int8, 5121: np.uint8, 5122: np.int16,
              5123: np.uint16, 5125: np.uint32, 5126: np.float32}
COMP_SIZE = {5120: 1, 5121: 1, 5122: 2, 5123: 2, 5125: 4, 5126: 4}
NCOMP = {"SCALAR": 1, "VEC2": 2, "VEC3": 3, "VEC4": 4, "MAT2": 4, "MAT3": 9, "MAT4": 16}
MODES = {0: "POINTS", 1: "LINES", 2: "LINE_LOOP", 3: "LINE_STRIP",
         4: "TRIANGLES", 5: "TRIANGLE_STRIP", 6: "TRIANGLE_FAN"}


# --------------------------------------------------------------------------
# GLB 读写
# --------------------------------------------------------------------------
def read_glb(path):
    with open(path, "rb") as f:
        magic, ver, length = struct.unpack("<III", f.read(12))
        if magic != 0x46546C67:
            raise ValueError("不是有效的 GLB 文件")
        chunks = {}
        while f.tell() < length:
            clen, ctype = struct.unpack("<II", f.read(8))
            chunks[ctype] = f.read(clen)
    gltf = json.loads(chunks[GLTF_JSON].decode("utf-8"))
    return gltf, chunks.get(GLTF_BIN, b"")


def write_glb(path, gltf, bin_data):
    j = json.dumps(gltf, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
    j += b" " * ((4 - len(j) % 4) % 4)
    b = bytes(bin_data)
    b += b"\x00" * ((4 - len(b) % 4) % 4)
    total = 12 + 8 + len(j) + 8 + len(b)
    with open(path, "wb") as f:
        f.write(struct.pack("<III", 0x46546C67, 2, total))
        f.write(struct.pack("<II", len(j), GLTF_JSON))
        f.write(j)
        f.write(struct.pack("<II", len(b), GLTF_BIN))
        f.write(b)


# --------------------------------------------------------------------------
# 世界变换与包围盒
# --------------------------------------------------------------------------
def node_local_matrix(n):
    if "matrix" in n:
        return np.array(n["matrix"], dtype=np.float64).reshape(4, 4, order="F")
    T = np.eye(4)
    if "translation" in n:
        T[:3, 3] = n["translation"]
    R = np.eye(4)
    if "rotation" in n:
        x, y, z, w = n["rotation"]
        R[:3, :3] = np.array([
            [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
            [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
            [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)]])
    S = np.eye(4)
    if "scale" in n:
        S[0, 0], S[1, 1], S[2, 2] = n["scale"]
    return T @ R @ S


def world_matrices(gltf):
    nodes = gltf["nodes"]
    world = [None] * len(nodes)

    def walk(i, parent):
        world[i] = parent @ node_local_matrix(nodes[i])
        for c in nodes[i].get("children", []):
            walk(c, world[i])

    for root in gltf["scenes"][gltf.get("scene", 0)]["nodes"]:
        walk(root, np.eye(4))
    return world


def mesh_node_bbox(gltf, world, node_idx):
    """用 POSITION accessor 的 min/max 变换出世界包围盒（不做顶点解码，很快）。"""
    n = gltf["nodes"][node_idx]
    lo = np.array([np.inf] * 3)
    hi = np.array([-np.inf] * 3)
    for pr in gltf["meshes"][n["mesh"]]["primitives"]:
        ai = pr.get("attributes", {}).get("POSITION")
        if ai is None:
            continue
        a = gltf["accessors"][ai]
        if "min" not in a:
            continue
        mn = np.array(a["min"], float)
        mx = np.array(a["max"], float)
        corners = np.array([[x, y, z, 1.0] for x in (mn[0], mx[0])
                            for y in (mn[1], mx[1]) for z in (mn[2], mx[2])])
        wc = (world[node_idx] @ corners.T).T[:, :3]
        lo = np.minimum(lo, wc.min(axis=0))
        hi = np.maximum(hi, wc.max(axis=0))
    return lo, hi


# --------------------------------------------------------------------------
# 顶点解码 / 裁剪
# --------------------------------------------------------------------------
def decode_accessor(gltf, bin_data, ai):
    a = gltf["accessors"][ai]
    n = NCOMP[a["type"]]
    dt = COMP_DTYPE[a["componentType"]]
    bv = gltf["bufferViews"][a["bufferView"]]
    base = bv.get("byteOffset", 0) + a.get("byteOffset", 0)
    item = COMP_SIZE[a["componentType"]] * n
    stride = bv.get("byteStride") or item
    count = a["count"]
    if stride == item:
        arr = np.frombuffer(bin_data, dtype=dt, count=count * n, offset=base).reshape(count, n)
    else:
        arr = np.empty((count, n), dtype=dt)
        for i in range(count):
            arr[i] = np.frombuffer(bin_data, dtype=dt, count=n, offset=base + i * stride)
    arr = np.array(arr)
    if a.get("normalized") and dt != np.float32:
        info = np.iinfo(dt)
        arr = arr.astype(np.float32) / max(abs(info.min), info.max)
    else:
        arr = arr.astype(np.float64)
    return arr


def lerp_vert(a, b, t):
    return {k: a[k] + (b[k] - a[k]) * t for k in a}


def clip_half(verts, c, keep_le):
    out = []
    m = len(verts)
    for i in range(m):
        cur = verts[i]
        nxt = verts[(i + 1) % m]
        cy = cur["POSITION"][1]
        ny = nxt["POSITION"][1]
        cin = cy <= c if keep_le else cy >= c
        nin = ny <= c if keep_le else ny >= c
        if cin:
            out.append(cur)
        if cin != nin:
            t = (c - cy) / (ny - cy)
            out.append(lerp_vert(cur, nxt, t))
    return out


def clip_band(verts, indices, mode, lo, hi):
    """把一个图元裁剪到 lo<=y<=hi，返回 [('tri'|'line', [顶点...]), ...]。"""
    out = []
    if mode == 4:
        tri = indices if indices is not None else np.arange(len(verts))
        tri = tri.reshape(-1)
        for k in range(0, len(tri) - 2, 3):
            poly = [verts[tri[k]], verts[tri[k + 1]], verts[tri[k + 2]]]
            if hi is not None:
                poly = clip_half(poly, hi, True)
            if lo is not None:
                poly = clip_half(poly, lo, False)
            if len(poly) >= 3:
                out.append(("tri", poly))
    elif mode == 1:
        seg = indices if indices is not None else np.arange(len(verts))
        seg = seg.reshape(-1)
        for k in range(0, len(seg) - 1, 2):
            poly = [verts[seg[k]], verts[seg[k + 1]]]
            if hi is not None:
                poly = clip_half(poly, hi, True)
            if lo is not None:
                poly = clip_half(poly, lo, False)
            if len(poly) >= 2:
                out.append(("line", poly[:2]))
    else:
        raise NotImplementedError(f"暂不支持的图元模式 {MODES.get(mode, mode)}")
    return out


# --------------------------------------------------------------------------
# 输出文件构建
# --------------------------------------------------------------------------
class BandBuilder:
    def __init__(self, src):
        self.src = src
        self.gltf = {
            "asset": {"version": "2.0", "generator": "split_glb.py"},
            "scene": 0,
            "scenes": [{"nodes": []}],
            "nodes": [],
            "meshes": [],
            "accessors": [],
            "bufferViews": [],
            "materials": src.get("materials", []),
            "buffers": [{"byteLength": 0}],
        }
        for ext in ("extensionsUsed", "extensionsRequired"):
            if ext in src:
                self.gltf[ext] = src[ext]
        self.bin = bytearray()
        self._acc = {}
        self._bv = {}
        self._mesh = {}

    # ---- 原样拷贝 ----
    def _copy_bv(self, bi):
        if bi in self._bv:
            return self._bv[bi]
        bv = self.src["bufferViews"][bi]
        off = bv.get("byteOffset", 0)
        data = self.src_bin[off:off + bv["byteLength"]]
        while len(self.bin) % 4:
            self.bin.append(0)
        new = {"buffer": 0, "byteOffset": len(self.bin), "byteLength": len(data)}
        if "target" in bv:
            new["target"] = bv["target"]
        if "byteStride" in bv:
            new["byteStride"] = bv["byteStride"]
        self.bin += data
        self._bv[bi] = len(self.gltf["bufferViews"])
        self.gltf["bufferViews"].append(new)
        return self._bv[bi]

    def _copy_acc(self, ai):
        if ai in self._acc:
            return self._acc[ai]
        a = self.src["accessors"][ai]
        new = {k: a[k] for k in ("componentType", "count", "type", "normalized", "min", "max")
               if k in a}
        new["bufferView"] = self._copy_bv(a["bufferView"])
        new["byteOffset"] = a.get("byteOffset", 0)
        self._acc[ai] = len(self.gltf["accessors"])
        self.gltf["accessors"].append(new)
        return self._acc[ai]

    def _copy_mesh(self, mi):
        if mi in self._mesh:
            return self._mesh[mi]
        m = self.src["meshes"][mi]
        prims = []
        for pr in m["primitives"]:
            np_ = {"attributes": {k: self._copy_acc(v) for k, v in pr["attributes"].items()}}
            if "indices" in pr:
                np_["indices"] = self._copy_acc(pr["indices"])
            if "material" in pr:
                np_["material"] = pr["material"]
            if "mode" in pr:
                np_["mode"] = pr["mode"]
            prims.append(np_)
        self._mesh[mi] = len(self.gltf["meshes"])
        self.gltf["meshes"].append({"primitives": prims})
        return self._mesh[mi]

    def add_verbatim_node(self, node_idx, world):
        n = self.src["nodes"][node_idx]
        new = {
            "name": n.get("name", ""),
            "mesh": self._copy_mesh(n["mesh"]),
            "matrix": [float(x) for x in world.flatten(order="F")],
        }
        self.gltf["nodes"].append(new)
        self.gltf["scenes"][0]["nodes"].append(len(self.gltf["nodes"]) - 1)

    # ---- 裁剪件 ----
    def _add_float_accessor(self, arr, type_str, with_minmax=True):
        arr = np.asarray(arr, dtype=np.float32)
        while len(self.bin) % 4:
            self.bin.append(0)
        off = len(self.bin)
        self.bin += arr.tobytes()
        bv = {"buffer": 0, "byteOffset": off, "byteLength": arr.nbytes}
        self.gltf["bufferViews"].append(bv)
        acc = {"bufferView": len(self.gltf["bufferViews"]) - 1, "componentType": 5126,
               "count": len(arr), "type": type_str}
        if with_minmax and len(arr):
            acc["min"] = arr.min(axis=0).tolist()
            acc["max"] = arr.max(axis=0).tolist()
        self.gltf["accessors"].append(acc)
        return len(self.gltf["accessors"]) - 1

    def add_clipped_node(self, node_idx, world, lo, hi):
        src = self.src
        node = src["nodes"][node_idx]
        mesh = src["meshes"][node["mesh"]]
        W = world
        W3 = W[:3, :3]
        t = W[:3, 3]
        N = np.linalg.inv(W3).T
        prims = []
        for pr in mesh["primitives"]:
            mode = pr.get("mode", 4)
            P = decode_accessor(src, self.src_bin, pr["attributes"]["POSITION"])
            Pw = (W3 @ P.T).T + t
            has_n = "NORMAL" in pr["attributes"]
            if has_n:
                Nw = (N @ decode_accessor(src, self.src_bin, pr["attributes"]["NORMAL"]).T).T
                norms = np.linalg.norm(Nw, axis=1, keepdims=True)
                Nw = Nw / np.where(norms == 0, 1, norms)
            idx = decode_accessor(src, self.src_bin, pr["indices"]).reshape(-1).astype(np.int64) \
                if "indices" in pr else None
            verts = []
            for i in range(len(Pw)):
                v = {"POSITION": Pw[i]}
                if has_n:
                    v["NORMAL"] = Nw[i]
                verts.append(v)
            polys = clip_band(verts, idx, mode, lo, hi)
            if not polys:
                continue
            pos, nor = [], []
            for kind, poly in polys:
                if kind == "tri":
                    for j in range(1, len(poly) - 1):
                        for v in (poly[0], poly[j], poly[j + 1]):
                            pos.append(v["POSITION"])
                            if has_n:
                                nor.append(v["NORMAL"])
                else:
                    for v in poly[:2]:
                        pos.append(v["POSITION"])
                        if has_n:
                            nor.append(v["NORMAL"])
            if not pos:
                continue
            attrs = {"POSITION": self._add_float_accessor(pos, "VEC3")}
            if has_n:
                attrs["NORMAL"] = self._add_float_accessor(nor, "VEC3", with_minmax=False)
            p = {"attributes": attrs, "mode": mode}
            if "material" in pr:
                p["material"] = pr["material"]
            prims.append(p)
        if not prims:
            return False
        self.gltf["meshes"].append({"primitives": prims})
        self.gltf["nodes"].append({
            "name": node.get("name", ""),
            "mesh": len(self.gltf["meshes"]) - 1,
        })
        self.gltf["scenes"][0]["nodes"].append(len(self.gltf["nodes"]) - 1)
        return True

    def finish(self, path):
        self.gltf["buffers"][0]["byteLength"] = len(self.bin)
        write_glb(path, self.gltf, self.bin)


# --------------------------------------------------------------------------
def main():
    ap = argparse.ArgumentParser(description="按平面把 GLB 拆成多个模块并裁剪共享件")
    ap.add_argument("--input", required=True)
    ap.add_argument("--outdir", required=True)
    ap.add_argument("--prefix", default=None)
    ap.add_argument("--axis", default="y", choices=["x", "y", "z"])
    ap.add_argument("--planes", type=float, nargs="+", required=True,
                    help="分割面坐标，例如 --planes 16 32 会生成 3 个模块")
    args = ap.parse_args()

    src, src_bin = read_glb(args.input)
    world = world_matrices(src)
    ax = "xyz".index(args.axis)
    planes = sorted(args.planes)
    nband = len(planes) + 1

    # 每个 mesh 节点的世界包围盒
    nodes_info = []
    for i, n in enumerate(src["nodes"]):
        if "mesh" not in n:
            continue
        lo, hi = mesh_node_bbox(src, world, i)
        if not np.isfinite(lo).all():
            continue
        nodes_info.append((i, lo, hi, (lo[ax] + hi[ax]) / 2))

    # 跨分割面的零件（需要裁剪）：对每个分割面分别判定
    def crosses(idx, plane):
        _, lo, hi, _ = nodes_info[idx]
        return lo[ax] < plane < hi[ax]

    global_idx = [k for k in range(len(nodes_info))
                  if all(crosses(k, p) for p in planes)]
    global_set = {nodes_info[k][0] for k in global_idx}
    print(f"mesh 节点 {len(nodes_info)}，其中跨越全部分割面的共享件 {len(global_set)} 个：")
    for k in global_idx:
        ni, lo, hi, _ = nodes_info[k]
        print(f"  - {src['nodes'][ni].get('name','')!r}  {args.axis}: {lo[ax]:.1f}..{hi[ax]:.1f}")

    prefix = args.prefix or os.path.splitext(os.path.basename(args.input))[0]
    os.makedirs(args.outdir, exist_ok=True)

    for b in range(nband):
        lo_p = planes[b - 1] if b > 0 else None
        hi_p = planes[b] if b < nband - 1 else None
        builder = BandBuilder(src)
        builder.src_bin = src_bin

        for k, (ni, nlo, nhi, c) in enumerate(nodes_info):
            if ni in global_set:
                continue
            if (lo_p is None or c > lo_p) and (hi_p is None or c <= hi_p):
                builder.add_verbatim_node(ni, world[ni])

        clipped = 0
        for k in global_idx:
            ni = nodes_info[k][0]
            if builder.add_clipped_node(ni, world[ni], lo_p, hi_p):
                clipped += 1

        name = f"{prefix}_模块{b + 1}.glb"
        path = os.path.join(args.outdir, name)
        builder.finish(path)
        size = os.path.getsize(path) / 1048576
        print(f"[{b+1}/{nband}] {name}: 节点 {len(builder.gltf['nodes'])} "
              f"(裁剪件 {clipped})  几何体 {len(builder.gltf['meshes'])}  {size:.1f} MB")


if __name__ == "__main__":
    main()