"""OCC-native STEP → GLB export script.

Reads a STEP file via OCP/XCAF (preserving part names and embedded colors),
tessellates with BRepMesh or GMSH Frontal-Delaunay, optionally applies
per-part hex colors, and writes a binary GLB in meters (Y-up, glTF convention).

No Blender required. Uses the same OCP bindings that cadquery ships with.

Usage:
    python3 export_step_to_gltf.py \
        --step_path /path/to/file.stp \
        --output_path /path/to/output.glb \
        [--linear_deflection 0.1] \
        [--angular_deflection 0.5] \
        [--tessellation_engine occ|gmsh] \
        [--color_map '{"RingInner": "#4C9BE8", "RingOuter": "#E85B4C"}']

Exit 0 on success, exit 1 on failure.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import traceback
from pathlib import Path

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


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--step_path", required=True)
    parser.add_argument("--output_path", required=True)
    parser.add_argument(
        "--linear_deflection", type=float, default=0.1,
        help="OCC linear deflection for tessellation (mm). Smaller = finer mesh. Default 0.1",
    )
    parser.add_argument(
        "--angular_deflection", type=float, default=0.5,
        help="OCC angular deflection (radians). Default 0.5",
    )
    parser.add_argument(
        "--color_map", default="{}",
        help='JSON dict mapping part name → hex color, e.g. \'{"Ring": "#4C9BE8"}\'',
    )
    parser.add_argument(
        "--sharp_threshold", type=float, default=20.0,
        help="Dihedral angle threshold (degrees) for sharp B-rep edge detection. Default 20.0",
    )
    parser.add_argument(
        "--tessellation_engine", choices=["occ", "gmsh"], default="occ",
        help="Tessellation backend: 'occ' = BRepMesh (default), 'gmsh' = Frontal-Delaunay",
    )
    return parser.parse_args()


def _hex_to_occ_color(hex_color: str):
    """Convert '#RRGGBB' → Quantity_Color (linear float)."""
    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)
    r = int(h[0:2], 16) / 255.0
    g = int(h[2:4], 16) / 255.0
    b = int(h[4:6], 16) / 255.0
    return Quantity_Color(r, g, b, Quantity_TOC_RGB)


def _apply_color_map(shape_tool, color_tool, free_labels, color_map: dict) -> None:
    """Apply hex colors from color_map to matching shapes by name (case-insensitive substring)."""
    from OCP.TDF import TDF_LabelSequence
    from OCP.TDataStd import TDataStd_Name
    from OCP.XCAFDoc import XCAFDoc_ShapeTool

    # XCAFDoc_ColorType: XCAFDoc_ColorGen=0, XCAFDoc_ColorSurf=1, XCAFDoc_ColorCurv=2
    try:
        from OCP.XCAFDoc import XCAFDoc_ColorSurf as COLOR_SURF
    except ImportError:
        COLOR_SURF = 1  # integer fallback

    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), 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:
    """Assign palette colors to leaf shapes when no color_map is provided."""
    from OCP.TDF import TDF_LabelSequence
    from OCP.XCAFDoc import XCAFDoc_ShapeTool

    try:
        from OCP.XCAFDoc import XCAFDoc_ColorSurf as COLOR_SURF
    except ImportError:
        COLOR_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):
        occ_color = _hex_to_occ_color(PALETTE_HEX[idx % len(PALETTE_HEX)])
        color_tool.SetColor(label, occ_color, COLOR_SURF)


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

    For each edge shared by exactly 2 faces, evaluates the dihedral angle using
    PCurve-based surface normal evaluation. When the angle exceeds the threshold,
    samples the analytical 3D curve uniformly at 0.3mm intervals via
    GCPnts_UniformAbscissa. Consecutive sample pairs give the KD-tree in
    export_gltf.py enough density to find and mark the correct Blender mesh edges.

    Note: BRep_Tool.Polygon3D_s() and PolygonOnTriangulation_s() return None in
    XCAF compound context — the tessellation data is stored on component instances,
    not on the compound edges. Curve sampling bypasses this entirely.

    Args:
        shape: OCC TopoDS_Shape (tessellated with BRepMesh_IncrementalMesh)
        sharp_threshold_deg: dihedral angle threshold in degrees (default 20°)

    Returns:
        List of [[x0,y0,z0],[x1,y1,z1]] consecutive segment pairs in mm (OCC
        coordinate space, Z-up). Must be called BEFORE mm→m scaling.
    """
    import math as _math

    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 0.3mm — well below the KD-tree TOL=0.5mm in export_gltf.py.
    # Tessellation vertex spacing for default deflection is ~0.78-1.55mm, so at
    # least one consecutive sample pair will straddle each tessellation edge.
    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)
        if faces.Size() < 2:
            n_checked += 1
            continue

        face_shapes = list(faces)
        if len(face_shapes) < 2:
            n_checked += 1
            continue

        n_checked += 1
        try:
            edge = _TopoDS.Edge_s(edge_shape)
            face1 = _TopoDS.Face_s(face_shapes[0])
            face2 = _TopoDS.Face_s(face_shapes[1])

            # PCurve-based normal evaluation at edge midpoint
            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))
            # Use exterior (supplement) angle so convex and concave edges both work
            if angle_deg > 90.0:
                angle_deg = 180.0 - angle_deg

            if angle_deg <= sharp_threshold_deg:
                continue  # smooth transition — skip

            n_sharp += 1

            # Sample the analytical 3D curve at fixed arc-length intervals.
            # GCPnts_UniformAbscissa works on the exact B-rep curve regardless of
            # whether tessellation polygon data is stored on the edge or not.
            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):
                            t = sampler.Parameter(j)
                            p = curve3d.Value(t)
                            pts.append([round(p.X(), 4), round(p.Y(), 4), round(p.Z(), 4)])
            except Exception:
                pts = []

            if len(pts) < 2:
                continue

            # Consecutive segment pairs — KD-tree in export_gltf.py maps each
            # endpoint to its nearest Blender vertex; if they differ and share a
            # mesh edge, that edge is marked sharp+seam.
            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 _tessellate_with_gmsh(shape, linear_deflection: float, angular_deflection: float) -> None:
    """Tessellate an OCC TopoDS_Shape using GMSH Frontal-Delaunay mesher.

    Writes the resulting Poly_Triangulation back to each TopoDS_Face via
    BRep_Builder.UpdateFace(), so the shape's tessellation data is available
    to RWGltf_CafWriter exactly like BRepMesh output.

    GMSH surface tags correspond 1:1 to faces in TopExp_Explorer(FACE) order
    after importShapes() from a .brep file — no coordinate-based matching needed.

    Falls back silently to BRepMesh for any face that GMSH cannot mesh (degenerate
    geometry, e.g. degenerate poles).

    Args:
        shape: tessellated in-place (OCC TopoDS_Shape)
        linear_deflection: controls CharacteristicLengthMax (mm)
        angular_deflection: controls minimum circle subdivision points (rad)
    """
    import math as _math
    import tempfile
    import gmsh

    from OCP.BRep import BRep_Builder
    from OCP.BRepTools import BRepTools
    from OCP.BRepMesh import BRepMesh_IncrementalMesh
    from OCP.Poly import Poly_Triangulation, Poly_Array1OfTriangle, Poly_Triangle
    from OCP.TColgp import TColgp_Array1OfPnt
    from OCP.TopExp import TopExp_Explorer
    from OCP.TopAbs import TopAbs_FACE
    from OCP.TopoDS import TopoDS as _TopoDS
    from OCP.gp import gp_Pnt

    # Write shape to temporary .brep for GMSH import
    with tempfile.NamedTemporaryFile(suffix=".brep", delete=False) as tmp:
        brep_path = tmp.name

    n_faces_gmsh = 0
    n_faces_fallback = 0
    n_triangles_total = 0

    try:
        BRepTools.Write_s(shape, brep_path)

        import os as _os
        n_threads = min(_os.cpu_count() or 1, 16)  # cap at 16 — sweet spot on benchmark
        gmsh.initialize()
        gmsh.option.setNumber("General.Terminal", 0)       # suppress console output
        gmsh.option.setNumber("General.NumThreads", n_threads)    # enable OpenMP parallelism
        gmsh.option.setNumber("Mesh.MaxNumThreads1D", n_threads)  # parallel edge meshing
        gmsh.option.setNumber("Mesh.MaxNumThreads2D", n_threads)  # parallel surface meshing
        gmsh.option.setNumber("Mesh.Algorithm", 6)         # Frontal-Delaunay 2D
        gmsh.option.setNumber("Mesh.RecombineAll", 0)      # keep triangles (no quads)
        # CharacteristicLength is an edge LENGTH target; OCC linear_deflection is a surface
        # DEVIATION tolerance. Empirically: OCC 0.1mm deflection on a 50mm cylinder produces
        # ~5mm edge lengths. Scale by 50× to match OCC density (target ≤120% of OCC file size).
        # MinimumCirclePoints: OCC angular_deflection=0.1rad → effectively ~12 uniform pts/circle.
        # Cap at 12 to avoid GMSH generating 3–5× more edges than OCC on cylindrical surfaces.
        gmsh.option.setNumber("Mesh.CharacteristicLengthMin", linear_deflection * 0.5)
        gmsh.option.setNumber("Mesh.CharacteristicLengthMax", linear_deflection * 50.0)
        min_circle_pts = min(12, max(6, int(_math.ceil(2.0 * _math.pi / max(angular_deflection, 0.01)))))
        gmsh.option.setNumber("Mesh.MinimumCirclePoints", min_circle_pts)
        gmsh.option.setNumber("Mesh.MinimumCurvePoints", 3)
        # Reduce noise from GMSH warnings
        gmsh.option.setNumber("General.Verbosity", 1)

        gmsh.model.add("shape")
        gmsh.model.occ.importShapes(brep_path)
        gmsh.model.occ.synchronize()
        gmsh.model.mesh.generate(2)

        # Build lookup: surface tag → list of (node_coords, triangle_node_tags)
        # GMSH surface tags from importShapes correspond 1:1 to TopExp_Explorer(FACE) order
        surface_tags = [tag for (_, tag) in gmsh.model.getEntities(2)]

        # Collect per-surface mesh data
        surface_mesh: dict[int, tuple] = {}  # tag → (nodes_xyz, triangles)
        for stag in surface_tags:
            try:
                node_tags, coords, _ = gmsh.model.mesh.getNodes(dim=2, tag=stag, includeBoundary=True)
                if len(node_tags) == 0:
                    continue
                # coords is flat [x0,y0,z0, x1,y1,z1, ...]
                node_map = {}  # gmsh tag → 1-based local index
                pts_xyz = []
                for i, ntag in enumerate(node_tags):
                    x, y, z = coords[3*i], coords[3*i+1], coords[3*i+2]
                    node_map[ntag] = i + 1  # 1-based
                    pts_xyz.append((x, y, z))

                elem_types, elem_tags, elem_node_tags = gmsh.model.mesh.getElements(dim=2, tag=stag)
                tris = []
                for etype, etags, entags in zip(elem_types, elem_tags, elem_node_tags):
                    if etype != 2:  # type 2 = triangle
                        continue
                    n_elems = len(etags)
                    for k in range(n_elems):
                        a = node_map.get(entags[3*k])
                        b = node_map.get(entags[3*k+1])
                        c = node_map.get(entags[3*k+2])
                        if a and b and c:
                            tris.append((a, b, c))

                if pts_xyz and tris:
                    surface_mesh[stag] = (pts_xyz, tris)
            except Exception:
                continue

    except Exception as _gmsh_err:
        print(f"WARNING: GMSH failed ({_gmsh_err}), falling back to BRepMesh", file=sys.stderr)
        # Fallback: full BRepMesh on the whole shape
        BRepMesh_IncrementalMesh(shape, linear_deflection, False, angular_deflection, True)
        return
    finally:
        try:
            gmsh.finalize()
        except Exception:
            pass
        try:
            Path(brep_path).unlink(missing_ok=True)
        except Exception:
            pass

    # Map GMSH surface data back to OCC faces via TopExp_Explorer (same order as importShapes)
    builder = BRep_Builder()
    explorer = TopExp_Explorer(shape, TopAbs_FACE)
    face_index = 0  # 0-based → maps to surface_tags[face_index]

    while explorer.More():
        face = _TopoDS.Face_s(explorer.Current())
        if face_index < len(surface_tags):
            stag = surface_tags[face_index]
            mesh_data = surface_mesh.get(stag)
            if mesh_data:
                pts_xyz, tris = mesh_data
                n_nodes = len(pts_xyz)
                n_tris = len(tris)
                try:
                    arr_pts = TColgp_Array1OfPnt(1, n_nodes)
                    for idx, (x, y, z) in enumerate(pts_xyz, 1):
                        arr_pts.SetValue(idx, gp_Pnt(x, y, z))

                    arr_tris = Poly_Array1OfTriangle(1, n_tris)
                    for idx, (a, b, c) in enumerate(tris, 1):
                        arr_tris.SetValue(idx, Poly_Triangle(a, b, c))

                    tri = Poly_Triangulation(arr_pts, arr_tris)
                    builder.UpdateFace(face, tri)
                    n_faces_gmsh += 1
                    n_triangles_total += n_tris
                except Exception as _e:
                    # Fallback for this face only
                    BRepMesh_IncrementalMesh(face, linear_deflection, False, angular_deflection, False)
                    n_faces_fallback += 1
            else:
                # No GMSH mesh for this surface → BRepMesh fallback for this face
                BRepMesh_IncrementalMesh(face, linear_deflection, False, angular_deflection, False)
                n_faces_fallback += 1
        else:
            BRepMesh_IncrementalMesh(face, linear_deflection, False, angular_deflection, False)
            n_faces_fallback += 1

        face_index += 1
        explorer.Next()

    print(
        f"GMSH tessellation: {n_faces_gmsh} faces meshed, "
        f"{n_faces_fallback} BRepMesh fallback, "
        f"{n_triangles_total} triangles total"
        f" (threads={n_threads})"
    )


