cdsl-cad/backend/engine/cdsl_engine/legacy/llm_engine.py

"""build123d 绘图引擎：执行 build_pack → STEP。"""

from __future__ import annotations

import builtins
import json
import math
import subprocess
import sys
from pathlib import Path
from typing import Any

# 保留内置 float，防止被 build123d 上下文 shadow
_f = builtins.float

from build123d import (  # noqa: E402
    Align,
    Axis,
    BuildPart,
    BuildSketch,
    Circle,
    Cone,
    Cylinder,
    Edge,
    Face,
    Location,
    Locations,
    Mode,
    Plane,
    Polygon,
    Sphere,
    Vector,
    Wire,
    export_step,
    extrude,
    import_step,
    revolve,
)


# Keep this in sync with the execution branches in run_engine_plan. The
# agent-facing schema and its parity test prevent unsupported names reaching
# this low-level dispatcher.
SUPPORTED_ATOMIC_IDS = frozenset({
    "extrude_add_blind",
    "extrude_add_two_sided",
    "extrude_cut_blind",
    "revolve_add",
    "revolve_cut",
    "hole_blind",
    "hole_countersink",
    "hole_counterbore",
    "sphere_add",
    "reference_plane",
    "reference_axis",
})


def _load(path: Path) -> dict[str, Any]:
    return json.loads(path.read_text(encoding="utf-8"))


def _plane_from_workplane(wp: dict[str, Any]) -> Plane:
    o = wp.get("origin_mm") or [0, 0, 0]
    x = wp.get("x_dir") or [1, 0, 0]
    n = wp.get("normal") or [0, 0, 1]
    return Plane(
        origin=Vector(_f(o[0]), _f(o[1]), _f(o[2])),
        x_dir=Vector(_f(x[0]), _f(x[1]), _f(x[2])),
        z_dir=Vector(_f(n[0]), _f(n[1]), _f(n[2])),
    )


def _axis_from_params(axis: dict[str, Any]) -> Axis:
    o = axis.get("origin_mm") or [0, 0, 0]
    d = axis.get("direction") or [1, 0, 0]
    return Axis(
        origin=Vector(_f(o[0]), _f(o[1]), _f(o[2])),
        direction=Vector(_f(d[0]), _f(d[1]), _f(d[2])),
    )


def _ordered_profile_points(sketch: dict[str, Any]) -> list[tuple[float, float]]:
    entities = sketch.get("entities") or []
    line_loop = [
        i for i, e in enumerate(entities) if e["type"] == "line" and not e.get("construction")
    ]
    if not line_loop:
        raise ValueError(f"sketch {sketch.get('id')}: no profile lines")
    pts: list[tuple[float, float]] = []
    for i in line_loop:
        e = entities[i]
        s = (_f(e["start"][0]), _f(e["start"][1]))
        en = (_f(e["end"][0]), _f(e["end"][1]))
        if not pts:
            pts.append(s)
        if abs(pts[-1][0] - s[0]) + abs(pts[-1][1] - s[1]) > 1e-4:
            if abs(pts[-1][0] - en[0]) + abs(pts[-1][1] - en[1]) <= 1e-4:
                s, en = en, s
            else:
                pts.append(s)
        pts.append(en)
    if abs(pts[0][0] - pts[-1][0]) + abs(pts[0][1] - pts[-1][1]) > 1e-4:
        pts.append(pts[0])
    return pts


def _arc_midpoint(edge: dict[str, Any], p1: Vector, p2: Vector, center: Vector, radius: float) -> Vector:
    """Return a point on the intended directed arc for ``make_three_point_arc``.

    Legacy contour data has no sweep direction and retains its prior shortest
    arc behavior.  Evidence-v2 analytic contours carry ``clockwise`` so a
    major arc or a clockwise arc cannot be silently inverted by the adapter.
    """
    v1 = p1 - center
    v2 = p2 - center
    if v1.length < 1e-9 or v2.length < 1e-9:
        return (p1 + p2) / 2
    n = Vector(*(edge.get("normal") or [0, 0, 1]))
    if n.length < 1e-9:
        n = v1.cross(v2)
    if n.length < 1e-9:
        n = Vector(0, 0, 1)
    n = n.normalized()
    v1n = v1.normalized() * radius
    if "clockwise" not in edge:
        bisector = v1n + v2.normalized() * radius
        if bisector.length < 1e-9:
            bisector = n.cross(v1n)
        return center + bisector.normalized() * radius
    sweep = math.atan2(n.dot(v1.cross(v2)), v1.dot(v2))
    if bool(edge["clockwise"]):
        if sweep >= 0:
            sweep -= math.tau
    elif sweep <= 0:
        sweep += math.tau
    half = sweep / 2
    midpoint_vector = v1n * math.cos(half) + n.cross(v1n) * math.sin(half)
    return center + midpoint_vector


