#!/usr/bin/env python3
"""
服务端 GPU：PyTorch3D point_face_distance（CUDA）计算点到各零件 STL 三角面的最近距离，
与 three-mesh-bvh 语义一致（点到三角形）。无 CUDA / 无依赖时退出码 2，由 Node 回退 CPU。

依赖（需与 CUDA 版本匹配自行安装）:
  pip install torch trimesh numpy
  pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable"
  # 或官方 wheel: https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md
"""
from __future__ import annotations

import argparse
import json
import os
import sys


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--points", required=True, help="raw float32 little-endian, N*3 bytes")
    ap.add_argument("--manifest-dir", required=True, help="含 manifest.json 与各 .stl")
    ap.add_argument("--out-dir", required=True, help="写入 assignments.int32 / distances.f32 / result.json")
    ap.add_argument("--chunk", type=int, default=65536, help="每批点数，显存不足时调小")
    args = ap.parse_args()

    try:
        import numpy as np
        import torch
        import trimesh
        from pytorch3d.loss.point_mesh_distance import point_face_distance
    except ImportError as e:  # noqa: PERF203
        print(json.dumps({"ok": False, "error": f"import: {e}", "fallback": True}), flush=True)
        sys.exit(2)

    if not torch.cuda.is_available():
        print(
            json.dumps({"ok": False, "error": "no CUDA device", "fallback": True}),
            flush=True,
        )
        sys.exit(2)

    device = torch.device("cuda:0")
    pts_path = args.points
    man_dir = args.manifest_dir.rstrip(os.sep)
    out_dir = args.out_dir
    os.makedirs(out_dir, exist_ok=True)

    with open(pts_path, "rb") as f:
        raw = f.read()

    n_float = len(raw) // 4
    if n_float % 3 != 0:
        print(json.dumps({"ok": False, "error": "points file size not multiple of 12"}), flush=True)
        sys.exit(1)
    pts_np = np.frombuffer(raw, dtype=np.float32).reshape(-1, 3)
    n_pt = int(pts_np.shape[0])

    man_path = os.path.join(man_dir, "manifest.json")
    with open(man_path, encoding="utf-8") as f:
        man = json.load(f)
    parts = man.get("parts") or []
    if not parts:
        print(json.dumps({"ok": False, "error": "empty manifest.parts"}), flush=True)
        sys.exit(1)

    stl_infos: list[tuple[torch.Tensor, torch.Tensor]] = []
    for p in parts:
        rel = p.get("stlRelative") or p.get("stl_relative")
        if not rel:
            print(json.dumps({"ok": False, "error": "part missing stlRelative"}), flush=True)
            sys.exit(1)
        stl_path = os.path.join(man_dir, rel)
        if not os.path.isfile(stl_path):
            print(json.dumps({"ok": False, "error": f"missing {stl_path}"}), flush=True)
            sys.exit(1)
        mesh = trimesh.load(stl_path, force="mesh")
        if hasattr(mesh, "geometry"):
            mesh = trimesh.util.concatenate(tuple(mesh.geometry.values()))
        verts = torch.tensor(mesh.vertices, dtype=torch.float32, device=device)
        faces = torch.tensor(mesh.faces, dtype=torch.int64, device=device)
        if faces.numel() == 0:
            print(json.dumps({"ok": False, "error": f"empty mesh {rel}"}), flush=True)
            sys.exit(1)
        stl_infos.append((verts, faces))

    n_parts = len(stl_infos)
    chunk = max(1024, int(args.chunk))

    # 增量 argmin，避免 (N × 零件数) 距离矩阵占满显存
    best_d = torch.full((n_pt,), float("inf"), dtype=torch.float32, device=device)
    best_i = torch.zeros((n_pt,), dtype=torch.int64, device=device)
    p_first = torch.tensor([0], dtype=torch.int64, device=device)
    t_first = torch.tensor([0], dtype=torch.int64, device=device)

    for pi in range(n_parts):
        verts, faces = stl_infos[pi]
        tris = verts[faces]  # (F, 3, 3)
        for s in range(0, n_pt, chunk):
            end = min(s + chunk, n_pt)
            block = torch.tensor(pts_np[s:end], dtype=torch.float32, device=device)
            mp = end - s
            d_sq = point_face_distance(
                block,
                p_first,
                tris,
                t_first,
                mp,
                5e-3,
            )
            d = torch.sqrt(torch.clamp(d_sq, min=0.0))
            sl = slice(s, end)
            improve = d < best_d[sl]
            best_d[sl] = torch.where(improve, d, best_d[sl])
            best_i[sl] = torch.where(
                improve,
                torch.full((mp,), pi, dtype=torch.int64, device=device),
                best_i[sl],
            )

    assign = best_i.to(torch.int32).cpu().numpy()
    min_d = best_d.to(torch.float32).cpu().numpy()

    a_path = os.path.join(out_dir, "assignments.int32")
    d_path = os.path.join(out_dir, "distances.f32")
    assign.tofile(a_path)
    min_d.tofile(d_path)

    meta = {
        "ok": True,
        "n": int(n_pt),
        "parts": n_parts,
        "assignmentsFile": a_path,
        "distancesFile": d_path,
        "device": str(device),
    }
    with open(os.path.join(out_dir, "result.json"), "w", encoding="utf-8") as f:
        json.dump(meta, f)
    print(json.dumps(meta), flush=True)


if __name__ == "__main__":
    main()