def _collect_part_key_mappings(shape_tool, free_labels) -> tuple[dict[str, str], dict[str, list[str]]]:
    """Return canonical and occurrence-aware part-key mappings for leaf XCAF parts."""
    import hashlib as _hashlib
    import re as _re
    from OCP.TDF import TDF_LabelSequence
    from OCP.TDataStd import TDataStd_Name
    from OCP.XCAFDoc import XCAFDoc_ShapeTool

    _af_re = _re.compile(r"_AF\d+$", _re.IGNORECASE)

    def _slug(source_name: str, xcaf_path: str = "", existing_keys: set[str] | None = None) -> str:
        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]

        if existing_keys is None:
            return slug

        key = slug
        n = 2
        while key in existing_keys:
            key = f"{slug}_{n}"
            n += 1
        existing_keys.add(key)
        return key

    part_key_map: dict[str, str] = {}
    part_key_occurrences: dict[str, list[str]] = {}
    existing_keys: set[str] = set()

    def _collect(label, path: str = "") -> None:
        name_attr = TDataStd_Name()
        name = ""
        if label.FindAttribute(TDataStd_Name.GetID_s(), name_attr):
            name = name_attr.Get().ToExtString()

        from OCP.TDF import TDF_Label as _TDF_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

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

        xcaf_path = f"{path}/{name}" if name else f"{path}/unnamed"
        if components.Length() == 0:
            normalized = _af_re.sub("", name) if name else ""
            if normalized:
                part_key = _slug(name, xcaf_path, existing_keys)
                part_key_map.setdefault(normalized, part_key)
                part_key_occurrences.setdefault(normalized, []).append(part_key)
            return

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

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

    return part_key_map, part_key_occurrences


