"""STEP → USD exporter for Schaeffler Automat.

Reads a STEP file via OCP/XCAF (preserving part names + embedded colors),
tessellates with BRepMesh, builds a USD stage with one UsdGeomMesh per leaf
part, and writes a .usd file.

Coordinate system: OCC is mm Z-up. USD stage is authored in mm Y-up
(matching glTF / Blender convention). metersPerUnit=0.001 is set so Blender
handles the mm→m conversion on import — no explicit scaling applied here.

Usage:
    python3 export_step_to_usd.py \\
        --step_path /path/to/file.stp \\
        --output_path /path/to/output.usd \\
        [--linear_deflection 0.03] \\
        [--angular_deflection 0.05] \\
        [--color_map '{"Ring": "#4C9BE8"}'] \\
        [--sharp_threshold 20.0] \\
        [--cad_file_id uuid]

Exit 0 on success, exit 1 on failure.
Prints MANIFEST_JSON: {...} to stdout before exit.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
import re
import sys
import traceback
from pathlib import Path


# ── CLI ───────────────────────────────────────────────────────────────────────

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--step_path", required=True)
    p.add_argument("--output_path", required=True)
    p.add_argument("--linear_deflection", type=float, default=0.03)
    p.add_argument("--angular_deflection", type=float, default=0.05)
    p.add_argument("--color_map", default="{}")
    p.add_argument("--sharp_threshold", type=float, default=20.0)
    p.add_argument("--cad_file_id", default="")
    return p.parse_args()


# ── Part key generation ───────────────────────────────────────────────────────

_AF_RE = re.compile(r'_AF\d+$', re.IGNORECASE)


def _generate_part_key(xcaf_path: str, source_name: str, existing_keys: set) -> str:
    """Deterministic slug, max 64 chars, unique within assembly."""
    base = _AF_RE.sub('', source_name) if source_name else ''
    base = re.sub(r'([a-z])([A-Z])', r'\1_\2', base)
    slug = re.sub(r'[^a-z0-9]+', '_', base.lower()).strip('_')
    if not slug:
        slug = f"part_{hashlib.sha256(xcaf_path.encode()).hexdigest()[:8]}"
    slug = slug[:50]
    key = slug
    n = 2
    while key in existing_keys:
        key = f"{slug}_{n}"
        n += 1
    existing_keys.add(key)
    return key


# ── Color helpers ─────────────────────────────────────────────────────────────

PALETTE_HEX = [
    "#4C9BE8", "#E85B4C", "#4CBE72", "#E8A84C", "#A04CE8",
    "#4CD4E8", "#E84CA8", "#7EC850", "#E86B30", "#5088C8",
]


def _occ_color_to_hex(occ_color) -> str:
    r = int(occ_color.Red() * 255)
    g = int(occ_color.Green() * 255)
    b = int(occ_color.Blue() * 255)
    return f"#{r:02X}{g:02X}{b:02X}"


def _hex_to_occ_color(hex_color: str):
    from OCP.Quantity import Quantity_Color, Quantity_TOC_RGB
    h = hex_color.lstrip("#")
    if len(h) < 6:
        return Quantity_Color(0.7, 0.7, 0.7, Quantity_TOC_RGB)
    return Quantity_Color(
        int(h[0:2], 16) / 255.0,
        int(h[2:4], 16) / 255.0,
        int(h[4:6], 16) / 255.0,
        Quantity_TOC_RGB,
    )


def _hex_to_rgb01(hex_color: str) -> tuple:
    h = hex_color.lstrip('#')
    if len(h) < 6:
        return (0.7, 0.7, 0.7)
    return (int(h[0:2], 16) / 255.0, int(h[2:4], 16) / 255.0, int(h[4:6], 16) / 255.0)


def _get_shape_color(color_tool, shape) -> str | None:
    """Return hex color for an OCC shape (surface color preferred)."""
    from OCP.Quantity import Quantity_Color
    try:
        from OCP.XCAFDoc import XCAFDoc_ColorSurf as _SURF
        from OCP.XCAFDoc import XCAFDoc_ColorGen as _GEN
    except ImportError:
        _SURF = 1
        _GEN = 0
    occ_color = Quantity_Color()
    if color_tool.GetColor(shape, _SURF, occ_color):
        return _occ_color_to_hex(occ_color)
    if color_tool.GetColor(shape, _GEN, occ_color):
        return _occ_color_to_hex(occ_color)
    return None


# ── XCAF color application ────────────────────────────────────────────────────

def _apply_color_map(shape_tool, color_tool, free_labels, color_map: dict) -> None:
    from OCP.TDF import TDF_LabelSequence
    from OCP.TDataStd import TDataStd_Name
    from OCP.XCAFDoc import XCAFDoc_ShapeTool
    try:
        from OCP.XCAFDoc import XCAFDoc_ColorSurf as _SURF
    except ImportError:
        _SURF = 1

    def _visit(label) -> None:
        name_attr = TDataStd_Name()
        name = ""
        if label.FindAttribute(TDataStd_Name.GetID_s(), name_attr):
            name = name_attr.Get().ToExtString()
        if name:
            for part_name, hex_color in color_map.items():
                if part_name.lower() in name.lower() or name.lower() in part_name.lower():
                    color_tool.SetColor(label, _hex_to_occ_color(hex_color), _SURF)
                    break
        components = TDF_LabelSequence()
        XCAFDoc_ShapeTool.GetComponents_s(label, components)
        for i in range(1, components.Length() + 1):
            _visit(components.Value(i))

    for i in range(1, free_labels.Length() + 1):
        _visit(free_labels.Value(i))


def _apply_palette_colors(shape_tool, color_tool, free_labels) -> None:
    from OCP.TDF import TDF_LabelSequence
    from OCP.XCAFDoc import XCAFDoc_ShapeTool
    try:
        from OCP.XCAFDoc import XCAFDoc_ColorSurf as _SURF
    except ImportError:
        _SURF = 1

    leaves: list = []

    def _collect(label) -> None:
        components = TDF_LabelSequence()
        XCAFDoc_ShapeTool.GetComponents_s(label, components)
        if components.Length() == 0:
            leaves.append(label)
        else:
            for i in range(1, components.Length() + 1):
                _collect(components.Value(i))

    for i in range(1, free_labels.Length() + 1):
        _collect(free_labels.Value(i))

    for idx, label in enumerate(leaves):
        color_tool.SetColor(label, _hex_to_occ_color(PALETTE_HEX[idx % len(PALETTE_HEX)]), _SURF)


# ── Sharp edge extraction (inlined from export_step_to_gltf.py) ──────────────

def _extract_sharp_edge_pairs(shape, sharp_threshold_deg: float = 20.0) -> list:
    """Extract sharp B-rep edges as dense curve-sample segment pairs (mm, Z-up).

    Ported from export_step_to_gltf.py to avoid importing that module
    (its top-level code runs main() on import).
    """
    from OCP.TopTools import TopTools_IndexedDataMapOfShapeListOfShape
    from OCP.TopExp import TopExp as _TopExp
    from OCP.TopAbs import TopAbs_EDGE, TopAbs_FACE, TopAbs_FORWARD
    from OCP.TopoDS import TopoDS as _TopoDS
    from OCP.BRepAdaptor import BRepAdaptor_Surface, BRepAdaptor_Curve2d, BRepAdaptor_Curve
    from OCP.BRepLProp import BRepLProp_SLProps
    from OCP.GCPnts import GCPnts_UniformAbscissa

    edge_face_map = TopTools_IndexedDataMapOfShapeListOfShape()
    _TopExp.MapShapesAndAncestors_s(shape, TopAbs_EDGE, TopAbs_FACE, edge_face_map)

    sharp_pairs: list = []
    n_checked = 0
    n_sharp = 0
    SAMPLE_STEP_MM = 0.3

    for i in range(1, edge_face_map.Extent() + 1):
        edge_shape = edge_face_map.FindKey(i)
        faces = edge_face_map.FindFromIndex(i)
        n_checked += 1
        if faces.Size() < 2:
            continue
        face_shapes = list(faces)
        if len(face_shapes) < 2:
            continue
        try:
            edge = _TopoDS.Edge_s(edge_shape)
            face1 = _TopoDS.Face_s(face_shapes[0])
            face2 = _TopoDS.Face_s(face_shapes[1])

            c2d_1 = BRepAdaptor_Curve2d(edge, face1)
            uv1 = c2d_1.Value((c2d_1.FirstParameter() + c2d_1.LastParameter()) / 2.0)
            surf1 = BRepAdaptor_Surface(face1)
            props1 = BRepLProp_SLProps(surf1, uv1.X(), uv1.Y(), 1, 1e-6)
            if not props1.IsNormalDefined():
                continue
            n1 = props1.Normal()
            if face1.Orientation() != TopAbs_FORWARD:
                n1.Reverse()

            c2d_2 = BRepAdaptor_Curve2d(edge, face2)
            uv2 = c2d_2.Value((c2d_2.FirstParameter() + c2d_2.LastParameter()) / 2.0)
            surf2 = BRepAdaptor_Surface(face2)
            props2 = BRepLProp_SLProps(surf2, uv2.X(), uv2.Y(), 1, 1e-6)
            if not props2.IsNormalDefined():
                continue
            n2 = props2.Normal()
            if face2.Orientation() != TopAbs_FORWARD:
                n2.Reverse()

            cos_angle = max(-1.0, min(1.0, n1.Dot(n2)))
            angle_deg = math.degrees(math.acos(cos_angle))
            if angle_deg > 90.0:
                angle_deg = 180.0 - angle_deg
            if angle_deg <= sharp_threshold_deg:
                continue

            n_sharp += 1
            pts: list = []
            try:
                curve3d = BRepAdaptor_Curve(edge)
                f_param = curve3d.FirstParameter()
                l_param = curve3d.LastParameter()
                if math.isfinite(f_param) and math.isfinite(l_param):
                    sampler = GCPnts_UniformAbscissa()
                    sampler.Initialize(curve3d, SAMPLE_STEP_MM, 1e-6)
                    if sampler.IsDone() and sampler.NbPoints() >= 2:
                        for j in range(1, sampler.NbPoints() + 1):
                            p = curve3d.Value(sampler.Parameter(j))
                            pts.append([round(p.X(), 4), round(p.Y(), 4), round(p.Z(), 4)])
            except Exception:
                pts = []

            if len(pts) < 2:
                continue
            for k in range(len(pts) - 1):
                sharp_pairs.append([pts[k], pts[k + 1]])
        except Exception:
            continue

    print(
        f"Sharp edge extraction: {n_checked} edges checked, "
        f"{n_sharp} sharp (>{sharp_threshold_deg:.0f}°), "
        f"{len(sharp_pairs)} segment pairs total"
    )
    return sharp_pairs


def _extract_seam_edge_pairs(shape) -> list:
    """Extract seam edges (periodic-surface boundary edges) as segment pairs (mm, Z-up).

    Seam edges are detected via BRep_Tool.IsClosed_s(edge) — edges that are
    topologically closed (start == end vertex). This includes the UV seams of
    periodic surfaces (cylinders, cones, spheres) but also full circles on flat
    faces and bore rims.
    TODO: Use ShapeAnalysis_Edge().IsSeam(edge, face) to restrict to true UV seams
          when UV-unwrapped texture mapping is needed (future phase).
    """
    from OCP.BRep import BRep_Tool
    from OCP.TopExp import TopExp_Explorer
    from OCP.TopAbs import TopAbs_EDGE
    from OCP.BRepAdaptor import BRepAdaptor_Curve
    from OCP.GCPnts import GCPnts_UniformAbscissa

    seam_pairs: list = []
    n_seam = 0
    exp = TopExp_Explorer(shape, TopAbs_EDGE)
    while exp.More():
        edge = exp.Current()
        exp.Next()
        if not BRep_Tool.IsClosed_s(edge):
            continue
        try:
            curve = BRepAdaptor_Curve(edge)
            # Use arc-length step (0.3 mm) matching the sharp edge sampler,
            # so segments are short enough for _world_to_index_pairs (tol=0.5 mm).
            sampler = GCPnts_UniformAbscissa()
            sampler.Initialize(curve, 0.3, 1e-6)
            if not sampler.IsDone() or sampler.NbPoints() < 2:
                continue
            pts = []
            for i in range(1, sampler.NbPoints() + 1):
                p = curve.Value(sampler.Parameter(i))
                pts.append([p.X(), p.Y(), p.Z()])
            for k in range(len(pts) - 1):
                seam_pairs.append([pts[k], pts[k + 1]])
            n_seam += 1
        except Exception:
            continue
    print(f"Seam edge extraction: {n_seam} seam edges, {len(seam_pairs)} segment pairs total")
    return seam_pairs


# ── XCAF traversal ────────────────────────────────────────────────────────────

def _traverse_xcaf(shape_tool, color_tool, label, path_prefix, existing_keys, depth=0):
    """Yield one dict per leaf shape in the XCAF hierarchy.

    Transform composition: `GetShape_s(reference_label)` returns the shape with
    the reference's own location already composed in. For standard Schaeffler flat
    assemblies (1–2 levels deep) this is correct. Deeply nested sub-assembly
    transforms (3+ levels) accumulate naturally because each recursive call
    receives a component label from the *referred* definition, so each level's
    location is composed by the next GetShape_s call.
    """
    from OCP.TDF import TDF_LabelSequence, TDF_Label
    from OCP.TDataStd import TDataStd_Name
    from OCP.XCAFDoc import XCAFDoc_ShapeTool

    name_attr = TDataStd_Name()
    source_name = ""
    if label.FindAttribute(TDataStd_Name.GetID_s(), name_attr):
        source_name = name_attr.Get().ToExtString()

    xcaf_path = (f"{path_prefix}/{source_name}" if source_name
                 else f"{path_prefix}/unnamed_{depth}")

    # Follow references to get the definition label (for sub-assembly detection)
    actual_label = label
    if XCAFDoc_ShapeTool.IsReference_s(label):
        ref_label = TDF_Label()
        if XCAFDoc_ShapeTool.GetReferredShape_s(label, ref_label):
            actual_label = ref_label

    components = TDF_LabelSequence()
    XCAFDoc_ShapeTool.GetComponents_s(actual_label, components)

    if components.Length() == 0:
        shape = shape_tool.GetShape_s(label)
        if shape.IsNull():
            shape = shape_tool.GetShape_s(actual_label)
        if shape.IsNull():
            return

        part_key = _generate_part_key(xcaf_path, source_name, existing_keys)
        color = _get_shape_color(color_tool, shape)

        yield {
            'shape': shape,
            'source_name': source_name,
            'xcaf_path': xcaf_path,
            'part_key': part_key,
            'color': color,
        }
    else:
        for i in range(1, components.Length() + 1):
            yield from _traverse_xcaf(
                shape_tool, color_tool, components.Value(i),
                xcaf_path, existing_keys, depth + 1,
            )


# ── Mesh geometry extraction ──────────────────────────────────────────────────

def _extract_mesh(shape) -> tuple[list, list]:
    """Return (vertices, triangles) from a tessellated OCC shape.

    Vertices are in OCC space (mm, Z-up).
    Triangles are 0-based index triples.
    """
    from OCP.TopExp import TopExp_Explorer
    from OCP.TopAbs import TopAbs_FACE, TopAbs_REVERSED
    from OCP.TopoDS import TopoDS
    from OCP.BRep import BRep_Tool
    from OCP.TopLoc import TopLoc_Location

    vertices: list = []
    triangles: list = []
    v_offset = 0

    shape_trsf = shape.Location().Transformation()
    shape_has_loc = not shape.Location().IsIdentity()

    exp = TopExp_Explorer(shape, TopAbs_FACE)
    while exp.More():
        face = TopoDS.Face_s(exp.Current())
        face_loc = TopLoc_Location()
        tri = BRep_Tool.Triangulation_s(face, face_loc)

        if tri is not None and tri.NbNodes() > 0:
            reversed_face = (face.Orientation() == TopAbs_REVERSED)
            face_has_loc = not face_loc.IsIdentity()

            for i in range(1, tri.NbNodes() + 1):
                node = tri.Node(i)
                if face_has_loc:
                    # face_loc from BRep_Tool.Triangulation_s already encodes the
                    # instance placement for compound-tessellated shapes — applying
                    # shape_loc on top would double-transform every vertex.
                    node = node.Transformed(face_loc.Transformation())
                elif shape_has_loc:
                    # Only fall back to shape_loc when face_loc is identity (e.g.
                    # shapes tessellated individually rather than as a compound).
                    node = node.Transformed(shape_trsf)
                vertices.append((node.X(), node.Y(), node.Z()))

            for i in range(1, tri.NbTriangles() + 1):
                n1, n2, n3 = tri.Triangle(i).Get()
                v0 = n1 - 1 + v_offset
                v1 = n2 - 1 + v_offset
                v2 = n3 - 1 + v_offset
                triangles.append((v0, v2, v1) if reversed_face else (v0, v1, v2))

            v_offset += tri.NbNodes()

        exp.Next()

    return vertices, triangles


# ── Index-space sharp edge mapping ────────────────────────────────────────────

def _world_to_index_pairs(vertices: list, world_pairs: list, tol: float = 0.5) -> list:
    """Map world-space (mm, Z-up) segment pairs → local vertex index pairs."""
    def _k(x, y, z):
        return (round(x / tol) * tol, round(y / tol) * tol, round(z / tol) * tol)

    coord_map: dict = {}
    for idx, (x, y, z) in enumerate(vertices):
        k = _k(x, y, z)
        if k not in coord_map:
            coord_map[k] = idx

    result = []
    for p0, p1 in world_pairs:
        i0 = coord_map.get(_k(p0[0], p0[1], p0[2]))
        i1 = coord_map.get(_k(p1[0], p1[1], p1[2]))
        if i0 is not None and i1 is not None and i0 != i1:
            result.append((i0, i1))
    return result


# ── USD prim name sanitizer ───────────────────────────────────────────────────

def _prim_name(name: str) -> str:
    safe = re.sub(r'[^A-Za-z0-9_]', '_', name)
    if safe and safe[0].isdigit():
        safe = f"_{safe}"
    return safe or "unnamed"


# ── Main ──────────────────────────────────────────────────────────────────────

def main() -> None:
    args = parse_args()
    color_map: dict = json.loads(args.color_map)

    step_path = Path(args.step_path)
    output_path = Path(args.output_path)

    if not step_path.exists():
        print(f"ERROR: STEP file not found: {step_path}", file=sys.stderr)
        sys.exit(1)

    output_path.parent.mkdir(parents=True, exist_ok=True)

    # ── OCC / XCAF imports ────────────────────────────────────────────────────
    from OCP.STEPCAFControl import STEPCAFControl_Reader
    from OCP.TDocStd import TDocStd_Document
    from OCP.XCAFApp import XCAFApp_Application
    from OCP.XCAFDoc import XCAFDoc_DocumentTool
    from OCP.TCollection import TCollection_ExtendedString
    from OCP.TDF import TDF_LabelSequence
    from OCP.BRepMesh import BRepMesh_IncrementalMesh
    from OCP.IFSelect import IFSelect_RetDone

    # ── pxr imports ───────────────────────────────────────────────────────────
    from pxr import Usd, UsdGeom, UsdShade, Sdf, Vt, Gf

    # ── Read STEP ─────────────────────────────────────────────────────────────
    app = XCAFApp_Application.GetApplication_s()
    doc = TDocStd_Document(TCollection_ExtendedString("MDTV-CAF"))
    app.InitDocument(doc)

    reader = STEPCAFControl_Reader()
    reader.SetNameMode(True)
    reader.SetColorMode(True)
    reader.SetLayerMode(True)
    status = reader.ReadFile(str(step_path))
    if status != IFSelect_RetDone:
        print(f"ERROR: STEPCAFControl_Reader failed (status={status})", file=sys.stderr)
        sys.exit(1)
    reader.Transfer(doc)

    shape_tool = XCAFDoc_DocumentTool.ShapeTool_s(doc.Main())
    color_tool = XCAFDoc_DocumentTool.ColorTool_s(doc.Main())

    free_labels = TDF_LabelSequence()
    shape_tool.GetFreeShapes(free_labels)
    print(
        f"Found {free_labels.Length()} root shape(s), tessellating "
        f"(linear={args.linear_deflection}mm, angular={args.angular_deflection}rad) …"
    )

    # ── Tessellate ────────────────────────────────────────────────────────────
    for i in range(1, free_labels.Length() + 1):
        shape = shape_tool.GetShape_s(free_labels.Value(i))
        if not shape.IsNull():
            BRepMesh_IncrementalMesh(
                shape, args.linear_deflection, False, args.angular_deflection, True
            )
    print("Tessellation complete.")

    # ── Sharp edge pairs (world-space mm, Z-up) ───────────────────────────────
    sharp_pairs_mm: list = []
    try:
        for i in range(1, free_labels.Length() + 1):
            root_shape = shape_tool.GetShape_s(free_labels.Value(i))
            if not root_shape.IsNull():
                sharp_pairs_mm.extend(
                    _extract_sharp_edge_pairs(root_shape, args.sharp_threshold)
                )
        print(f"Total sharp segment pairs: {len(sharp_pairs_mm)}")
    except Exception as exc:
        print(f"WARNING: sharp edge extraction failed (non-fatal): {exc}", file=sys.stderr)

    # ── Seam edge pairs (world-space mm, Z-up) ────────────────────────────────
    seam_pairs_mm: list = []
    try:
        for i in range(1, free_labels.Length() + 1):
            root_shape = shape_tool.GetShape_s(free_labels.Value(i))
            if not root_shape.IsNull():
                seam_pairs_mm.extend(_extract_seam_edge_pairs(root_shape))
        print(f"Total seam segment pairs: {len(seam_pairs_mm)}")
    except Exception as exc:
        print(f"WARNING: seam edge extraction failed (non-fatal): {exc}", file=sys.stderr)

    # ── Apply colors ──────────────────────────────────────────────────────────
    if color_map:
        try:
            _apply_color_map(shape_tool, color_tool, free_labels, color_map)
            print(f"Applied color_map ({len(color_map)} entries)")
        except Exception as exc:
            print(f"WARNING: color_map application failed (non-fatal): {exc}", file=sys.stderr)
    else:
        try:
            _apply_palette_colors(shape_tool, color_tool, free_labels)
            print("Applied palette colors")
        except Exception as exc:
            print(f"WARNING: palette colors failed (non-fatal): {exc}", file=sys.stderr)

    # ── Create USD stage ──────────────────────────────────────────────────────
    stage = Usd.Stage.CreateNew(str(output_path))
    UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.y)
    UsdGeom.SetStageMetersPerUnit(stage, 0.001)  # mm; Blender handles m conversion on import

    root_prim = UsdGeom.Xform.Define(stage, "/Root")
    stage.SetDefaultPrim(root_prim.GetPrim())
    UsdGeom.Xform.Define(stage, "/Root/Assembly")
    stage.DefinePrim("/Root/Looks", "Scope")

    # ── Walk XCAF tree → author USD prims ─────────────────────────────────────
    existing_keys: set = set()
    manifest_parts: list = []
    n_parts = 0
    n_empty = 0

    for root_idx in range(1, free_labels.Length() + 1):
        root_label = free_labels.Value(root_idx)

        from OCP.TDataStd import TDataStd_Name as _Name
        _na = _Name()
        root_src = ""
        if root_label.FindAttribute(_Name.GetID_s(), _na):
            root_src = _na.Get().ToExtString()
        node_name = _prim_name(root_src or f"Root{root_idx}")
        node_path = f"/Root/Assembly/{node_name}"
        UsdGeom.Xform.Define(stage, node_path)

        for part in _traverse_xcaf(shape_tool, color_tool, root_label, "", existing_keys):
            source_name = part['source_name']
            part_key = part['part_key']
            hex_color = part['color']
            shape = part['shape']
            xcaf_path = part['xcaf_path']

            # color_map override (substring match)
            for map_name, map_hex in color_map.items():
                if (map_name.lower() in source_name.lower()
                        or source_name.lower() in map_name.lower()):
                    hex_color = map_hex
                    break
            if not hex_color:
                hex_color = PALETTE_HEX[n_parts % len(PALETTE_HEX)]

            vertices, triangles = _extract_mesh(shape)
            if not vertices or not triangles:
                n_empty += 1
                continue

            part_path = f"{node_path}/{part_key}"
            # Name the Mesh prim after part_key so Blender imports it with the
            # part name directly (Blender collapses single-child Xform+Mesh into
            # just the Mesh object, using the mesh prim's leaf name as object name).
            mesh_path = f"{part_path}/{part_key}"

            # ── Xform prim ────────────────────────────────────────────────
            xform = UsdGeom.Xform.Define(stage, part_path)
            prim = xform.GetPrim()
            prim.SetCustomDataByKey("schaeffler:partKey", part_key)
            prim.SetCustomDataByKey("schaeffler:sourceName", source_name)
            prim.SetCustomDataByKey("schaeffler:sourceAssemblyPath", xcaf_path)
            prim.SetCustomDataByKey("schaeffler:sourceColor", hex_color)
            prim.SetCustomDataByKey("schaeffler:tessellation:linearDeflectionMm",
                                    args.linear_deflection)
            prim.SetCustomDataByKey("schaeffler:tessellation:angularDeflectionRad",
                                    args.angular_deflection)
            if args.cad_file_id:
                prim.SetCustomDataByKey("schaeffler:cadFileId", args.cad_file_id)

            # ── UsdGeomMesh ───────────────────────────────────────────────
            mesh = UsdGeom.Mesh.Define(stage, mesh_path)
            mesh.CreateSubdivisionSchemeAttr(UsdGeom.Tokens.none)

            # OCC (X, Y, Z) mm Z-up → USD (X, -Z, Y) mm Y-up
            mesh.CreatePointsAttr(Vt.Vec3fArray([
                Gf.Vec3f(x, -z, y) for (x, y, z) in vertices
            ]))
            mesh.CreateFaceVertexCountsAttr(Vt.IntArray([3] * len(triangles)))
            mesh.CreateFaceVertexIndicesAttr(
                Vt.IntArray([idx for tri in triangles for idx in tri])
            )
            r, g, b = _hex_to_rgb01(hex_color)
            mesh.CreateDisplayColorAttr(Vt.Vec3fArray([Gf.Vec3f(r, g, b)]))

            # ── Index-space sharp + seam edge primvars ───────────────────
            # Lookup is in OCC Z-up space; pairs are also Z-up — no swap needed.
            # Both `vertices` and `*_pairs_mm` are in OCC Z-up mm space with the
            # full per-shape location already applied — same coordinate frame required
            # by _world_to_index_pairs for the nearest-vertex lookup (tol=0.5 mm).
            primvars_api = UsdGeom.PrimvarsAPI(mesh)
            if sharp_pairs_mm:
                idx_pairs = _world_to_index_pairs(vertices, sharp_pairs_mm)
                if idx_pairs:
                    pv = primvars_api.CreatePrimvar(
                        "schaeffler:sharpEdgeVertexPairs",
                        Sdf.ValueTypeNames.Int2Array,
                        UsdGeom.Tokens.constant,
                    )
                    pv.Set(Vt.Vec2iArray([Gf.Vec2i(a, b) for a, b in idx_pairs]))
            if seam_pairs_mm:
                seam_idx_pairs = _world_to_index_pairs(vertices, seam_pairs_mm)
                if seam_idx_pairs:
                    pv_seam = primvars_api.CreatePrimvar(
                        "schaeffler:seamEdgeVertexPairs",
                        Sdf.ValueTypeNames.Int2Array,
                        UsdGeom.Tokens.constant,
                    )
                    pv_seam.Set(Vt.Vec2iArray([Gf.Vec2i(a, b) for a, b in seam_idx_pairs]))

            # ── Material placeholder + binding ────────────────────────────
            mat_name = _prim_name(source_name) if source_name else f"mat_{part_key}"
            mat_usd_path = f"/Root/Looks/{mat_name}"
            if not stage.GetPrimAtPath(mat_usd_path):
                UsdShade.Material.Define(stage, mat_usd_path)
            UsdShade.MaterialBindingAPI(mesh.GetPrim()).Bind(
                UsdShade.Material(stage.GetPrimAtPath(mat_usd_path))
            )

            manifest_parts.append({
                "part_key": part_key,
                "source_name": source_name,
                "prim_path": part_path,
            })
            n_parts += 1

    stage.Save()

    sz = output_path.stat().st_size // 1024 if output_path.exists() else 0
    print(f"USD exported: {output_path.name} ({sz} KB), "
          f"{n_parts} parts, {n_empty} empty shapes skipped")

    # ── Stdout manifest (one line, parsed by Celery task) ─────────────────────
    print(f"MANIFEST_JSON: {json.dumps({'parts': manifest_parts})}")


try:
    main()
except SystemExit:
    raise
except Exception:
    traceback.print_exc()
    sys.exit(1)