src/streamlit_app.py

# app.py
import io
import math
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import List, Tuple, Dict, Iterable

import streamlit as st
import ezdxf
import cadquery as cq
from cadquery import exporters

Point2D = Tuple[float, float]


# =========================
# Data structures
# =========================
@dataclass
class PolyProfile:
    points: List[Point2D]
    layer: str
    source: str


@dataclass
class CircleProfile:
    center: Point2D
    radius: float
    layer: str
    source: str


# =========================
# Geometry helpers
# =========================
def dist(a: Point2D, b: Point2D) -> float:
    return math.hypot(a[0] - b[0], a[1] - b[1])


def almost_same(a: Point2D, b: Point2D, tol: float = 1e-4) -> bool:
    return dist(a, b) <= tol


def round_pt(p: Point2D, ndigits: int = 4) -> Point2D:
    return (round(float(p[0]), ndigits), round(float(p[1]), ndigits))


def area_polygon(points: List[Point2D]) -> float:
    """Signed area * 0.5"""
    if len(points) < 3:
        return 0.0
    s = 0.0
    for i in range(len(points)):
        x1, y1 = points[i]
        x2, y2 = points[(i + 1) % len(points)]
        s += x1 * y2 - x2 * y1
    return s / 2.0


def ensure_closed(points: List[Point2D], tol: float = 1e-4) -> List[Point2D]:
    if not points:
        return points
    if not almost_same(points[0], points[-1], tol):
        return points + [points[0]]
    return points


def dedupe_sequential(points: List[Point2D], tol: float = 1e-7) -> List[Point2D]:
    if not points:
        return points
    out = [points[0]]
    for p in points[1:]:
        if dist(p, out[-1]) > tol:
            out.append(p)
    if len(out) > 1 and dist(out[0], out[-1]) <= tol:
        out[-1] = out[0]
    return out


def sample_arc(center: Point2D, radius: float, start_deg: float, end_deg: float, segments: int = 48) -> List[Point2D]:
    """DXF ARC is CCW from start_angle to end_angle."""
    start = math.radians(start_deg)
    end = math.radians(end_deg)
    if end < start:
        end += 2 * math.pi
    pts = []
    for i in range(segments + 1):
        t = start + (end - start) * i / segments
        pts.append((center[0] + radius * math.cos(t), center[1] + radius * math.sin(t)))
    return pts


# =========================
# DXF parsing
# =========================
def load_dxf_from_upload(uploaded_file) -> ezdxf.document.Drawing:
    suffix = ".dxf"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tf:
        raw = uploaded_file.getvalue()
        tf.write(raw)
        temp_path = tf.name
    return ezdxf.readfile(temp_path, encoding="cp932")


def get_layers(doc: ezdxf.document.Drawing) -> List[str]:
    layers = set()
    for e in doc.modelspace():
        if hasattr(e.dxf, "layer"):
            layers.add(e.dxf.layer)
    return sorted(layers)