def _collect_part_key_map(shape_tool, free_labels) -> dict[str, str]:
    part_key_map, _ = _collect_part_key_mappings(shape_tool, free_labels)
    return part_key_map


def _normalize_occ_node_name(raw_name: str) -> str:
    """Collapse exporter-only suffixes back to their semantic OCC source name."""
    import re as _re

    name = (raw_name or "").strip()
    name = _re.sub(r"\.\d{3}$", "", name)

    previous = None
    while previous != name:
        previous = name
        name = _re.sub(r"_AF\d+(_ASM)?_?$", "", name, flags=_re.IGNORECASE)
    return name


def _slugify_occ_name(raw_name: str) -> str:
    import re as _re

    base = _re.sub(r"([a-z])([A-Z])", r"\1_\2", raw_name or "")
    return _re.sub(r"[^a-z0-9]+", "_", base.lower()).strip("_")


def _is_instance_variant_name(mesh_name: str, semantic_name: str) -> bool:
    if not mesh_name or not semantic_name:
        return False
    if mesh_name == semantic_name:
        return True
    if mesh_name.startswith(f"{semantic_name}_"):
        return True
    return mesh_name.rstrip("0123456789_") == semantic_name


def _read_vec(node: dict, key: str, default: tuple[float, ...]) -> tuple[float, ...]:
    value = node.get(key)
    if not isinstance(value, list) or len(value) != len(default):
        return default
    result: list[float] = []
    for idx, fallback in enumerate(default):
        item = value[idx]
        result.append(float(item) if isinstance(item, (int, float)) else fallback)
    return tuple(result)


