src/llm.py

from dataclasses import dataclass, field, asdict
from typing import List, Tuple, Optional, Iterable, Dict, Any
import cadquery as cq
import math
import json
import base64
from pathlib import Path

from openai import OpenAI
import os


Point2D = Tuple[float, float]


# =========================
# Input profile definitions
# =========================

@dataclass
class PolyProfile:
    points: List[Point2D]


@dataclass
class CircleProfile:
    center: Point2D
    radius: float


@dataclass
class Contour:
    kind: str  # "poly" | "circle"
    points: List[Point2D]  # 閉ループ点列（最後は先頭と同じ）
    area_abs: float
    sample_point: Point2D
    wire: Optional[cq.Wire] = None
    children: List["Contour"] = field(default_factory=list)
    parent: Optional["Contour"] = None
    depth: int = 0


# =========================
# Geometry utilities
# =========================

def dedupe_sequential(pts: List[Point2D], eps: float = 1e-9) -> List[Point2D]:
    if not pts:
        return pts
    out = [pts[0]]
    for p in pts[1:]:
        if abs(p[0] - out[-1][0]) > eps or abs(p[1] - out[-1][1]) > eps:
            out.append(p)
    return out


def ensure_closed(pts: List[Point2D], eps: float = 1e-9) -> List[Point2D]:
    if not pts:
        return pts
    if abs(pts[0][0] - pts[-1][0]) > eps or abs(pts[0][1] - pts[-1][1]) > eps:
        return pts + [pts[0]]
    return pts


def signed_area_polygon(pts_closed: List[Point2D]) -> float:
    pts = pts_closed[:-1]
    a = 0.0
    n = len(pts)
    for i in range(n):
        x1, y1 = pts[i]
        x2, y2 = pts[(i + 1) % n]
        a += x1 * y2 - x2 * y1
    return 0.5 * a


def polygon_centroid(pts_closed: List[Point2D]) -> Point2D:
    pts = pts_closed[:-1]
    area2 = 0.0
    cx = 0.0
    cy = 0.0
    n = len(pts)

    for i in range(n):
        x1, y1 = pts[i]
        x2, y2 = pts[(i + 1) % n]
        cross = x1 * y2 - x2 * y1
        area2 += cross
        cx += (x1 + x2) * cross
        cy += (y1 + y2) * cross

    if abs(area2) < 1e-12:
        sx = sum(x for x, _ in pts)
        sy = sum(y for _, y in pts)
        return (sx / len(pts), sy / len(pts))

    cx /= (3.0 * area2)
    cy /= (3.0 * area2)
    return (cx, cy)


def point_on_segment(pt: Point2D, a: Point2D, b: Point2D, eps: float = 1e-9) -> bool:
    px, py = pt
    ax, ay = a
    bx, by = b

    cross = (px - ax) * (by - ay) - (py - ay) * (bx - ax)
    if abs(cross) > eps:
        return False

    dot = (px - ax) * (bx - ax) + (py - ay) * (by - ay)
    if dot < -eps:
        return False

    sq_len = (bx - ax) ** 2 + (by - ay) ** 2
    if dot - sq_len > eps:
        return False

    return True


def point_in_polygon(pt: Point2D, poly_closed: List[Point2D], include_boundary: bool = True) -> bool:
    x, y = pt
    poly = poly_closed[:-1]
    inside = False
    n = len(poly)

    for i in range(n):
        p1 = poly[i]
        p2 = poly[(i + 1) % n]

        if include_boundary and point_on_segment(pt, p1, p2):
            return True

        x1, y1 = p1
        x2, y2 = p2
        intersects = ((y1 > y) != (y2 > y)) and (
            x < (x2 - x1) * (y - y1) / ((y2 - y1) + 1e-30) + x1
        )
        if intersects:
            inside = not inside

    return inside


