"""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 mirroring the full XCAF
assembly hierarchy (intermediate Xform prims with local transforms, leaf
Mesh prims with definition-space geometry), and writes a .usd file.

Coordinate system: OCC is mm Z-up. USD stage is Z-up with a coordinate
swap Xform on /Root/Assembly: (X_occ, Y_occ, Z_occ) → (X, -Z, Y).
metersPerUnit=0.001 is set so Blender handles mm→m on import.

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] \\
        [--material_map '{"part_name": "SCHAEFFLER_010101_Steel-Bare", ...}']

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="")
    p.add_argument("--material_map", 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 + hierarchical USD authoring ──────────────────────────────

def _get_label_name(label) -> str:
    """Extract the source name string from a TDF_Label."""
    from OCP.TDataStd import TDataStd_Name
    name_attr = TDataStd_Name()
    if label.FindAttribute(TDataStd_Name.GetID_s(), name_attr):
        return name_attr.Get().ToExtString()
    return ""


def _occ_trsf_to_usd_matrix(trsf):
    """Convert an OCC gp_Trsf (column-vector) to a USD Gf.Matrix4d (row-vector).

    OCC uses p' = R·p + t (column-vector convention).
    USD uses p' = p·M (row-vector convention).
    So M = [R^T | 0; t^T | 1].
    """
    from pxr import Gf
    return Gf.Matrix4d(
        trsf.Value(1, 1), trsf.Value(2, 1), trsf.Value(3, 1), 0,
        trsf.Value(1, 2), trsf.Value(2, 2), trsf.Value(3, 2), 0,
        trsf.Value(1, 3), trsf.Value(2, 3), trsf.Value(3, 3), 0,
        trsf.Value(1, 4), trsf.Value(2, 4), trsf.Value(3, 4), 1,
    )


def _author_xcaf_to_usd(
    stage, shape_tool, color_tool, label,
    usd_parent_path, xcaf_path_prefix,
    existing_keys, mat_map_lower, color_map, args,
    manifest_parts, counters, prim_names_at_level,
    depth=0,
):
    """Recursively author USD prims mirroring the XCAF hierarchy.

    Assembly nodes → UsdGeom.Xform with local transform from component Location.
    Leaf shapes   → Xform + Mesh with definition-space geometry.
    Sharp/seam edges are extracted per-part from the definition shape.

    The local transform for each node comes from GetShape_s(label).Location(),
    which returns ONLY this label's own placement (not composed with parents).
    USD scene graph composition handles the full parent-to-leaf chain.
    """
    from OCP.TDF import TDF_LabelSequence, TDF_Label
    from OCP.XCAFDoc import XCAFDoc_ShapeTool
    from OCP.TopLoc import TopLoc_Location
    from pxr import UsdGeom, UsdShade, Sdf, Vt, Gf

    source_name = _get_label_name(label)
    xcaf_path = (f"{xcaf_path_prefix}/{source_name}" if source_name
                 else f"{xcaf_path_prefix}/unnamed_{depth}")

    # Get local transform from this label's shape Location.
    # GetShape_s(label) returns the shape with ONLY this label's own Location
    # (not composed with parent locations).
    label_shape = shape_tool.GetShape_s(label)
    if label_shape.IsNull():
        return
    local_loc = label_shape.Location()
    has_local_trsf = not local_loc.IsIdentity()

    # Follow reference to definition label
    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

    # Check for sub-components on the definition
    components = TDF_LabelSequence()
    XCAFDoc_ShapeTool.GetComponents_s(actual_label, components)

    if components.Length() > 0:
        # ── ASSEMBLY NODE ──────────────────────────────────────────────
        raw_name = _prim_name(source_name or f"asm_{depth}")
        unique_name = raw_name
        n = 2
        while unique_name in prim_names_at_level:
            unique_name = f"{raw_name}_{n}"
            n += 1
        prim_names_at_level.add(unique_name)

        xform_path = f"{usd_parent_path}/{unique_name}"
        xform = UsdGeom.Xform.Define(stage, xform_path)

        if has_local_trsf:
            xform.AddTransformOp().Set(
                _occ_trsf_to_usd_matrix(local_loc.Transformation()))

        prim = xform.GetPrim()
        prim.SetCustomDataByKey("schaeffler:sourceName", source_name)
        prim.SetCustomDataByKey("schaeffler:sourceAssemblyPath", xcaf_path)

        print(f"  {'  ' * depth}[asm] {source_name} → {xform_path}"
              f"{' (transform)' if has_local_trsf else ''}")

        child_names: set = set()
        for i in range(1, components.Length() + 1):
            _author_xcaf_to_usd(
                stage, shape_tool, color_tool, components.Value(i),
                xform_path, xcaf_path,
                existing_keys, mat_map_lower, color_map, args,
                manifest_parts, counters, child_names, depth + 1,
            )
    else:
        # ── LEAF SHAPE ─────────────────────────────────────────────────
        # Get definition shape without instance location
        def_shape = shape_tool.GetShape_s(actual_label)
        if def_shape.IsNull():
            return
        # Strip any residual location so _extract_mesh yields definition-space
        # vertices (face_loc only, no instance placement).
        bare_def = def_shape.Located(TopLoc_Location())

        part_key = _generate_part_key(xcaf_path, source_name, existing_keys)
        hex_color = _get_shape_color(color_tool, label_shape)
        if not hex_color:
            hex_color = _get_shape_color(color_tool, def_shape)

        # color_map override (substring match)
        if source_name:
            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[counters['n_parts'] % len(PALETTE_HEX)]

        # Extract mesh from definition shape (face_loc only, no instance placement)
        vertices, triangles = _extract_mesh(bare_def)
        if not vertices or not triangles:
            counters['n_empty'] += 1
            return

        # Ensure unique prim name at this level
        raw_name = _prim_name(part_key)
        unique_name = raw_name
        n = 2
        while unique_name in prim_names_at_level:
            unique_name = f"{raw_name}_{n}"
            n += 1
        prim_names_at_level.add(unique_name)

        part_path = f"{usd_parent_path}/{unique_name}"
        # Name the Mesh prim after part_key so Blender imports it with the
        # part name directly (Blender collapses single-child Xform+Mesh).
        mesh_path = f"{part_path}/{part_key}"

        # ── Xform prim with local transform ────────────────────────
        xform = UsdGeom.Xform.Define(stage, part_path)
        if has_local_trsf:
            xform.AddTransformOp().Set(
                _occ_trsf_to_usd_matrix(local_loc.Transformation()))

        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)

        # Vertices in OCC definition space (mm, Z-up).
        # The /Root/Assembly Xform carries the OCC→USD coordinate swap
        # so no per-vertex (X,-Z,Y) transformation is needed here.
        mesh.CreatePointsAttr(Vt.Vec3fArray([
            Gf.Vec3f(x, y, z) 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)]))

        # ── Material metadata on mesh prim (customData) ───────────
        mesh_prim = mesh.GetPrim()
        mesh_prim.SetCustomDataByKey("schaeffler:partKey", part_key)
        mesh_prim.SetCustomDataByKey("schaeffler:sourceName", source_name)

        canonical_mat = _lookup_material(source_name, part_key, mat_map_lower)
        if canonical_mat:
            mesh_prim.SetCustomDataByKey(
                "schaeffler:canonicalMaterialName", canonical_mat)

        primvars_api = UsdGeom.PrimvarsAPI(mesh)

        # ── Per-part sharp + seam edge primvars (definition space) ─
        try:
            sharp_pairs = _extract_sharp_edge_pairs(bare_def, args.sharp_threshold)
            if sharp_pairs:
                idx_pairs = _world_to_index_pairs(vertices, sharp_pairs)
                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]))
        except Exception as exc:
            print(f"WARNING: sharp edge extraction for {part_key}: {exc}")

        try:
            seam_pairs = _extract_seam_edge_pairs(bare_def)
            if seam_pairs:
                seam_idx_pairs = _world_to_index_pairs(vertices, seam_pairs)
                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]))
        except Exception as exc:
            print(f"WARNING: seam edge extraction for {part_key}: {exc}")

        # ── Material binding ──────────────────────────────────────
        if canonical_mat:
            mat_prim_name = _prim_name(canonical_mat)
        else:
            mat_prim_name = (_prim_name(source_name) if source_name
                            else f"mat_{part_key}")
        mat_usd_path = f"/Root/Looks/{mat_prim_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,
            "canonical_material": canonical_mat,
        })
        counters['n_parts'] += 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.

    Transform strategy: strip the shape's own Location before exploring faces
    so that face_loc from BRep_Tool.Triangulation_s is always relative to the
    shape's DEFINITION space (not contaminated by instance placement).  Then
    uniformly apply the shape's Location to every vertex.  This avoids both
    double-transform (when face_loc already includes placement) and missing-
    transform (when face_loc is identity but shape has placement).
    """
    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()

    # Strip instance placement so face exploration yields definition-space locs
    bare = shape.Located(TopLoc_Location())

    exp = TopExp_Explorer(bare, 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)
                # Step 1: face_loc — definition-space transform (face within shape)
                if face_has_loc:
                    node = node.Transformed(face_loc.Transformation())
                # Step 2: shape_loc — instance placement (shape within assembly)
                if shape_has_loc:
                    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"


