#!/usr/bin/env python3
"""按"查看器标签"缩放 GLB 中某个零件的几何体（绕其包围盒中心）。

默认：长度 = 平面内最长的轴，宽度 = 第二长的轴。缩放以中心为基准，
所以 width=3 表示两边各 +100%（总宽变 3 倍），length=1.5 表示两边各 +25%。

用法：
  python3 viewer/tools/scale_part.py --input a.glb --label 区域_143 --width 3 --length 1.5 --inplace
"""
import argparse, os, sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from split_glb import read_glb, write_glb, NCOMP, COMP_DTYPE, COMP_SIZE
import numpy as np


def viewer_labels(gltf):
    used, labels = {}, {}
    for i, n in enumerate(gltf["nodes"]):
        nm = n.get("name", "")
        if not nm:
            continue
        if nm in used:
            used[nm] += 1
            labels[i] = f"{nm}_{used[nm]}"
        else:
            used[nm] = 0
            labels[i] = nm
    return labels


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True)
    ap.add_argument("--label", required=True, help="查看器标签，如 区域_143")
    ap.add_argument("--width", type=float, default=1.0, help="宽度方向缩放倍数（最短的平面轴）")
    ap.add_argument("--length", type=float, default=1.0, help="长度方向缩放倍数（最长的平面轴）")
    ap.add_argument("--output")
    ap.add_argument("--inplace", action="store_true")
    args = ap.parse_args()

    gltf, bin_ = read_glb(args.input)
    labels = viewer_labels(gltf)
    node_idx = next((i for i, l in labels.items() if l == args.label), None)
    if node_idx is None:
        print(f"未找到标签 {args.label!r}")
        return
    node = gltf["nodes"][node_idx]
    mesh = gltf["meshes"][node["mesh"]]

    changed = 0
    for pr in mesh["primitives"]:
        ai = pr["attributes"]["POSITION"]
        acc = gltf["accessors"][ai]
        if acc["componentType"] != 5126 or acc["type"] != "VEC3":
            print("仅支持 float VEC3 的 POSITION")
            return
        bv = gltf["bufferViews"][acc["bufferView"]]
        base = bv.get("byteOffset", 0) + acc.get("byteOffset", 0)
        n = acc["count"]
        arr = np.frombuffer(bin_, dtype="<f4", count=n * 3, offset=base).reshape(n, 3).copy()

        lo, hi = arr.min(0), arr.max(0)
        center = (lo + hi) / 2
        ext = hi - lo
        order = np.argsort(ext)[::-1]          # 最长的两个平面轴
        length_ax, width_ax = int(order[0]), int(order[1])
        print(f"尺寸 {ext[0]:.2f} x {ext[1]:.2f} x {ext[2]:.2f}")
        print(f"长度轴={'XYZ'[length_ax]} (x{args.length})   宽度轴={'XYZ'[width_ax]} (x{args.width})")

        arr[:, length_ax] = center[length_ax] + (arr[:, length_ax] - center[length_ax]) * args.length
        arr[:, width_ax] = center[width_ax] + (arr[:, width_ax] - center[width_ax]) * args.width

        # 回写字节
        buf = bytearray(bin_)
        buf[base:base + arr.nbytes] = np.ascontiguousarray(arr, dtype="<f4").tobytes()
        bin_ = bytes(buf)

        acc["min"] = arr.min(0).tolist()
        acc["max"] = arr.max(0).tolist()
        changed += 1
        print(f"新尺寸 {arr.max(0)[0]-arr.min(0)[0]:.2f} x {arr.max(0)[1]-arr.min(0)[1]:.2f} "
              f"x {arr.max(0)[2]-arr.min(0)[2]:.2f}")

    if not changed:
        print("没有可缩放的图元")
        return
    out = args.output or (args.input if args.inplace else args.input + ".scaled.glb")
    if args.inplace:
        tmp = out + ".tmp"
        write_glb(tmp, gltf, bin_)
        os.replace(tmp, out)
    else:
        write_glb(out, gltf, bin_)
    print(f"已写入 {out} ({os.path.getsize(out)/1048576:.1f} MB)")


if __name__ == "__main__":
    main()