"""STEP to STL converter for Flamenco tasks.

Usage: python convert_step.py <step_path> <stl_path> <quality>
  quality: 'low' or 'high'

Produces:
  - Combined STL at <stl_path> (for fallback)
  - Per-part STLs in <stl_path_without_ext>_parts/ with manifest.json
"""
import sys
import os
import json
import time


def _export_per_part_stls(step_path, parts_dir, quality):
    """Export one STL per named STEP leaf shape using OCP XCAF.

    Creates parts_dir with individual STL files and a manifest.json:
      {"parts": [{"index": 0, "name": "PartName", "file": "00_PartName.stl"}, ...]}

    Returns the manifest list, or empty list on failure.
    """
    tol = 0.01 if quality == "high" else 0.3
    angular_tol = 0.05 if quality == "high" else 0.3

    try:
        from OCP.STEPCAFControl import STEPCAFControl_Reader
        from OCP.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ShapeTool
        from OCP.TDataStd import TDataStd_Name
        from OCP.TDF import TDF_Label as TDF_Label_cls, TDF_LabelSequence
        from OCP.XCAFApp import XCAFApp_Application
        from OCP.TDocStd import TDocStd_Document
        from OCP.TCollection import TCollection_ExtendedString
        from OCP.IFSelect import IFSelect_RetDone
        import cadquery as cq
    except ImportError as e:
        print(f"[convert_step] per-part export skipped (import error): {e}")
        return []

    # Read STEP with XCAF
    app = XCAFApp_Application.GetApplication_s()
    doc = TDocStd_Document(TCollection_ExtendedString("XmlOcaf"))
    app.InitDocument(doc)

    reader = STEPCAFControl_Reader()
    reader.SetNameMode(True)
    status = reader.ReadFile(str(step_path))
    if status != IFSelect_RetDone:
        print(f"[convert_step] XCAF reader failed with status {status}")
        return []

    if not reader.Transfer(doc):
        print("[convert_step] XCAF transfer failed")
        return []

    shape_tool = XCAFDoc_DocumentTool.ShapeTool_s(doc.Main())
    name_id = TDataStd_Name.GetID_s()

    # Recursively collect leaf shapes with their names
    leaves = []  # list of (name, TopoDS_Shape)

    def _get_label_name(label):
        """Extract name string from a TDF_Label."""
        name_attr = TDataStd_Name()
        if label.FindAttribute(name_id, name_attr):
            return name_attr.Get().ToExtString()
        return ""

    def _collect_leaves(label):
        """Recursively collect leaf (simple shape) labels."""
        if XCAFDoc_ShapeTool.IsAssembly_s(label):
            # Get components of this assembly
            components = TDF_LabelSequence()
            XCAFDoc_ShapeTool.GetComponents_s(label, components)
            for i in range(1, components.Length() + 1):
                comp_label = components.Value(i)
                if XCAFDoc_ShapeTool.IsReference_s(comp_label):
                    ref_label = TDF_Label_cls()
                    XCAFDoc_ShapeTool.GetReferredShape_s(comp_label, ref_label)
                    # Use the component name (instance name), fall back to referred shape name
                    comp_name = _get_label_name(comp_label)
                    ref_name = _get_label_name(ref_label)
                    # Prefer referred shape name — matches material_map keys
                    name = ref_name or comp_name
                    if XCAFDoc_ShapeTool.IsAssembly_s(ref_label):
                        _collect_leaves(ref_label)
                    elif XCAFDoc_ShapeTool.IsSimpleShape_s(ref_label):
                        # Use comp_label shape — includes instance transform (position)
                        shape = XCAFDoc_ShapeTool.GetShape_s(comp_label)
                        leaves.append((name or f"unnamed_{len(leaves)}", shape))
                else:
                    _collect_leaves(comp_label)
        elif XCAFDoc_ShapeTool.IsSimpleShape_s(label):
            name = _get_label_name(label)
            shape = XCAFDoc_ShapeTool.GetShape_s(label)
            leaves.append((name or f"unnamed_{len(leaves)}", shape))

    # Get top-level free shapes
    top_labels = TDF_LabelSequence()
    shape_tool.GetFreeShapes(top_labels)
    for i in range(1, top_labels.Length() + 1):
        _collect_leaves(top_labels.Value(i))

    if not leaves:
        print("[convert_step] no leaf shapes found via XCAF")
        return []

    # Export each leaf shape as individual STL
    os.makedirs(parts_dir, exist_ok=True)
    manifest = []

    for idx, (name, shape) in enumerate(leaves):
        # Sanitize filename: replace problematic chars
        safe_name = name.replace("/", "_").replace("\\", "_").replace(" ", "_")
        filename = f"{idx:02d}_{safe_name}.stl"
        filepath = os.path.join(parts_dir, filename)

        try:
            cq_shape = cq.Shape(shape)
            cq_shape.exportStl(filepath, tolerance=tol, angularTolerance=angular_tol)
            manifest.append({"index": idx, "name": name, "file": filename})
        except Exception as e:
            print(f"[convert_step] WARNING: failed to export part '{name}': {e}")

    # Write manifest
    manifest_path = os.path.join(parts_dir, "manifest.json")
    with open(manifest_path, "w") as f:
        json.dump({"parts": manifest}, f, indent=2)

    total_size = sum(
        os.path.getsize(os.path.join(parts_dir, p["file"]))
        for p in manifest
        if os.path.exists(os.path.join(parts_dir, p["file"]))
    )
    print(f"[convert_step] exported {len(manifest)} per-part STLs "
          f"({total_size / 1024:.0f} KB total) to {parts_dir}")

    return manifest