# ── Material map lookup (mirrors _blender_materials.build_mat_map_lower) ─────

def _build_mat_map_lower(material_map: dict) -> dict:
    """Build a lowercased material_map with AF-stripped and slug variants.

    Same normalization as _blender_materials.build_mat_map_lower() so that
    source_name → canonical material name lookup works consistently.
    """
    mat_map_lower: dict = {}
    for k, v in material_map.items():
        kl = k.lower().strip()
        mat_map_lower[kl] = v
        # Slug variant: replace non-alphanumeric with '_' (same as _generate_part_key)
        slug_key = re.sub(r'[^a-z0-9]+', '_', kl).strip('_')
        if slug_key and slug_key != kl:
            mat_map_lower.setdefault(slug_key, v)
        # Strip OCC assembly-frame suffixes: _AF0, _AF0_1, _AF0_1_AF0, etc.
        stripped = re.sub(r'(_af\d+(_\d+)?)+$', '', kl)
        if stripped != kl:
            mat_map_lower.setdefault(stripped, v)
            slug_stripped = re.sub(r'[^a-z0-9]+', '_', stripped).strip('_')
            if slug_stripped and slug_stripped != stripped:
                mat_map_lower.setdefault(slug_stripped, v)
    return mat_map_lower


def _lookup_material(source_name: str, part_key: str, mat_map_lower: dict) -> str | None:
    """Look up canonical material name for a part, trying multiple key variants."""
    if not mat_map_lower:
        return None
    # Try source_name (lowered)
    sn = source_name.lower().strip()
    if sn in mat_map_lower:
        return mat_map_lower[sn]
    # Try AF-stripped source_name
    stripped = re.sub(r'(_af\d+(_\d+)?)+$', '', sn, flags=re.IGNORECASE)
    if stripped != sn and stripped in mat_map_lower:
        return mat_map_lower[stripped]
    # Try slug of source_name (matches part_key generation logic)
    slug = re.sub(r'[^a-z0-9]+', '_', sn).strip('_')
    if slug and slug in mat_map_lower:
        return mat_map_lower[slug]
    # Try part_key directly
    pk = part_key.lower().strip()
    if pk in mat_map_lower:
        return mat_map_lower[pk]
    # Prefix fallback: longest key that starts with or is started by part_key
    for key in sorted(mat_map_lower.keys(), key=len, reverse=True):
        if len(key) >= 5 and len(pk) >= 5 and (pk.startswith(key) or key.startswith(pk)):
            return mat_map_lower[key]
    return None


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

