src/loader/step_loader.py

"""STEP file loading with XDE support and Shape Healing."""
from __future__ import annotations

import logging
from dataclasses import dataclass, field
from pathlib import Path

from OCP.IFSelect import IFSelect_RetDone
from OCP.STEPCAFControl import STEPCAFControl_Reader
from OCP.STEPControl import STEPControl_Reader
from OCP.ShapeFix import ShapeFix_Shape, ShapeFix_ShapeTolerance
from OCP.TCollection import TCollection_ExtendedString
from OCP.TDocStd import TDocStd_Document
from OCP.TopoDS import TopoDS_Shape
from OCP.XCAFApp import XCAFApp_Application
from OCP.XCAFDoc import XCAFDoc_ShapeTool
from OCP.TDF import TDF_LabelSequence

logger = logging.getLogger(__name__)


@dataclass
class StepFileInfo:
    """Metadata about a loaded STEP file."""
    path: Path
    shape: TopoDS_Shape
    doc: TDocStd_Document | None = None
    shape_tool: XCAFDoc_ShapeTool | None = None
    reader: STEPControl_Reader | None = None
    protocol: str = ""
    unit: str = "mm"
    num_roots: int = 0
    extra: dict = field(default_factory=dict)


def _heal_shape(shape: TopoDS_Shape, tolerance: float = 0.01) -> TopoDS_Shape:
    """Apply shape healing to fix geometry issues."""
    fixer = ShapeFix_Shape(shape)
    fixer.SetPrecision(tolerance)
    fixer.SetMaxTolerance(tolerance * 10)
    fixer.Perform()

    tol_fixer = ShapeFix_ShapeTolerance()
    tol_fixer.SetTolerance(fixer.Shape(), tolerance)

    logger.info("Shape healing completed")
    return fixer.Shape()


def _detect_protocol(reader: STEPControl_Reader) -> str:
    """Detect AP protocol from STEP file."""
    try:
        ws = reader.WS()
        model = ws.Model()
        if model is not None:
            header = str(model.Header()) if hasattr(model, "Header") else ""
            for ap in ("AP214", "AP203", "AP242"):
                if ap.lower() in header.lower() or ap in header:
                    return ap
    except Exception:
        pass
    return "unknown"


def _load_xde(file_path: Path) -> tuple[TDocStd_Document, XCAFDoc_ShapeTool, str]:
    """Load STEP into XDE document and return (doc, shape_tool, protocol)."""
    app = XCAFApp_Application.GetApplication_s()
    doc = TDocStd_Document(TCollection_ExtendedString("MDTV-XCAF"))
    app.InitDocument(doc)

    xde_reader = STEPCAFControl_Reader()
    xde_reader.SetNameMode(True)
    xde_reader.SetColorMode(True)
    xde_reader.SetLayerMode(True)

    status = xde_reader.ReadFile(str(file_path))
    if status != IFSelect_RetDone:
        raise RuntimeError(f"Failed to read STEP file: {file_path} (status={status})")

    protocol = _detect_protocol(xde_reader.Reader())

    if not xde_reader.Transfer(doc):
        raise RuntimeError(f"Failed to transfer STEP data: {file_path}")

    shape_tool = XCAFDoc_ShapeTool.Set_s(doc.Main())
    return doc, shape_tool, protocol


def _load_basic(file_path: Path) -> tuple[TopoDS_Shape, str, int, STEPControl_Reader]:
    """Fallback: load with basic STEPControl_Reader."""
    reader = STEPControl_Reader()
    status = reader.ReadFile(str(file_path))
    if status != IFSelect_RetDone:
        raise RuntimeError(f"Failed to read STEP file: {file_path}")

    protocol = _detect_protocol(reader)
    num_roots = reader.NbRootsForTransfer()
    reader.TransferRoots()
    shape = reader.OneShape()
    return shape, protocol, num_roots, reader


def load_step(file_path: str | Path, heal: bool = True) -> StepFileInfo:
    """Load a STEP file, trying XDE reader first then falling back to basic.

    Args:
        file_path: Path to the .stp/.step file.
        heal: Whether to apply shape healing.

    Returns:
        StepFileInfo with shape, XDE document (if available), and metadata.
    """
    file_path = Path(file_path)
    if not file_path.exists():
        raise FileNotFoundError(f"STEP file not found: {file_path}")

    logger.info("Loading STEP file: %s (%.1f MB)", file_path, file_path.stat().st_size / 1e6)

    doc = None
    shape_tool = None
    shape = None
    basic_reader = None
    protocol = "unknown"
    num_roots = 0

    # Try XDE reader for assembly structure
    try:
        doc, shape_tool, protocol = _load_xde(file_path)

        labels = TDF_LabelSequence()
        shape_tool.GetFreeShapes(labels)
        num_roots = labels.Length()
        logger.info("XDE reader: %d free shape(s), protocol=%s", num_roots, protocol)

        if num_roots > 0:
            if num_roots == 1:
                shape = shape_tool.GetShape(labels.Value(1))
            else:
                from OCP.BRep import BRep_Builder
                from OCP.TopoDS import TopoDS_Compound
                builder = BRep_Builder()
                compound = TopoDS_Compound()
                builder.MakeCompound(compound)
                for i in range(1, num_roots + 1):
                    builder.Add(compound, shape_tool.GetShape(labels.Value(i)))
                shape = compound
    except Exception as e:
        logger.warning("XDE reader failed: %s", e)

    # Fallback to basic reader if XDE produced no shapes
    if shape is None or shape.IsNull():
        logger.info("Falling back to basic STEPControl_Reader")
        shape, protocol, num_roots, basic_reader = _load_basic(file_path)

    if shape.IsNull():
        raise RuntimeError(f"Empty shape in STEP file: {file_path}")

    if heal:
        shape = _heal_shape(shape)

    return StepFileInfo(
        path=file_path,
        shape=shape,
        doc=doc,
        shape_tool=shape_tool,
        reader=basic_reader,
        protocol=protocol,
        num_roots=num_roots,
    )