def polygon_interior_sample_point(pts_closed: List[Point2D]) -> Point2D:
    c = polygon_centroid(pts_closed)
    if point_in_polygon(c, pts_closed, include_boundary=False):
        return c

    pts = pts_closed[:-1]
    if len(pts) >= 3:
        x0, y0 = pts[0]
        x1, y1 = pts[1]
        mid = ((x0 + x1) * 0.5, (y0 + y1) * 0.5)

        cc = polygon_centroid(pts_closed)
        sx = mid[0] * 0.9 + cc[0] * 0.1
        sy = mid[1] * 0.9 + cc[1] * 0.1
        probe = (sx, sy)
        if point_in_polygon(probe, pts_closed, include_boundary=False):
            return probe

    return c


def contour_contains(a: "Contour", b: "Contour", eps: float = 1e-9) -> bool:
    if a.kind == "poly":
        return point_in_polygon(b.sample_point, a.points, include_boundary=False)
    elif a.kind == "circle":
        cx, cy = a.sample_point
        rx = a.points[0][0] - cx
        ry = a.points[0][1] - cy
        r = math.hypot(rx, ry)
        bx, by = b.sample_point
        return math.hypot(bx - cx, by - cy) < (r - eps)
    return False


def make_poly_contour(points: List[Point2D], scale: float, min_area: float) -> Optional[Contour]:
    pts = [(x * scale, y * scale) for x, y in points]
    pts = dedupe_sequential(pts)
    pts = ensure_closed(pts)

    if len(pts) < 4:
        return None

    area = signed_area_polygon(pts)
    if abs(area) < min_area:
        return None

    wire = cq.Workplane("XY").polyline(pts[:-1]).close().wire().val()
    sample = polygon_interior_sample_point(pts)

    return Contour(
        kind="poly",
        points=pts,
        area_abs=abs(area),
        sample_point=sample,
        wire=wire,
    )


def make_circle_contour(center: Point2D, radius: float, scale: float, min_area: float) -> Optional[Contour]:
    cx, cy = center[0] * scale, center[1] * scale
    r = radius * scale
    area = math.pi * r * r

    if area < min_area or r <= 0:
        return None

    wire = cq.Workplane("XY").center(cx, cy).circle(r).wire().val()
    pts = [
        (cx + r, cy),
        (cx, cy + r),
        (cx - r, cy),
        (cx, cy - r),
        (cx + r, cy),
    ]

    return Contour(
        kind="circle",
        points=pts,
        area_abs=area,
        sample_point=(cx, cy),
        wire=wire,
    )


def build_nesting(contours: List[Contour]) -> List[Contour]:
    for c in contours:
        c.children = []
        c.parent = None
        c.depth = 0

    contours_sorted = sorted(contours, key=lambda c: c.area_abs)

    for child in contours_sorted:
        parent_candidates = [
            cand for cand in contours_sorted
            if cand is not child
            and cand.area_abs > child.area_abs
            and contour_contains(cand, child)
        ]
        if parent_candidates:
            parent = min(parent_candidates, key=lambda c: c.area_abs)
            child.parent = parent
            parent.children.append(child)

    roots = [c for c in contours if c.parent is None]

    def assign_depth(node: Contour, depth: int):
        node.depth = depth
        for ch in node.children:
            assign_depth(ch, depth + 1)

    for r in roots:
        assign_depth(r, 0)

    return roots


def face_from_contour_tree(node: Contour) -> cq.Face:
    if node.wire is None:
        raise ValueError("Contour has no wire")
    outer = node.wire
    holes = [ch.wire for ch in node.children if ch.wire is not None]
    return cq.Face.makeFromWires(outer, holes)


def collect_material_nodes(roots: List[Contour]) -> List[Contour]:
    out: List[Contour] = []

    def walk(n: Contour):
        if n.depth % 2 == 0:
            out.append(n)
        for ch in n.children:
            walk(ch)

    for r in roots:
        walk(r)

    return out


def union_workplanes(solids: List[cq.Workplane]) -> cq.Workplane:
    if not solids:
        raise ValueError("No solids to union")
    model = solids[0]
    for s in solids[1:]:
        model = model.union(s)
    return model