def _have_matching_local_transforms(left: dict, right: dict) -> bool:
    position_left = _read_vec(left, "translation", (0.0, 0.0, 0.0))
    position_right = _read_vec(right, "translation", (0.0, 0.0, 0.0))
    scale_left = _read_vec(left, "scale", (1.0, 1.0, 1.0))
    scale_right = _read_vec(right, "scale", (1.0, 1.0, 1.0))
    rotation_left = _read_vec(left, "rotation", (0.0, 0.0, 0.0, 1.0))
    rotation_right = _read_vec(right, "rotation", (0.0, 0.0, 0.0, 1.0))

    epsilon = 1e-5
    if any(abs(a - b) > epsilon for a, b in zip(position_left, position_right)):
        return False
    if any(abs(a - b) > epsilon for a, b in zip(scale_left, scale_right)):
        return False

    dot = sum(a * b for a, b in zip(rotation_left, rotation_right))
    return abs(1 - abs(dot)) <= 1e-4


def _part_key_priority(source_name: str, part_key: str) -> tuple[int, int, int]:
    import re as _re

    penalty = 0
    if _re.search(r"_AF\d+(_ASM)?_?$", source_name, flags=_re.IGNORECASE):
        penalty += 100
    if _re.search(r"_af\d+$", part_key, flags=_re.IGNORECASE):
        penalty += 100
    if _re.search(r"_\d+_\d+$", source_name):
        penalty += 25
    return (penalty, len(source_name), len(part_key))