def main():
    if len(sys.argv) < 4:
        print("Usage: convert_step.py <step_path> <stl_path> <quality>")
        sys.exit(1)

    step_path = sys.argv[1]
    stl_path = sys.argv[2]
    quality = sys.argv[3]

    if not os.path.isfile(step_path):
        print(f"ERROR: STEP file not found: {step_path}")
        sys.exit(1)

    os.makedirs(os.path.dirname(stl_path), exist_ok=True)

    # Cache hit: skip re-conversion if STL already exists and is non-empty
    if os.path.isfile(stl_path) and os.path.getsize(stl_path) > 0:
        size_kb = os.path.getsize(stl_path) / 1024
        print(f"[convert_step] Cache hit: {stl_path} ({size_kb:.0f} KB) — skipping STEP conversion")
        stl_stem = os.path.splitext(stl_path)[0]
        parts_dir = stl_stem + "_parts"
        manifest_path = os.path.join(parts_dir, "manifest.json")
        if not os.path.isfile(manifest_path):
            print("[convert_step] Per-part STLs missing — exporting from STEP")
            t1 = time.time()
            try:
                manifest = _export_per_part_stls(step_path, parts_dir, quality)
                if manifest:
                    print(f"[convert_step] per-part export took {time.time() - t1:.1f}s")
                else:
                    print("[convert_step] per-part export empty — combined STL only")
            except Exception as e:
                print(f"[convert_step] per-part export failed (non-fatal): {e}")
        else:
            print(f"[convert_step] Per-part STLs exist: {parts_dir}")
        return

    print(f"Converting STEP -> STL: {step_path}")
    print(f"Quality: {quality}")
    t0 = time.time()

    import cadquery as cq

    tol = 0.01 if quality == "high" else 0.3
    angular_tol = 0.05 if quality == "high" else 0.3

    result = cq.importers.importStep(step_path)
    cq.exporters.export(
        result,
        stl_path,
        exportType="STL",
        tolerance=tol,
        angularTolerance=angular_tol,
    )

    elapsed = time.time() - t0
    size_kb = os.path.getsize(stl_path) / 1024
    print(f"STL written: {stl_path} ({size_kb:.0f} KB, {elapsed:.1f}s)")

    # Export per-part STLs alongside the combined STL (non-fatal)
    stl_stem = os.path.splitext(stl_path)[0]
    parts_dir = stl_stem + "_parts"
    t1 = time.time()
    try:
        manifest = _export_per_part_stls(step_path, parts_dir, quality)
        if manifest:
            print(f"[convert_step] per-part export took {time.time() - t1:.1f}s")
        else:
            print("[convert_step] per-part export failed or empty — combined STL only")
    except Exception as e:
        print(f"[convert_step] per-part export failed (non-fatal): {e}")


if __name__ == "__main__":
    main()
