core/executor.py

"""
Safe CadQuery code execution engine.
Executes LLM-generated CadQuery code in a sandboxed namespace,
validates the result, and exports to STEP/STL.
"""

import io
import sys
import traceback
from pathlib import Path
from typing import Any, Optional

import cadquery as cq
from pydantic import BaseModel, Field


class ExecutionResult(BaseModel):
    """Result of executing a CadQuery script."""
    model_config = {"arbitrary_types_allowed": True}

    success: bool
    result: Optional[Any] = Field(default=None, exclude=True)  # cq.Workplane
    code: str = Field(default="", exclude=True)
    error: Optional[str] = None
    stdout: str = Field(default="", exclude=True)
    volume: float = Field(default=0.0, serialization_alias="volume_mm3")
    bounding_box: tuple = Field(default=(), serialization_alias="bounding_box_mm")
    face_count: int = 0
    edge_count: int = 0

    def summary(self) -> str:
        if not self.success:
            return f"FAILED: {self.error}"
        bb = self.bounding_box
        return (
            f"OK | Volume: {self.volume:.1f} mm³ | "
            f"BBox: {bb[0]:.1f}×{bb[1]:.1f}×{bb[2]:.1f} mm | "
            f"Faces: {self.face_count} | Edges: {self.edge_count}"
        )



# Allowed imports in the sandboxed namespace
SAFE_NAMESPACE = {
    "cq": cq,
    "cadquery": cq,
    "math": __import__("math"),
    "__builtins__": {
        "range": range,
        "len": len,
        "abs": abs,
        "min": min,
        "max": max,
        "round": round,
        "int": int,
        "float": float,
        "tuple": tuple,
        "list": list,
        "True": True,
        "False": False,
        "None": None,
        "print": print,
        "enumerate": enumerate,
        "zip": zip,
    },
}


def sanitize_code(code: str) -> str:
    """Clean up LLM output — strip markdown fences, trailing whitespace,
    and redundant import statements (already in namespace)."""
    code = code.strip()

    # Remove markdown code fences if present
    if code.startswith("```python"):
        code = code[len("```python"):]
    elif code.startswith("```"):
        code = code[3:]
    if code.endswith("```"):
        code = code[:-3]

    # Strip import lines for modules already in namespace
    lines = code.strip().splitlines()
    cleaned = []
    for line in lines:
        stripped = line.strip()
        # Keep the line unless it's a redundant cadquery/math import
        if stripped.startswith("import cadquery") or stripped.startswith("from cadquery"):
            continue
        if stripped == "import math":
            continue
        cleaned.append(line)

    return "\n".join(cleaned).strip()


def execute_cadquery(code: str) -> ExecutionResult:
    """
    Execute a CadQuery script string and return the result.
    The script must assign its output to a variable called `result`.
    """
    code = sanitize_code(code)

    # Capture stdout
    old_stdout = sys.stdout
    sys.stdout = captured = io.StringIO()

    namespace = dict(SAFE_NAMESPACE)

    try:
        exec(code, namespace)  # noqa: S102
    except Exception:
        sys.stdout = old_stdout
        return ExecutionResult(
            success=False,
            code=code,
            error=traceback.format_exc(),
            stdout=captured.getvalue(),
        )

    sys.stdout = old_stdout
    stdout_text = captured.getvalue()

    # Extract the result
    result_obj = namespace.get("result")
    if result_obj is None:
        return ExecutionResult(
            success=False,
            code=code,
            error="Script did not assign a value to `result`.",
            stdout=stdout_text,
        )

    if not isinstance(result_obj, cq.Workplane):
        return ExecutionResult(
            success=False,
            code=code,
            error=f"Expected cq.Workplane, got {type(result_obj).__name__}",
            stdout=stdout_text,
        )

    # Extract geometry metadata
    try:
        shape = result_obj.val()
        bb = result_obj.val().BoundingBox()
        bbox_dims = (bb.xlen, bb.ylen, bb.zlen)
        volume = shape.Volume()
        faces = len(result_obj.faces().vals())
        edges = len(result_obj.edges().vals())
    except Exception as e:
        return ExecutionResult(
            success=False,
            code=code,
            error=f"Geometry extraction failed: {e}",
            stdout=stdout_text,
        )

    return ExecutionResult(
        success=True,
        result=result_obj,
        code=code,
        stdout=stdout_text,
        volume=volume,
        bounding_box=bbox_dims,
        face_count=faces,
        edge_count=edges,
    )


def export_step(result: cq.Workplane, path: str | Path) -> Path:
    """Export a CadQuery workplane to STEP format."""
    path = Path(path)
    cq.exporters.export(result, str(path), exportType="STEP")
    return path


def export_stl(result: cq.Workplane, path: str | Path, tolerance: float = 0.01) -> Path:
    """Export a CadQuery workplane to STL format."""
    path = Path(path)
    cq.exporters.export(result, str(path), exportType="STL", tolerance=tolerance)
    return path


def export_3mf(result: cq.Workplane, path: str | Path) -> Path:
    """Export a CadQuery workplane to 3MF format (slicer-ready)."""
    path = Path(path)
    cq.exporters.export(result, str(path), exportType="3MF")
    return path


def export_all(result: cq.Workplane, base_path: str | Path) -> dict[str, Path]:
    """Export to STEP, STL, and 3MF."""
    base = Path(base_path)
    return {
        "step": export_step(result, base.with_suffix(".step")),
        "stl": export_stl(result, base.with_suffix(".stl")),
        "3mf": export_3mf(result, base.with_suffix(".3mf")),
    }