def _build_exact_part_key_lookup(part_key_map: dict[str, str]) -> dict[str, str]:
    return {
        _slugify_occ_name(raw_name): part_key
        for raw_name, part_key in part_key_map.items()
        if raw_name and part_key
    }


def _build_canonical_part_key_lookup(part_key_map: dict[str, str]) -> dict[str, str]:
    canonical: dict[str, tuple[tuple[int, int, int], str]] = {}
    for raw_name, part_key in part_key_map.items():
        normalized = _normalize_occ_node_name(raw_name)
        if not normalized or not part_key:
            continue
        slug = _slugify_occ_name(normalized)
        if not slug:
            continue
        candidate = (_part_key_priority(raw_name, part_key), part_key)
        current = canonical.get(slug)
        if current is None or candidate[0] < current[0]:
            canonical[slug] = candidate
    return {slug: part_key for slug, (_, part_key) in canonical.items()}


def _score_semantic_sibling(
    mesh_name: str,
    semantic_name: str,
    part_key: str,
    *,
    transforms_match: bool,
) -> int:
    import re as _re

    transform_bonus = 10_000 if transforms_match else 0
    canonical_bonus = 1_000 if not _re.search(r"_af\d+$", part_key, flags=_re.IGNORECASE) else 0
    exact_bonus = 100 if mesh_name.rstrip("0123456789_") == semantic_name else 0
    specificity_bonus = len(semantic_name)
    return transform_bonus + canonical_bonus + exact_bonus + specificity_bonus


def _find_semantic_sibling_part_key(
    node_index: int,
    nodes: list[dict],
    parent_by_child: dict[int, int],
    canonical_lookup: dict[str, str],
) -> str | None:
    node = nodes[node_index]
    if node.get("mesh") is None:
        return None

    parent_index = parent_by_child.get(node_index)
    if parent_index is None:
        return None

    mesh_name = _normalize_occ_node_name(node.get("name", ""))
    if not mesh_name:
        return None

    best_match: tuple[int, str] | None = None
    parent = nodes[parent_index]
    for sibling_index in parent.get("children") or []:
        if sibling_index == node_index:
            continue
        sibling = nodes[sibling_index]
        if sibling.get("mesh") is not None:
            continue

        semantic_name = _normalize_occ_node_name(sibling.get("name", ""))
        if not _is_instance_variant_name(mesh_name, semantic_name):
            continue

        sibling_part_key = canonical_lookup.get(_slugify_occ_name(semantic_name))
        if not sibling_part_key:
            continue

        score = _score_semantic_sibling(
            mesh_name,
            semantic_name,
            sibling_part_key,
            transforms_match=_have_matching_local_transforms(node, sibling),
        )
        candidate = (score, sibling_part_key)
        if best_match is None or candidate[0] > best_match[0]:
            best_match = candidate

    return best_match[1] if best_match else None