def extract_profiles(doc, target_layers: List[str], tol: float = 1e-4):
    """
    まずは扱いやすい形だけ抽出:
    - 閉じた LWPOLYLINE / POLYLINE
    - CIRCLE
    - LINE の閉ループ化
    - ARC はプレビュー用に保持（面化は現状しない）
    """
    msp = doc.modelspace()
    selected = []
    for e in msp:
        layer = getattr(e.dxf, "layer", "0")
        if (not target_layers) or (layer in target_layers):
            selected.append(e)

    poly_profiles: List[PolyProfile] = []
    circle_profiles: List[CircleProfile] = []
    line_segments: List[Tuple[Point2D, Point2D, str]] = []
    preview_arcs: List[List[Point2D]] = []
    warnings: List[str] = []

    for e in selected:
        dxftype = e.dxftype()
        layer = getattr(e.dxf, "layer", "0")

        try:
            if dxftype == "LWPOLYLINE":
                pts = [(float(x), float(y)) for x, y, *_ in e.get_points()]
                pts = dedupe_sequential(pts)
                if e.closed or (len(pts) >= 3 and almost_same(pts[0], pts[-1], tol)):
                    pts = ensure_closed(pts, tol)
                    poly_profiles.append(PolyProfile(points=pts, layer=layer, source=dxftype))
                else:
                    # 開いた LWPOLYLINE は今回はスキップ
                    warnings.append(f"Open LWPOLYLINE skipped on layer '{layer}'")

            elif dxftype == "POLYLINE":
                pts = [(float(v.dxf.location.x), float(v.dxf.location.y)) for v in e.vertices]
                pts = dedupe_sequential(pts)
                is_closed = False
                if hasattr(e, "is_closed"):
                    is_closed = bool(e.is_closed)
                if is_closed or (len(pts) >= 3 and almost_same(pts[0], pts[-1], tol)):
                    pts = ensure_closed(pts, tol)
                    poly_profiles.append(PolyProfile(points=pts, layer=layer, source=dxftype))
                else:
                    warnings.append(f"Open POLYLINE skipped on layer '{layer}'")

            elif dxftype == "CIRCLE":
                c = e.dxf.center
                circle_profiles.append(
                    CircleProfile(
                        center=(float(c.x), float(c.y)),
                        radius=float(e.dxf.radius),
                        layer=layer,
                        source=dxftype,
                    )
                )

            elif dxftype == "LINE":
                s = e.dxf.start
                t = e.dxf.end
                line_segments.append(((float(s.x), float(s.y)), (float(t.x), float(t.y)), layer))

            elif dxftype == "ARC":
                c = e.dxf.center
                pts = sample_arc(
                    center=(float(c.x), float(c.y)),
                    radius=float(e.dxf.radius),
                    start_deg=float(e.dxf.start_angle),
                    end_deg=float(e.dxf.end_angle),
                    segments=48,
                )
                preview_arcs.append(pts)
                warnings.append(f"ARC detected on layer '{layer}' (preview only, not solidified)")

            else:
                # 必要ならここに ELLIPSE, SPLINE, HATCH を足す
                pass

        except Exception as ex:
            warnings.append(f"{dxftype} parse error on layer '{layer}': {ex}")

    # LINE 群から簡易ループ生成
    line_loops, line_warnings = chain_segments_to_closed_loops(line_segments, tol=tol)
    warnings.extend(line_warnings)
    for loop_pts, layer in line_loops:
        poly_profiles.append(PolyProfile(points=loop_pts, layer=layer, source="LINE_LOOP"))

    return poly_profiles, circle_profiles, preview_arcs, warnings


def chain_segments_to_closed_loops(
    segments: List[Tuple[Point2D, Point2D, str]],
    tol: float = 1e-4
) -> Tuple[List[Tuple[List[Point2D], str]], List[str]]:
    """
    かなり素朴なループ化:
    同一レイヤー内で端点がつながる線分を順につなぎ、
    閉じたら 1 ループとして返す。
    """
    warnings = []
    by_layer: Dict[str, List[Tuple[Point2D, Point2D]]] = {}
    for s, e, layer in segments:
        by_layer.setdefault(layer, []).append((s, e))

    loops: List[Tuple[List[Point2D], str]] = []

    for layer, segs in by_layer.items():
        unused = segs[:]

        while unused:
            s0, e0 = unused.pop(0)
            chain = [s0, e0]
            changed = True

            while changed:
                changed = False
                i = 0
                while i < len(unused):
                    a, b = unused[i]

                    if almost_same(chain[-1], a, tol):
                        chain.append(b)
                        unused.pop(i)
                        changed = True
                        continue
                    elif almost_same(chain[-1], b, tol):
                        chain.append(a)
                        unused.pop(i)
                        changed = True
                        continue
                    elif almost_same(chain[0], b, tol):
                        chain = [a] + chain
                        unused.pop(i)
                        changed = True
                        continue
                    elif almost_same(chain[0], a, tol):
                        chain = [b] + chain
                        unused.pop(i)
                        changed = True
                        continue

                    i += 1

            chain = dedupe_sequential(chain)
            if len(chain) >= 4 and almost_same(chain[0], chain[-1], tol):
                loops.append((ensure_closed(chain, tol), layer))
            else:
                warnings.append(f"Open LINE chain skipped on layer '{layer}'")

    return loops, warnings