def main() -> None:
    args = parse_args()
    color_map: dict = json.loads(args.color_map)
    raw_material_map: dict = json.loads(args.material_map)
    mat_map_lower = _build_mat_map_lower(raw_material_map) if raw_material_map else {}
    if mat_map_lower:
        print(f"Material map: {len(raw_material_map)} entries ({len(mat_map_lower)} with variants)")

    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.")

    # ── 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.z)
    UsdGeom.SetStageMetersPerUnit(stage, 0.001)  # mm; Blender handles m conversion on import

    root_prim = UsdGeom.Xform.Define(stage, "/Root")
    stage.SetDefaultPrim(root_prim.GetPrim())

    # /Root/Assembly carries the OCC→Blender coordinate swap.
    # OCC is mm Z-up Y-forward; Blender/GLB convention is Z-up Y-backward.
    # Transform: (X_occ, Y_occ, Z_occ) → (X, -Z, Y)  (= Rx(-90°)).
    # Authored as a USD row-vector matrix on the Assembly Xform so that all
    # child XCAF transforms (authored in OCC space) are correctly composed.
    assembly_xform = UsdGeom.Xform.Define(stage, "/Root/Assembly")
    assembly_xform.AddTransformOp().Set(Gf.Matrix4d(
        1,  0,  0, 0,
        0,  0,  1, 0,
        0, -1,  0, 0,
        0,  0,  0, 1,
    ))

    stage.DefinePrim("/Root/Looks", "Scope")

    # ── Walk XCAF tree → author USD prims (hierarchical) ──────────────────
    # Sharp/seam edges are extracted per-part inside _author_xcaf_to_usd
    # (in definition space, matching definition-space mesh vertices).
    existing_keys: set = set()
    manifest_parts: list = []
    counters = {"n_parts": 0, "n_empty": 0}

    for root_idx in range(1, free_labels.Length() + 1):
        root_label = free_labels.Value(root_idx)
        root_names: set = set()
        _author_xcaf_to_usd(
            stage, shape_tool, color_tool, root_label,
            "/Root/Assembly", "",
            existing_keys, mat_map_lower, color_map, args,
            manifest_parts, counters, root_names,
        )

    n_parts = counters["n_parts"]
    n_empty = counters["n_empty"]
    stage.Save()

    sz = output_path.stat().st_size // 1024 if output_path.exists() else 0
    n_mat_assigned = sum(1 for p in manifest_parts if p.get("canonical_material"))
    print(f"USD exported: {output_path.name} ({sz} KB), "
          f"{n_parts} parts, {n_empty} empty shapes skipped, "
          f"{n_mat_assigned}/{n_parts} material primvars written")

    # ── 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)