def _inject_glb_extras(
    glb_path: Path,
    extras: dict,
    part_key_map: dict | None = None,
    part_key_occurrences: dict[str, list[str]] | None = None,
) -> None:
    """Patch a GLB binary to add/update scenes[0].extras JSON field."""
    import struct as _struct

    data = glb_path.read_bytes()
    json_len = _struct.unpack_from("<I", data, 12)[0]
    json_type = _struct.unpack_from("<I", data, 16)[0]
    if json_type != 0x4E4F534A:
        print(
            "WARNING: _inject_glb_extras: unexpected chunk type, skipping extras injection",
            file=sys.stderr,
        )
        return

    j = json.loads(data[20: 20 + json_len])
    if "scenes" in j and j["scenes"]:
        existing = j["scenes"][0].get("extras") or {}
        existing.update(extras)
        j["scenes"][0]["extras"] = existing
    else:
        j.setdefault("extras", {}).update(extras)

    if part_key_map:
        exact_lookup = _build_exact_part_key_lookup(part_key_map)
        canonical_lookup = _build_canonical_part_key_lookup(part_key_map)
        occurrence_lookup: dict[str, list[str]] = {}
        if part_key_occurrences:
            for raw_name, part_keys in part_key_occurrences.items():
                normalized = _normalize_occ_node_name(raw_name)
                if not normalized:
                    continue
                cleaned_keys = [part_key for part_key in part_keys if part_key]
                if cleaned_keys:
                    occurrence_lookup[normalized] = cleaned_keys

        occurrence_indices: dict[str, int] = {}
        parent_by_child: dict[int, int] = {}
        nodes = j.get("nodes", [])
        for parent_index, parent in enumerate(nodes):
            for child_index in parent.get("children") or []:
                if isinstance(child_index, int):
                    parent_by_child[child_index] = parent_index

        n_stamped = 0
        for node_index, node in enumerate(nodes):
            raw_name = node.get("name", "")
            if not raw_name:
                continue

            normalized_raw = _normalize_occ_node_name(raw_name)
            canonical_part_key = canonical_lookup.get(_slugify_occ_name(normalized_raw))
            exact_part_key = exact_lookup.get(_slugify_occ_name(raw_name))

            if node.get("mesh") is not None:
                occurrence_part_key = None
                occurrence_keys = occurrence_lookup.get(normalized_raw)
                if occurrence_keys:
                    occurrence_index = occurrence_indices.get(normalized_raw, 0)
                    if occurrence_index < len(occurrence_keys):
                        occurrence_part_key = occurrence_keys[occurrence_index]
                        occurrence_indices[normalized_raw] = occurrence_index + 1

                semantic_sibling_part_key = _find_semantic_sibling_part_key(
                    node_index,
                    nodes,
                    parent_by_child,
                    canonical_lookup,
                )
                part_key = (
                    occurrence_part_key
                    or exact_part_key
                    or canonical_part_key
                    or semantic_sibling_part_key
                )
            else:
                part_key = canonical_part_key or exact_part_key

            if part_key:
                node.setdefault("extras", {})["partKey"] = part_key
                n_stamped += 1

        print(f"Stamped partKey extras on {n_stamped} GLB nodes")

    new_json = json.dumps(j, separators=(",", ":"))
    pad = (4 - len(new_json) % 4) % 4
    new_json_bytes = new_json.encode() + b" " * pad

    rest = data[20 + json_len:]
    new_chunk = _struct.pack("<II", len(new_json_bytes), 0x4E4F534A) + new_json_bytes
    new_total = 12 + len(new_chunk) + len(rest)
    new_header = _struct.pack("<III", 0x46546C67, 2, new_total)
    glb_path.write_bytes(new_header + new_chunk + rest)


def _prepare_atomic_export_path(output_path: Path) -> Path:
    """Return a non-existent temp file path adjacent to the final GLB target."""
    output_path.parent.mkdir(parents=True, exist_ok=True)
    temp_path = output_path.with_name(
        f".{output_path.stem}.{os.getpid()}.tmp{output_path.suffix}"
    )
    temp_path.unlink(missing_ok=True)
    return temp_path