# =========================
# CadQuery build
# =========================
# def build_solid(
#     poly_profiles: List[PolyProfile],
#     circle_profiles: List[CircleProfile],
#     height: float,
#     scale: float = 1.0,
#     min_area: float = 1e-6,
# ):
#     """
#     各閉輪郭を独立したソリッドとして Z 方向に押し出し、最後に union する。

#     制約:
#     - 輪郭の内外判定はしない
#     - 穴（内周）はサポートしない
#     - 自己交差ポリゴンは未対応
#     - 不正形状に対する例外処理は未実装
#     """
#     solids = []

#     for prof in poly_profiles:
#         pts = [(p[0] * scale, p[1] * scale) for p in prof.points]
#         pts = dedupe_sequential(pts)
#         pts = ensure_closed(pts)
#         if len(pts) < 4:
#             continue
#         if abs(area_polygon(pts[:-1])) < min_area:
#             continue

#         wp = cq.Workplane("XY").polyline(pts[:-1]).close().extrude(height)
#         solids.append(wp)

#     for c in circle_profiles:
#         wp = (
#             cq.Workplane("XY")
#             .center(c.center[0] * scale, c.center[1] * scale)
#             .circle(c.radius * scale)
#             .extrude(height)
#         )
#         solids.append(wp)

#     if not solids:
#         return None

#     #[TODO] 要修正
#     model = solids[0]
#     for s in solids[1:]:
#         model = model.union(s)

#     return model

from core import build_solid
from llm import generate_cadquery_with_openai, execute_generated_cadquery


def export_bytes(model) -> Tuple[bytes, bytes]:
    with tempfile.TemporaryDirectory() as td:
        stl_path = str(Path(td) / "result.stl")
        step_path = str(Path(td) / "result.step")
        exporters.export(model, stl_path)
        exporters.export(model, step_path)
        stl_bytes = Path(stl_path).read_bytes()
        step_bytes = Path(step_path).read_bytes()
    return stl_bytes, step_bytes


# =========================
# Preview
# =========================
def build_svg_preview(poly_profiles, circle_profiles, arc_previews, width=900, height=700) -> str:
    all_pts = []

    for p in poly_profiles:
        all_pts.extend(p.points)
    for c in circle_profiles:
        cx, cy = c.center
        r = c.radius
        all_pts.extend([(cx - r, cy - r), (cx + r, cy + r)])
    for arc in arc_previews:
        all_pts.extend(arc)

    if not all_pts:
        return "<svg xmlns='http://www.w3.org/2000/svg' width='900' height='700'></svg>"

    xs = [p[0] for p in all_pts]
    ys = [p[1] for p in all_pts]
    min_x, max_x = min(xs), max(xs)
    min_y, max_y = min(ys), max(ys)

    dx = max(max_x - min_x, 1e-6)
    dy = max(max_y - min_y, 1e-6)

    margin = 20
    sx = (width - 2 * margin) / dx
    sy = (height - 2 * margin) / dy
    s = min(sx, sy)

    def tr(p):
        x = margin + (p[0] - min_x) * s
        y = height - (margin + (p[1] - min_y) * s)
        return x, y

    out = [
        f"<svg xmlns='http://www.w3.org/2000/svg' width='{width}' height='{height}' viewBox='0 0 {width} {height}'>",
        "<rect width='100%' height='100%' fill='white'/>"
    ]

    # polylines
    for prof in poly_profiles:
        pts = [tr(p) for p in prof.points]
        d = " ".join(f"{x:.2f},{y:.2f}" for x, y in pts)
        out.append(f"<polyline points='{d}' fill='none' stroke='black' stroke-width='1.5'/>")

    # circles
    for c in circle_profiles:
        cx, cy = tr(c.center)
        r = c.radius * s
        out.append(f"<circle cx='{cx:.2f}' cy='{cy:.2f}' r='{r:.2f}' fill='none' stroke='black' stroke-width='1.5'/>")

    # arcs preview
    for arc in arc_previews:
        pts = [tr(p) for p in arc]
        d = " ".join(f"{x:.2f},{y:.2f}" for x, y in pts)
        out.append(f"<polyline points='{d}' fill='none' stroke='gray' stroke-width='1' stroke-dasharray='4,2'/>")

    out.append("</svg>")
    return "\n".join(out)