def build_solid(
    poly_profiles: Iterable[PolyProfile],
    circle_profiles: Iterable[CircleProfile],
    height: float,
    scale: float = 1.0,
    min_area: float = 1e-6,
) -> Optional[cq.Workplane]:
    if height <= 0:
        raise ValueError("height must be positive")

    contours: List[Contour] = []

    for prof in poly_profiles:
        c = make_poly_contour(prof.points, scale=scale, min_area=min_area)
        if c is not None:
            contours.append(c)

    for cprof in circle_profiles:
        c = make_circle_contour(
            center=cprof.center,
            radius=cprof.radius,
            scale=scale,
            min_area=min_area,
        )
        if c is not None:
            contours.append(c)

    if not contours:
        return None

    roots = build_nesting(contours)
    material_nodes = collect_material_nodes(roots)

    solids: List[cq.Workplane] = []
    for node in material_nodes:
        face = face_from_contour_tree(node)
        solid = cq.Solid.extrudeLinear(face, cq.Vector(0, 0, height))
        solids.append(cq.Workplane("XY").newObject([solid]))

    if not solids:
        return None

    return union_workplanes(solids)


def export_step(model: cq.Workplane, path: str) -> None:
    cq.exporters.export(model, path)


def export_stl(model: cq.Workplane, path: str) -> None:
    cq.exporters.export(model, path)


# =========================
# OpenAI API integration
# =========================

CADQUERY_JSON_SCHEMA: Dict[str, Any] = {
    "name": "cadquery_generation",
    "schema": {
        "type": "object",
        "additionalProperties": False,
        "properties": {
            "design_intent": {"type": "string"},
            "assumptions": {
                "type": "array",
                "items": {"type": "string"},
            },
            "parameters": {
                "type": "array",
                "items": {
                    "type": "object",
                    "additionalProperties": False,
                    "properties": {
                        "name": {"type": "string"},
                        "value": {"type": ["number", "string", "boolean", "null"]},
                    },
                    "required": ["name", "value"],
                },
            },
            "cadquery_code": {"type": "string"},
        },
        "required": ["design_intent", "assumptions", "parameters", "cadquery_code"],
    },
    "strict": True,
}


def image_file_to_data_url(image_path: str) -> str:
    path = Path(image_path)
    suffix = path.suffix.lower()
    mime_map = {
        ".png": "image/png",
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".webp": "image/webp",
    }
    if suffix not in mime_map:
        raise ValueError(f"Unsupported image type: {suffix}")

    b64 = base64.b64encode(path.read_bytes()).decode("utf-8")
    return f"data:{mime_map[suffix]};base64,{b64}"


def profiles_to_geometry_json(
    poly_profiles: Iterable[PolyProfile],
    circle_profiles: Iterable[CircleProfile],
    scale: float = 1.0,
) -> Dict[str, Any]:
    return {
        "scale": scale,
        "polylines": [
            {"points": [[x, y] for x, y in prof.points]}
            for prof in poly_profiles
        ],
        "circles": [
            {"center": [c.center[0], c.center[1]], "radius": c.radius}
            for c in circle_profiles
        ],
    }


def extract_text_from_response_output_text(resp) -> str:
    """
    SDK差分にある程度強い取り出し。
    """
    if hasattr(resp, "output_text") and resp.output_text:
        return resp.output_text

    # fallback
    try:
        texts = []
        for item in resp.output:
            if getattr(item, "type", None) == "message":
                for c in item.content:
                    if getattr(c, "type", None) in ("output_text", "text"):
                        texts.append(c.text)
        return "\n".join(texts).strip()
    except Exception:
        return ""