def _finalize_atomic_export(temp_path: Path, output_path: Path) -> None:
    """Publish the validated temp GLB atomically to the canonical output path."""
    output_path.parent.mkdir(parents=True, exist_ok=True)
    temp_path.replace(output_path)


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

    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, TCollection_AsciiString
    from OCP.TDF import TDF_LabelSequence
    from OCP.BRepMesh import BRepMesh_IncrementalMesh
    from OCP.IFSelect import IFSelect_RetDone
    from OCP.Message import Message_ProgressRange

    # --- Init XDE document ---
    app = XCAFApp_Application.GetApplication_s()
    doc = TDocStd_Document(TCollection_ExtendedString("MDTV-CAF"))
    app.InitDocument(doc)

    # --- Read STEP into XDE (preserves part names + embedded colors) ---
    reader = STEPCAFControl_Reader()
    reader.SetNameMode(True)
    reader.SetColorMode(True)
    reader.SetLayerMode(True)
    status = reader.ReadFile(args.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())

    # --- Tessellate all free shapes (in mm — original STEP coordinates) ---
    # IMPORTANT: We do NOT use BRepBuilderAPI_Transform for mm→m scaling because
    # it destroys Poly_Triangulation data, and SetShape on a root XCAF label does
    # not propagate to component labels.  Instead we tessellate in mm, export the
    # GLB in mm, and post-process the GLB to add a root scale node (0.001).
    free_labels = TDF_LabelSequence()
    shape_tool.GetFreeShapes(free_labels)

    # Collect partKeyMap before tessellation (XCAF names are stable at this point)
    part_key_map, part_key_occurrences = _collect_part_key_mappings(shape_tool, free_labels)
    print(f"partKeyMap: {len(part_key_map)} unique part names collected")

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

    engine = getattr(args, "tessellation_engine", "occ")
    if engine == "gmsh":
        # GMSH: tessellate each solid individually to cap peak RAM usage.
        # Strategy:
        #   1. BRepMesh baseline on full root_shape — tessellates ALL face types
        #      (solids, shells, free faces). Ensures nothing is skipped.
        #   2. GMSH override per unique SOLID — better seam topology.
        #      Overrides the BRepMesh triangulation on solid faces only.
        #      REVERSED solids (mirrored instances) keep BRepMesh to avoid
        #      GMSH inverted-Jacobian issues.
        # Deduplication uses IsSame() (TShape pointer comparison) — NOT id(TShape())
        # because OCP creates a new Python wrapper per TShape() call, making id() unreliable.
        from OCP.TopExp import TopExp_Explorer as _Explorer
        from OCP.TopAbs import TopAbs_SOLID as _SOLID, TopAbs_SHELL as _SHELL, TopAbs_REVERSED as _REVERSED
        from OCP.TopLoc import TopLoc_Location as _TopLoc_Location

        for i in range(1, free_labels.Length() + 1):
            root_shape = shape_tool.GetShape_s(free_labels.Value(i))
            if root_shape.IsNull():
                continue

            # Step 1: BRepMesh baseline — catches non-solid shapes (free faces, shells)
            # that TopExp_Explorer(SOLID) would miss. Also provides fallback for any
            # solid that GMSH fails to tessellate.
            BRepMesh_IncrementalMesh(
                root_shape,
                args.linear_deflection,
                False,
                args.angular_deflection,
                True,
            )

            # Step 2: GMSH override for SOLID shapes (better seam topology)
            _seen_shapes: list = []

            solids = []
            exp = _Explorer(root_shape, _SOLID)
            while exp.More():
                solids.append(exp.Current())
                exp.Next()

            if not solids:
                exp = _Explorer(root_shape, _SHELL)
                while exp.More():
                    solids.append(exp.Current())
                    exp.Next()

            from OCP.TopoDS import TopoDS_Compound as _Compound
            from OCP.BRep import BRep_Builder as _BBuilder

            eligible = []
            for solid in solids:
                if solid.Orientation() == _REVERSED:
                    continue
                if any(solid.IsSame(s) for s in _seen_shapes):
                    continue
                eligible.append(solid.Located(_TopLoc_Location()))
                _seen_shapes.append(solid)

            if eligible:
                if len(eligible) == 1:
                    _tessellate_with_gmsh(eligible[0], args.linear_deflection, args.angular_deflection)
                else:
                    compound = _Compound()
                    bb = _BBuilder()
                    bb.MakeCompound(compound)
                    for s in eligible:
                        bb.Add(compound, s)
                    _tessellate_with_gmsh(compound, args.linear_deflection, args.angular_deflection)
    else:
        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,   # isRelative
                    args.angular_deflection,
                    True,    # isInParallel
                )

    # --- Extract sharp B-rep edge pairs (coords in mm, same as tessellation) ---
    sharp_pairs: 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():
                pairs = _extract_sharp_edge_pairs(root_shape, args.sharp_threshold)
                sharp_pairs.extend(pairs)
        print(f"Total OCC sharp segment pairs: {len(sharp_pairs)}")
    except Exception as _exc:
        print(f"WARNING: sharp edge extraction failed (non-fatal): {_exc}", file=sys.stderr)
        sharp_pairs = []

    # --- Apply colors ---
    if color_map:
        _apply_color_map(shape_tool, color_tool, free_labels, color_map)
        print(f"Applied color_map ({len(color_map)} entries)")
    else:
        _apply_palette_colors(shape_tool, color_tool, free_labels)
        print("Applied palette colors (no color_map provided)")

    # --- Export GLB via RWGltf_CafWriter (in mm, Z-up → Y-up handled by writer) ---
    from OCP.RWGltf import RWGltf_CafWriter

    out = Path(args.output_path)
    temp_out = _prepare_atomic_export_path(out)

    writer = RWGltf_CafWriter(TCollection_AsciiString(str(temp_out)), True)  # True = binary GLB
    # MergeFaces=True merges per-face triangulations into a single buffer per shape.
    # Without this, RWGltf_CafWriter fails to find per-face Poly_Triangulation data
    # from the XCAF component hierarchy and falls back to degenerate meshes (~2 verts/face).
    writer.SetMergeFaces(True)

    # Perform export
    try:
        from OCP.TColStd import TColStd_IndexedDataMapOfStringString
        metadata = TColStd_IndexedDataMapOfStringString()
        ok = writer.Perform(doc, metadata, Message_ProgressRange())
    except TypeError:
        # Older API without metadata dict
        ok = writer.Perform(doc, Message_ProgressRange())

    if not ok or not temp_out.exists() or temp_out.stat().st_size == 0:
        temp_out.unlink(missing_ok=True)
        print(f"ERROR: RWGltf_CafWriter.Perform returned ok={ok}, file exists={temp_out.exists()}",
              file=sys.stderr)
        sys.exit(1)

    print(f"GLB exported: {temp_out.name} ({temp_out.stat().st_size // 1024} KB)")

    # --- Inject sharp edge pairs and partKeyMap into GLB extras ---
    # Blender 5.0 reads scenes[0].extras as scene custom properties on import,
    # making the data available to export_gltf.py as bpy.context.scene["key"].
    # partKeyMap is read by Three.js in ThreeDViewer to resolve partKey from mesh name.
    try:
        extras_payload: dict = {}
        if sharp_pairs:
            extras_payload["hartomat_sharp_edge_pairs"] = sharp_pairs
            extras_payload["hartomat_sharp_threshold_deg"] = args.sharp_threshold
        if part_key_map:
            extras_payload["partKeyMap"] = part_key_map
        if extras_payload:
            _inject_glb_extras(
                temp_out,
                extras_payload,
                part_key_map=part_key_map if part_key_map else None,
                part_key_occurrences=part_key_occurrences if part_key_occurrences else None,
            )
            if sharp_pairs:
                print(f"Injected {len(sharp_pairs)} sharp edge segment pairs into GLB extras")
            if part_key_map:
                print(f"Injected partKeyMap ({len(part_key_map)} entries) into GLB extras")
    except Exception as _exc:
        print(f"WARNING: GLB extras injection failed (non-fatal): {_exc}", file=sys.stderr)

    _finalize_atomic_export(temp_out, out)
    print(f"GLB finalized: {out.name} ({out.stat().st_size // 1024} KB)")

    # NOTE: RWGltf_CafWriter reads unit metadata from the XDE document (set by
    # STEPCAFControl_Reader from the STEP file's SI_UNIT declarations) and converts
    # mm → m automatically.  It also handles Z-up → Y-up coordinate transform.
    # No additional scaling or BRepBuilderAPI_Transform is needed.


if __name__ == "__main__":
    try:
        main()
    except SystemExit:
        raise
    except Exception:
        traceback.print_exc()
        sys.exit(1)