# =========================
# Streamlit UI
# =========================
st.set_page_config(page_title="2D -> 3D", layout="wide")
st.title("3D Generation from 2D CAD")
st.caption("2Dの図面を入力し、3Dに変換します。")

with st.sidebar:
    st.header("設定")
    height = st.number_input("押し出し高さ", min_value=0.1, value=100.0, step=1.0)
    scale = st.number_input("スケール", min_value=0.0001, value=1.0, step=0.1, format="%.4f")
    tol = st.number_input("接続許容差", min_value=1e-6, value=1e-4, step=1e-4, format="%.6f")
    st.markdown(
        """
        対応:
        - Closed LWPOLYLINE
        - Closed POLYLINE
        - CIRCLE
        - LINE の閉ループ化

        制限:
        - ARC は現在プレビューのみ
        - 穴(HOLE)の自動判定は未対応
        - HATCH / SPLINE / ELLIPSE は未対応
        """
    )

st.write(st.get_option("server.enableXsrfProtection"))
uploaded = st.file_uploader("DXFファイルをアップロード", type=["dxf"])

if uploaded is not None:
    try:
        doc = load_dxf_from_upload(uploaded)
        layers = get_layers(doc)
        selected_layers = st.selectbox(
            "対象レイヤー",
            options=layers,
        )

        poly_profiles, circle_profiles, arc_previews, warnings = extract_profiles(
            doc, selected_layers, tol=tol
        )

        c1, c2, c3 = st.columns(3)
        c1.metric("閉ポリライン数", len(poly_profiles))
        c2.metric("円数", len(circle_profiles))
        c3.metric("警告数", len(warnings))

        svg = build_svg_preview(poly_profiles, circle_profiles, arc_previews)
        st.subheader("2Dプレビュー")
        st.components.v1.html(svg, height=720, scrolling=True)

        if warnings:
            with st.expander("警告 / スキップ情報"):
                for w in warnings:
                    st.write("- " + w)

        if st.button("3Dモデルを生成"):
            with st.spinner("3Dモデルを生成中..."):
                # model = build_solid(
                #     poly_profiles=poly_profiles,
                #     circle_profiles=circle_profiles,
                #     height=height,
                #     scale=scale,
                # )

                from utils import render_dxf_layer_to_png
                render_dxf_layer_to_png(
                    doc = doc,
                    layer_name=selected_layers,
                    output_path="layer.png",
                )
                st.image("layer.png")

                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",
                )
                model = execute_generated_cadquery(generation["cadquery_code"])

            if model is None:
                st.error("3D化できる閉輪郭が見つかりませんでした。")
            else:
                st.success("3Dモデルを生成しました。")
                st.code(f"solids = {len(poly_profiles) + len(circle_profiles)}")

                stl_bytes, step_bytes = export_bytes(model)
                col1, col2 = st.columns(2)
                with col1:
                    st.download_button(
                        "STLをダウンロード",
                        data=stl_bytes,
                        file_name="result.stl",
                        mime="model/stl",
                    )
                with col2:
                    st.download_button(
                        "STEPをダウンロード",
                        data=step_bytes,
                        file_name="result.step",
                        mime="application/step",
                    )

    except Exception as e:
        st.exception(e)
else:
    st.info("まず DXF ファイルをアップロードしてください。")