def generate_cadquery_with_openai(
    image_path: str,
    poly_profiles: Iterable[PolyProfile],
    circle_profiles: Iterable[CircleProfile],
    modeling_instruction: str,
    model: str = "gpt-4.1",
    scale: float = 1.0,
    api_key: Optional[str] = None,
) -> Dict[str, Any]:
    """
    画像 + 幾何情報 + 指示 から、CadQueryコードを JSON で生成する。
    """
    client = OpenAI(api_key=api_key)
    image_data_url = image_file_to_data_url(image_path)
    geometry_json = profiles_to_geometry_json(poly_profiles, circle_profiles, scale=scale)

    developer_prompt = """
You are a mechanical CAD reasoning assistant.
Return only JSON matching the schema.

Rules:
- Produce valid Python CadQuery code.
- The code must assign the final model to a variable named `result`.
- Prefer parametric variables at the top of the code.
- Use only cadquery as cq.
- Do not wrap the code in markdown fences.
- Infer design intent from both image and geometry JSON, but do not invent unsupported features.
- If uncertain, state assumptions explicitly and keep geometry simple.
"""

    user_text = f"""
以下の2D CAD レイヤー画像と幾何情報をもとに、設計意図を推定し、
CadQueryコードを生成してください。

# Modeling instruction
{modeling_instruction}

# Geometry JSON
{json.dumps(geometry_json, ensure_ascii=False, indent=2)}
"""

    resp = client.responses.create(
        model=model,
        input=[
            {
                "role": "developer",
                "content": [
                    {"type": "input_text", "text": developer_prompt}
                ],
            },
            {
                "role": "user",
                "content": [
                    {"type": "input_text", "text": user_text},
                    {"type": "input_image", "image_url": image_data_url, "detail": "high"},
                ],
            },
        ],
        text={
            "format": {
                "type": "json_schema",
                "name": CADQUERY_JSON_SCHEMA["name"],
                "schema": CADQUERY_JSON_SCHEMA["schema"],
                "strict": True,
            }
        },
    )

    raw = extract_text_from_response_output_text(resp)
    if not raw:
        raise RuntimeError("Model response did not contain text output")

    data = json.loads(raw)

    if "cadquery_code" not in data:
        raise RuntimeError("Structured output missing cadquery_code")

    return data


def execute_generated_cadquery(
    cadquery_code: str,
    extra_globals: Optional[Dict[str, Any]] = None,
) -> cq.Workplane:
    """
    生成コードを実行して result を取り出す。
    注意:
    - LLM生成コードの exec は危険なので、実運用では sandbox 推奨。
    """
    g: Dict[str, Any] = {"cq": cq, "__builtins__": __builtins__}
    if extra_globals:
        g.update(extra_globals)

    l: Dict[str, Any] = {}
    exec(cadquery_code, g, l)

    result = l.get("result", g.get("result"))
    if result is None:
        raise RuntimeError("Generated code did not assign `result`")

    if not isinstance(result, cq.Workplane):
        raise TypeError("Generated `result` is not a cadquery.Workplane")

    return result


# =========================
# Example usage
# =========================

if __name__ == "__main__":
    poly_profiles = [
        PolyProfile(points=[
            (0, 0),
            (100, 0),
            (100, 60),
            (0, 60),
        ])
    ]
    circle_profiles = [
        CircleProfile(center=(50, 30), radius=10)
    ]

    generation = generate_cadquery_with_openai(
        image_path="layer.png",
        poly_profiles=poly_profiles,
        circle_profiles=circle_profiles,
        modeling_instruction=(
            "これは2D輪郭レイヤーです。外形・穴・対称性を考慮し、"
            "板状部品として自然な3DモデルをCadQueryで作ってください。"
            "厚みは 8 mm を採用してください。"
        ),
        model="gpt-4.1",
        api_key=os.getenv("OPENAI_API_KEY")
    )

    print("=== design_intent ===")
    print(generation["design_intent"])
    print("=== assumptions ===")
    print(json.dumps(generation["assumptions"], ensure_ascii=False, indent=2))
    print("=== cadquery_code ===")
    print(generation["cadquery_code"])

    result = execute_generated_cadquery(generation["cadquery_code"])
    export_step(result, "generated_from_llm.step")
    export_stl(result, "generated_from_llm.stl")