def _face_from_contour_edges(edges_mm: list[dict[str, Any]], *, desired_normal: list[float] | None = None) -> Face:
    b123_edges: list[Edge] = []
    for e in edges_mm:
        p1 = Vector(*e["start_mm"])
        p2 = Vector(*e["end_mm"])
        if e.get("type") == "arc" and e.get("center_mm") and e.get("radius_mm") is not None:
            center = Vector(*e["center_mm"])
            r = _f(e["radius_mm"])
            v1 = p1 - center
            v2 = p2 - center
            if v1.length < 1e-9 or v2.length < 1e-9:
                b123_edges.append(Edge.make_line(p1, p2))
                continue
            mid = _arc_midpoint(e, p1, p2, center, r)
            try:
                b123_edges.append(Edge.make_three_point_arc(p1, mid, p2))
            except Exception:
                b123_edges.append(Edge.make_line(p1, p2))
        else:
            b123_edges.append(Edge.make_line(p1, p2))
    face = Face(Wire(b123_edges))
    if desired_normal is not None:
        dn = Vector(*desired_normal)
        if dn.length > 1e-9:
            fn = face.normal_at()
            if fn.dot(dn) < 0:
                # 重建反转的 Wire：边顺序反转 + 每条边起止点交换
                # 这样法向自然翻转，但每条边的几何方向不变（不同于 Face.Reversed）
                rev_edges: list[Edge] = []
                for e in reversed(edges_mm):
                    p1 = Vector(*e["end_mm"])
                    p2 = Vector(*e["start_mm"])
                    if e.get("type") == "arc" and e.get("center_mm") and e.get("radius_mm") is not None:
                        center = Vector(*e["center_mm"])
                        r = _f(e["radius_mm"])
                        v1 = p1 - center
                        v2 = p2 - center
                        if v1.length < 1e-9 or v2.length < 1e-9:
                            rev_edges.append(Edge.make_line(p1, p2))
                            continue
                        mid = _arc_midpoint(e, p1, p2, center, r)
                        try:
                            rev_edges.append(Edge.make_three_point_arc(p1, mid, p2))
                        except Exception:
                            rev_edges.append(Edge.make_line(p1, p2))
                    else:
                        rev_edges.append(Edge.make_line(p1, p2))
                face = Face(Wire(rev_edges))
    return face


def _amount(params: dict[str, Any], *, prefer_sign: str | None = None) -> float:
    dist = abs(_f(params["distance_mm"]))
    if prefer_sign == "plus":
        return dist
    if prefer_sign == "minus":
        return -dist
    return -dist if bool(params.get("reverse")) else dist


def _build_nested_circle_profiles(circles: list[dict[str, Any]]) -> None:
    """Build circular islands and holes from containment parity.

    A circle contained by one larger circle is a hole; a circle contained by
    two larger circles is an island again.  This preserves annular profiles
    without storing the heavy tessellated sketch regions from the SW export.
    """
    ordered = sorted(circles, key=lambda item: _f(item["radius_mm"]), reverse=True)
    tolerance = 1e-6
    for index, circle in enumerate(ordered):
        center = circle["center"]
        radius = _f(circle["radius_mm"])
        containing = 0
        for outer in ordered[:index]:
            outer_center = outer["center"]
            outer_radius = _f(outer["radius_mm"])
            distance = math.hypot(
                _f(center[0]) - _f(outer_center[0]),
                _f(center[1]) - _f(outer_center[1]),
            )
            if distance + radius <= outer_radius + tolerance:
                containing += 1
        mode = Mode.ADD if containing % 2 == 0 else Mode.SUBTRACT
        with Locations((_f(center[0]), _f(center[1]))):
            Circle(radius, mode=mode)


