"""USD import helper for Blender headless renders.

Runs inside Blender's Python environment (bpy available).
Imports a USD stage and restores seam + sharp edges from
schaeffler:*EdgeVertexPairs primvars (mapped as Blender mesh attributes
by Blender's built-in USD importer).

USD stage convention: mm Y-up, metersPerUnit=0.001.
Blender's USD importer respects metersPerUnit and scales objects to metres.
"""
import sys

import bpy      # type: ignore[import]
import bmesh    # type: ignore[import]
from mathutils import Vector  # type: ignore[import]


def import_usd_file(usd_path: str) -> list:
    """Import USD stage into current Blender scene.

    Returns list of imported mesh objects, centred at world origin.
    USD stage is mm Y-up with metersPerUnit=0.001 — Blender scales to metres.
    """
    bpy.ops.object.select_all(action='DESELECT')
    bpy.ops.wm.usd_import(filepath=usd_path)
    parts = [o for o in bpy.context.selected_objects if o.type == 'MESH']

    if not parts:
        print(f"[import_usd] ERROR: No mesh objects imported from {usd_path}")
        sys.exit(1)

    print(f"[import_usd] imported {len(parts)} part(s) from USD: "
          f"{[p.name for p in parts[:5]]}", flush=True)

    _rename_usd_objects(parts)

    # Restore seam + sharp edges from primvars mapped to Blender mesh attributes.
    # Blender's USD importer converts Int2Array primvars to INT32_2D attributes.
    # Attribute names: "schaeffler:sharpEdgeVertexPairs" / "schaeffler:seamEdgeVertexPairs"
    restored = 0
    for part in parts:
        restored += _restore_seam_sharp(part)
    if restored:
        print(f"[import_usd] restored seam/sharp on {restored} mesh(es)", flush=True)

    # Centre combined bbox at world origin (same as import_glb convention)
    all_corners = []
    for p in parts:
        all_corners.extend(p.matrix_world @ Vector(c) for c in p.bound_box)
    if all_corners:
        mins = Vector((min(v.x for v in all_corners),
                       min(v.y for v in all_corners),
                       min(v.z for v in all_corners)))
        maxs = Vector((max(v.x for v in all_corners),
                       max(v.y for v in all_corners),
                       max(v.z for v in all_corners)))
        center = (mins + maxs) * 0.5
        all_imported = list(bpy.context.selected_objects)
        root_objects = [o for o in all_imported if o.parent is None]
        for obj in root_objects:
            obj.location -= center

    return parts


def _rename_usd_objects(parts: list) -> None:
    """No-op: mesh prims are now named after part_key in export_step_to_usd.py.

    Blender 5.0 collapses single-child Xform+Mesh into just the Mesh object,
    using the mesh prim leaf name as the Blender object name. Since the mesh
    prim is now named after the part_key (not "Mesh"), Blender imports each
    object with the correct part name, making material matching work directly.

    The object["usd:path"] custom property is NOT set by Blender's importer,
    so the previous path-based rename approach did not work.
    """
    print(f"[import_usd] mesh objects named from USD prim paths: {[p.name for p in parts[:3]]!r}",
          flush=True)


def _restore_seam_sharp(obj) -> int:
    """Apply seam+sharp edges from USD primvars mapped as Blender mesh attributes.

    Blender's USD importer maps primvars:schaeffler:sharpEdgeVertexPairs (Int2Array)
    to a mesh attribute named "schaeffler:sharpEdgeVertexPairs" with type INT32_2D.
    Each attribute element has a .value property returning a 2-tuple (v0, v1).

    Returns 1 if any edge data was applied, 0 otherwise.
    """
    mesh = obj.data
    sharp_attr = mesh.attributes.get("schaeffler:sharpEdgeVertexPairs")
    seam_attr  = mesh.attributes.get("schaeffler:seamEdgeVertexPairs")
    if not sharp_attr and not seam_attr:
        return 0

    # Ensure single-user data block before bmesh edit
    if mesh.users > 1:
        obj.data = mesh.copy()
        mesh = obj.data

    bm = bmesh.new()
    bm.from_mesh(mesh)
    bm.verts.ensure_lookup_table()

    n_verts = len(bm.verts)

    def _apply_pairs(attr, mark_fn):
        applied = 0
        for elem in attr.data:
            v = elem.value  # 2-tuple for INT32_2D
            if len(v) >= 2 and 0 <= v[0] < n_verts and 0 <= v[1] < n_verts:
                edge = bm.edges.get([bm.verts[v[0]], bm.verts[v[1]]])
                if edge:
                    mark_fn(edge)
                    applied += 1
        return applied

    n_sharp = n_seam = 0
    if sharp_attr:
        n_sharp = _apply_pairs(sharp_attr, lambda e: setattr(e, 'smooth', False))
    if seam_attr:
        n_seam = _apply_pairs(seam_attr, lambda e: setattr(e, 'seam', True))

    bm.to_mesh(mesh)
    bm.free()

    if n_sharp or n_seam:
        print(f"[import_usd] {obj.name}: {n_sharp} sharp edges, {n_seam} seam edges",
              flush=True)
    return 1