def run_engine_plan(
    pack: dict[str, Any],
    out_step: Path,
    *,
    cut_sign: str = "from_params",
) -> dict[str, Any]:
    log: list[str] = []

    compiler_context = pack.get("compiler_context")
    if isinstance(compiler_context, dict):
        # 回退路径：使用本包 translator（不依赖外部 backend.src）
        try:
            from .translator import generate_build123d_code, get_part_name
        except ImportError:
            from translator import generate_build123d_code, get_part_name

        context = dict(compiler_context)
        context.setdefault("metadata", {})["part_name"] = str(pack.get("part_id") or out_step.stem)
        out_step.parent.mkdir(parents=True, exist_ok=True)
        completed = subprocess.run(
            [sys.executable, "-c", generate_build123d_code(context)],
            cwd=out_step.parent,
            capture_output=True,
            text=True,
            timeout=180,
        )
        if completed.returncode != 0:
            raise RuntimeError(
                f"exact compiler execution failed\nSTDOUT:\n{completed.stdout}\nSTDERR:\n{completed.stderr}"
            )
        generated_name = get_part_name({"part_name": context["metadata"]["part_name"]})
        generated = out_step.parent / f"{generated_name}.step"
        if generated != out_step and generated.exists():
            generated.replace(out_step)
        if not out_step.exists():
            raise RuntimeError(f"exact compiler did not generate {out_step}")
        solid = import_step(str(out_step))
        bb = solid.bounding_box()
        return {
            "out_step": str(out_step),
            "volume_mm3": _f(solid.volume),
            "bbox_mm": {
                "min": [bb.min.X, bb.min.Y, bb.min.Z],
                "max": [bb.max.X, bb.max.Y, bb.max.Z],
            },
            "engine": "translator_fallback",
        }

    with BuildPart() as part:
        for step in pack.get("steps") or []:
            atomic = step["atomic_id"]
            params = step["params"]
            sketch = step.get("sketch")
            sid = step.get("step_id")

            if atomic == "reference_plane":
                # Context features deliberately produce no solid.  They remain
                # executable plan steps so their dependencies are preserved and
                # can be registered by the session-based runtime.
                plane = _plane_from_workplane(params.get("plane") or {})
                log.append(
                    f"{sid}: reference_plane origin={tuple(plane.origin)} normal={tuple(plane.z_dir)}"
                )

            elif atomic == "reference_axis":
                axis = _axis_from_params(params.get("axis") or {})
                log.append(
                    f"{sid}: reference_axis origin={tuple(axis.position)} direction={tuple(axis.direction)}"
                )

            elif atomic == "sphere_add":
                radius = _f(params.get("radius_mm") or 0)
                center = params.get("center_mm") or [0, 0, 0]
                if radius <= 0 or len(center) != 3:
                    raise ValueError(f"{sid}: sphere_add requires a positive radius_mm and center_mm")
                with Locations((_f(center[0]), _f(center[1]), _f(center[2]))):
                    Sphere(radius, mode=Mode.ADD)
                log.append(f"{sid}: sphere_add radius={radius}")

            elif atomic in ("extrude_add_blind", "extrude_add_two_sided", "extrude_cut_blind"):
                if sketch is None:
                    raise ValueError(f"{sid}: missing sketch")
                plane = _plane_from_workplane(sketch.get("workplane") or {})
                mode = Mode.SUBTRACT if "cut" in atomic else Mode.ADD
                edges = sketch.get("contour_edges_mm") or []
                regions = sketch.get("contour_regions_mm") or []
                sign = cut_sign if "cut" in atomic else "from_params"

                circles = [
                    e
                    for e in (sketch.get("entities") or [])
                    if e.get("type") == "circle" and not e.get("construction")
                ]
                lines = [
                    e
                    for e in (sketch.get("entities") or [])
                    if e.get("type") == "line" and not e.get("construction")
                ]

                # 多区域轮廓（外环 + 孔）：由 shape generator 展开
                if regions:
                    faces = []
                    normal = (sketch.get("workplane") or {}).get("normal")
                    for reg in regions:
                        outer_edges = reg.get("outer") or []
                        if len(outer_edges) < 2:
                            continue
                        face = _face_from_contour_edges(outer_edges, desired_normal=normal)
                        for hole_edges in reg.get("holes") or []:
                            if len(hole_edges) < 2:
                                continue
                            hole = _face_from_contour_edges(hole_edges, desired_normal=normal)
                            face = face.cut(hole)
                        faces.append(face)
                    if not faces:
                        raise ValueError(f"{sid}: contour_regions_mm produced no faces")
                    if atomic == "extrude_add_two_sided":
                        d = abs(_f(params["distance_mm"]))
                        for face in faces:
                            extrude(to_extrude=face, amount=d, both=True, mode=Mode.ADD)
                    else:
                        amt = _amount(params, prefer_sign=None if sign == "from_params" else sign)
                        for face in faces:
                            extrude(to_extrude=face, amount=amt, mode=mode)
                    log.append(f"{sid}: {atomic} regions={len(faces)}")
                    continue

                # 切除：草图常含面外框线+圆孔；优先圆孔，避免误用外框整面切除
                prefer_circles = bool(circles) and atomic.startswith("extrude_cut")

                if prefer_circles:
                    with BuildSketch(plane):
                        for e in circles:
                            with Locations((_f(e["center"][0]), _f(e["center"][1]))):
                                Circle(_f(e["radius_mm"]))
                    if atomic == "extrude_add_two_sided":
                        d = abs(_f(params["distance_mm"]))
                        extrude(amount=d, both=True, mode=Mode.ADD)
                    else:
                        amt = _amount(params, prefer_sign=None if sign == "from_params" else sign)
                        extrude(amount=amt, mode=mode)
                    log.append(f"{sid}: {atomic} circle-only n={len(circles)}")
                elif len(edges) >= 2:
                    face = _face_from_contour_edges(edges, desired_normal=sketch.get("workplane", {}).get("normal"))
                    if atomic == "extrude_add_two_sided":
                        d = abs(_f(params["distance_mm"]))
                        extrude(to_extrude=face, amount=d, both=True, mode=Mode.ADD)
                        log.append(f"{sid}: extrude_two_sided both={d} contour")
                    else:
                        amt = _amount(params, prefer_sign=None if sign == "from_params" else sign)
                        extrude(to_extrude=face, amount=amt, mode=mode)
                        log.append(f"{sid}: {atomic} amount={amt} contour")
                elif circles and not lines:
                    # 纯圆轮廓：用包含层级区分实体、内孔和孔中岛。
                    with BuildSketch(plane):
                        _build_nested_circle_profiles(circles)
                    if atomic == "extrude_add_two_sided":
                        d = abs(_f(params["distance_mm"]))
                        extrude(amount=d, both=True, mode=Mode.ADD)
                    else:
                        amt = _amount(params, prefer_sign=None if sign == "from_params" else sign)
                        extrude(amount=amt, mode=mode)
                    log.append(f"{sid}: {atomic} circle-only n={len(circles)}")
                else:
                    with BuildSketch(plane):
                        pts = _ordered_profile_points(sketch)
                        poly = pts[:-1] if len(pts) >= 2 and pts[0] == pts[-1] else pts
                        Polygon(*poly)
                        for e in circles:
                            with Locations((_f(e["center"][0]), _f(e["center"][1]))):
                                Circle(_f(e["radius_mm"]), mode=Mode.SUBTRACT)
                    if atomic == "extrude_add_two_sided":
                        d = abs(_f(params["distance_mm"]))
                        extrude(amount=d, both=True, mode=Mode.ADD)
                        log.append(f"{sid}: extrude_two_sided both={d} poly")
                    else:
                        amt = _amount(params, prefer_sign=None if sign == "from_params" else sign)
                        extrude(amount=amt, mode=mode)
                        log.append(f"{sid}: {atomic} amount={amt} poly")

            elif atomic in ("revolve_add", "revolve_cut"):
                if sketch is None:
                    raise ValueError(f"{sid}: missing sketch")
                plane = _plane_from_workplane(sketch.get("workplane") or {})
                axis = _axis_from_params(params.get("axis") or {})
                angle = _f(params.get("angle_deg") or 360)
                mode = Mode.SUBTRACT if atomic == "revolve_cut" else Mode.ADD
                edges = sketch.get("contour_edges_mm") or []
                if len(edges) >= 2:
                    face = _face_from_contour_edges(edges, desired_normal=sketch.get("workplane", {}).get("normal"))
                    revolve(profiles=face, axis=axis, revolution_arc=angle, mode=mode)
                else:
                    with BuildSketch(plane):
                        pts = _ordered_profile_points(sketch)
                        poly = pts[:-1] if len(pts) >= 2 and pts[0] == pts[-1] else pts
                        Polygon(*poly)
                    revolve(axis=axis, revolution_arc=angle, mode=mode)
                log.append(f"{sid}: {atomic} angle={angle}")

            elif atomic in ("hole_blind", "hole_countersink", "hole_counterbore"):
                dia = _f(params.get("diameter_mm") or 0)
                depth = _f(params.get("depth_mm") or 0)
                positions = params.get("positions") or []
                if sketch is not None:
                    plane = _plane_from_workplane(sketch.get("workplane") or {})
                else:
                    plane = Plane.XY
                host_face = params.get("host_face") or {}
                frame = host_face.get("frame") or {}
                frame_origin = Vector(*(frame.get("origin_mm") or plane.origin.to_tuple()))
                frame_x = Vector(*(frame.get("x_dir") or plane.x_dir.to_tuple()))
                frame_y = Vector(*(frame.get("y_dir") or plane.y_dir.to_tuple()))
                normal = plane.z_dir.normalized()
                bb = part.part.bounding_box()
                part_center = Vector(
                    (bb.min.X + bb.max.X) / 2,
                    (bb.min.Y + bb.max.Y) / 2,
                    (bb.min.Z + bb.max.Z) / 2,
                )
                inward = normal if (part_center - frame_origin).dot(normal) >= 0 else -normal
                for pos in positions:
                    mm = pos.get("mm") or [0, 0, 0]
                    start = frame_origin + frame_x * _f(mm[0]) + frame_y * _f(mm[1])
                    cs_dia = _f(params.get("countersink_diameter_mm") or 0)
                    cs_angle = _f(params.get("countersink_angle_rad") or 0)
                    cb_dia = _f(params.get("counterbore_diameter_mm") or 0)
                    cb_depth = _f(params.get("counterbore_depth_mm") or 0)
                    cs_depth = (
                        ((cs_dia - dia) / 2) / math.tan(cs_angle / 2)
                        if cs_dia > dia and cs_angle > 0
                        else 0
                    )
                    base_offset = cs_depth + (cb_depth if cb_dia > dia else 0)
                    main_depth = max(0.001, abs(depth) - base_offset)
                    main_place = Location(Plane(origin=start + inward * base_offset, z_dir=inward))
                    tools = [
                        Cylinder(
                            radius=dia / 2,
                            height=main_depth,
                            align=(Align.CENTER, Align.CENTER, Align.MIN),
                            mode=Mode.PRIVATE,
                        ).move(main_place)
                    ]
                    if cb_dia > dia and cb_depth > 0:
                        tools.append(
                            Cylinder(
                                radius=cb_dia / 2,
                                height=cb_depth,
                                align=(Align.CENTER, Align.CENTER, Align.MIN),
                                mode=Mode.PRIVATE,
                            ).move(Location(Plane(origin=start, z_dir=inward)))
                        )
                    if cs_depth > 0:
                        tools.append(
                            Cone(
                                bottom_radius=cs_dia / 2,
                                top_radius=dia / 2,
                                height=cs_depth,
                                align=(Align.CENTER, Align.CENTER, Align.MIN),
                                mode=Mode.PRIVATE,
                            ).move(Location(Plane(origin=start, z_dir=inward)))
                        )
                    drill_angle = _f(params.get("drill_angle_rad") or 0)
                    if drill_angle > 0:
                        tip_depth = (dia / 2) / math.tan(drill_angle / 2)
                        tools.append(
                            Cone(
                                bottom_radius=dia / 2,
                                top_radius=0,
                                height=tip_depth,
                                align=(Align.CENTER, Align.CENTER, Align.MIN),
                                mode=Mode.PRIVATE,
                            ).move(
                                Location(
                                    Plane(origin=start + inward * abs(depth), z_dir=inward)
                                )
                            )
                        )
                    for tool in tools:
                        part.part = part.part.cut(tool)
                log.append(f"{sid}: {atomic} npos={len(positions)}")

            else:
                raise ValueError(f"unsupported atomic_id: {atomic}")

        solid = part.part

    out_step.parent.mkdir(parents=True, exist_ok=True)
    export_step(solid, str(out_step))
    bb = solid.bounding_box()
    return {
        "out_step": str(out_step),
        "volume_mm3": _f(solid.volume),
        "bbox_mm": {
            "min": [bb.min.X, bb.min.Y, bb.min.Z],
            "max": [bb.max.X, bb.max.Y, bb.max.Z],
        },
        "log": log,
        "cut_sign": cut_sign,
    }


def main() -> None:
    import argparse

    ap = argparse.ArgumentParser()
    ap.add_argument("--pack", type=Path, required=True)
    ap.add_argument("--out-step", type=Path, required=True)
    ap.add_argument("--report", type=Path, default=None)
    ap.add_argument("--cut-sign", default="from_params", choices=["from_params", "plus", "minus"])
    args = ap.parse_args()
    info = run_engine_plan(_load(args.pack), args.out_step, cut_sign=args.cut_sign)
    if args.report:
        args.report.write_text(json.dumps(info, ensure_ascii=False, indent=2), encoding="utf-8")
    print(
        json.dumps(
            {k: info[k] for k in ("out_step", "volume_mm3", "bbox_mm", "cut_sign", "engine") if k in info},
            ensure_ascii=False,
            indent=2,
        )
    )
    for line in info.get("log") or []:
        print(line)


if __name__ == "__main__":
    main()