app.py

import os
import re
import sys
import base64
import shutil
import mimetypes
import subprocess
import tempfile
from pathlib import Path
import time

# =========================
# Paths / environment
# =========================
BASE_DIR = Path.cwd().resolve()
GRADIO_TEMP_DIR = BASE_DIR / "gradio_tmp"
ARTIFACT_DIR = BASE_DIR / "outputs"

GRADIO_TEMP_DIR.mkdir(parents=True, exist_ok=True)
ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)

os.environ["GRADIO_TEMP_DIR"] = str(GRADIO_TEMP_DIR)

import gradio as gr
from anthropic import Anthropic
from openai import OpenAI

# =========================
# Clients
# =========================

ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

if not ANTHROPIC_API_KEY:
    raise RuntimeError("ANTHROPIC_API_KEY is not set")

if not OPENAI_API_KEY:
    raise RuntimeError("OPENAI_API_KEY is not set")

anthropic_client = Anthropic(api_key=ANTHROPIC_API_KEY)
openai_client = OpenAI(api_key=OPENAI_API_KEY)

# =========================
# Model options
# =========================

CLAUDE_MODELS = [
    "claude-opus-4-7",
    "claude-opus-4-6",
    "claude-sonnet-4-6",
    "claude-sonnet-4-5",
    "claude-haiku-4-5-20251001"
]

OPENAI_MODELS = [
    "gpt-5.4-pro-2026-03-05",
    "gpt-5.2-2025-12-11",
    "gpt-5-nano-2025-08-07",
]

DEFAULT_CLAUDE_MODEL = "claude-opus-4-7"
DEFAULT_OPENAI_MODEL = "gpt-5-nano-2025-08-07"

# =========================
# Helpers
# =========================

def strip_code_fences(text: str) -> str:
    text = text.strip()
    match = re.search(r"```(?:python)?\s*(.*?)```", text, re.DOTALL | re.IGNORECASE)
    if match:
        return match.group(1).strip()
    return text


def detect_media_type(image_path: str) -> str:
    media_type, _ = mimetypes.guess_type(image_path)
    return media_type or "application/octet-stream"


def copy_to_artifact(src_path: Path, suffix: str) -> str:
    dst = ARTIFACT_DIR / f"{next(tempfile._get_candidate_names())}{suffix}"
    shutil.copy2(src_path, dst)
    return str(dst.resolve())


def truncate_text(text: str, max_chars: int = 4000) -> str:
    if not text:
        return ""
    return text[-max_chars:]

def append_log(log: str, message: str) -> str:
    return (log + message.rstrip() + "\n").strip() + "\n"

# =========================
# LLM steps
# =========================

def generate_cad_code(text_prompt, image_path, step_path, claude_model):
    content = []

    spec_text = (text_prompt or "").strip()
    if spec_text:
        content.append(
            {
                "type": "text",
                "text": f"""Generate executable CadQuery Python code to build a 3D CAD model.

Requirements:
- Use CadQuery only
- Save files in the current working directory
- Export BOTH final files:
  - ./output.step
  - ./output.stl
- ALSO export progressive STL snapshots during construction:
  - ./snap_001.stl
  - ./snap_002.stl
  - ./snap_003.stl
  - ...
- The script must be directly executable with `python model.py`
- Do not use markdown fences
- Output ONLY valid Python code
- At the end, ensure the exports are actually executed

Important implementation rules:
- Build the model in multiple clear steps
- Keep the current shape in a variable named `result`
- After every major modeling step, export `result` to a new snapshot STL
- Use zero-padded numbering like snap_001.stl
- At the end, export the final `result` to output.step and output.stl

Specification:
{spec_text}
""",
            }
        )

    if image_path:
        with open(image_path, "rb") as f:
            img_b64 = base64.b64encode(f.read()).decode("utf-8")

        content.append(
            {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": detect_media_type(image_path),
                    "data": img_b64,
                },
            }
        )
    
    if step_path:
        uploaded_file = anthropic_client.beta.files.upload(
            file=(
                Path(step_path).name,
                open(step_path, "rb"),
                "text/plain",
            )
        )
        content.append(
            {
                "type": "document",
                "source": {
                    "type": "file",
                    "file_id": uploaded_file.id,
                },
                "title": "input.step",
                "context": "This is a STEP CAD file in plain text format.",
            }
        )

    if not content:
        raise gr.Error("Text CAD specification or image is required.")

    response = anthropic_client.beta.messages.create(
        model=claude_model,
        max_tokens=16000,
        betas=["files-api-2025-04-14"],
        messages=[
            {
                "role": "user", 
                "content": content
            }
        ],
    )

    code = "".join(
        block.text
        for block in response.content
        if getattr(block, "type", None) == "text"
    ).strip()

    code = strip_code_fences(code)

    if not code:
        raise gr.Error("Claude returned empty code.")

    return code

def translate_to_japanese(text, openai_model):
    if not text or not text.strip():
        return ""

    response = openai_client.responses.create(
        model=openai_model,
        instructions=(
            "You are a professional translator. "
            "Translate the user's English text into natural, concise Japanese. "
            "Preserve structure, numbering, and technical meaning. "
            "Return only the Japanese translation."
        ),
        input=text,
    )

    return response.output_text.strip()

def parse_plan_and_code(response_text: str):
    plan_text = ""
    code_text = response_text.strip()

    if "[PLAN]" in response_text and "[CODE]" in response_text:
        before_code, after_code = response_text.split("[CODE]", 1)
        _, plan_part = before_code.split("[PLAN]", 1)
        plan_text = plan_part.strip()
        code_text = after_code.strip()

    code_text = strip_code_fences(code_text)
    return plan_text, code_text


def generate_cad_artifacts(text_prompt, image_path, step_path, claude_model):
    content = []

    spec_text = (text_prompt or "").strip()
    if spec_text:
        content.append(
            {
                "type": "text",
                "text": f"""Generate a 3D CAD model in CadQuery.

Return output in exactly this format:

[PLAN]
Write 4 to 8 concrete build steps for the model.
Each step should explain:
- what geometry or feature is being created
- why it is needed
- any important design intent, constraint, or tradeoff

The PLAN should be specific and descriptive, suitable for showing as a "Build Progress" log.

[CODE]
Return executable Python code using CadQuery only.

Requirements for the code:
- Use CadQuery only
- Save files in the current working directory
- Export BOTH final files:
  - ./output.step
  - ./output.stl
- ALSO export progressive STL snapshots during construction:
  - ./snap_001.stl
  - ./snap_002.stl
  - ./snap_003.stl
  - ...
- The script must be directly executable with `python model.py`
- Do not use markdown fences
- Output ONLY the PLAN block and the CODE block
- At the end, ensure the exports are actually executed

Important implementation rules:
- Build the model in multiple clear steps
- Keep the current shape in a variable named `result`
- After every major modeling step, export `result` to a new snapshot STL
- Use zero-padded numbering like snap_001.stl
- At the end, export the final `result` to output.step and output.stl
- The number and order of snapshot exports should roughly align with the PLAN steps

Specification:
{spec_text}
""",
            }
        )

    if image_path:
        with open(image_path, "rb") as f:
            img_b64 = base64.b64encode(f.read()).decode("utf-8")

        content.append(
            {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": detect_media_type(image_path),
                    "data": img_b64,
                },
            }
        )

    if step_path:
        uploaded_file = anthropic_client.beta.files.upload(
            file=(
                Path(step_path).name,
                open(step_path, "rb"),
                "text/plain",
            )
        )
        content.append(
            {
                "type": "document",
                "source": {
                    "type": "file",
                    "file_id": uploaded_file.id,
                },
                "title": "input.step",
                "context": "This is a STEP CAD file in plain text format.",
            }
        )

    if not content:
        raise gr.Error("Text CAD specification or image is required.")

    response = anthropic_client.beta.messages.create(
        model=claude_model,
        max_tokens=16000,
        betas=["files-api-2025-04-14"],
        messages=[
            {
                "role": "user", 
                "content": content
            }
        ],
    )

    raw_text = "".join(
        block.text
        for block in response.content
        if getattr(block, "type", None) == "text"
    ).strip()

    if not raw_text:
        raise gr.Error("Claude returned empty output.")

    plan_text, code = parse_plan_and_code(raw_text)

    if not code:
        raise gr.Error("Claude returned empty code.")

    return plan_text, code

def run_cadquery(code):
    with tempfile.TemporaryDirectory(dir=str(GRADIO_TEMP_DIR)) as tmpdir:
        tmpdir = Path(tmpdir)
        script_path = tmpdir / "model.py"

        with open(script_path, "w", encoding="utf-8") as f:
            f.write(code)

        result = subprocess.run(
            [sys.executable, str(script_path)],
            cwd=str(tmpdir),
            capture_output=True,
            text=True,
            timeout=180,
        )

        if result.returncode != 0:
            raise gr.Error(
                "CadQuery execution failed.\n\n"
                f"STDOUT:\n{truncate_text(result.stdout)}\n\n"
                f"STDERR:\n{truncate_text(result.stderr)}"
            )

        step_path = tmpdir / "output.step"
        stl_path = tmpdir / "output.stl"

        files = [p.name for p in tmpdir.iterdir()]

        if not step_path.exists() and not stl_path.exists():
            raise gr.Error(
                "CAD script finished but did not generate output.step or output.stl.\n\n"
                f"Files found in temp dir: {files}\n\n"
                "Make sure the generated CadQuery code exports exactly "
                "./output.step and ./output.stl"
            )

        if not step_path.exists():
            raise gr.Error(
                "output.step was not created.\n\n"
                f"Files found in temp dir: {files}"
            )

        if not stl_path.exists():
            raise gr.Error(
                "output.stl was not created.\n\n"
                f"Files found in temp dir: {files}"
            )

        final_step = copy_to_artifact(step_path, ".step")
        final_stl = copy_to_artifact(stl_path, ".stl")

    return final_step, final_stl



def run_cadquery_progressive(code):
    with tempfile.TemporaryDirectory(dir=str(GRADIO_TEMP_DIR)) as tmpdir:
        tmpdir = Path(tmpdir)
        script_path = tmpdir / "model.py"

        with open(script_path, "w", encoding="utf-8") as f:
            f.write(code)

        proc = subprocess.Popen(
            [sys.executable, str(script_path)],
            cwd=str(tmpdir),
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )

        seen = set()
        last_snapshot_path = None

        while True:
            snap_files = sorted(tmpdir.glob("snap_*.stl"))
            for snap in snap_files:
                if snap.name not in seen:
                    seen.add(snap.name)
                    copied = copy_to_artifact(snap, ".stl")
                    last_snapshot_path = copied
                    yield {
                        "viewer": copied,
                        "status": f"Rendering step {snap.stem}...",
                        "step_file": None,
                        "review": None,
                        "done": False,
                    }

            ret = proc.poll()
            if ret is not None:
                break

            time.sleep(0.4)

        stdout, stderr = proc.communicate()

        if proc.returncode != 0:
            raise gr.Error(
                "CadQuery execution failed.\n\n"
                f"STDOUT:\n{truncate_text(stdout)}\n\n"
                f"STDERR:\n{truncate_text(stderr)}"
            )

        step_path = tmpdir / "output.step"
        stl_path = tmpdir / "output.stl"
        files = [p.name for p in tmpdir.iterdir()]

        if not step_path.exists() or not stl_path.exists():
            raise gr.Error(
                "CAD script finished but final output.step/output.stl was missing.\n\n"
                f"Files found: {files}"
            )

        final_step = copy_to_artifact(step_path, ".step")
        final_stl = copy_to_artifact(stl_path, ".stl")

        yield {
            "viewer": final_stl,
            "status": "Final model completed.",
            "step_file": final_step,
            "review": None,
            "done": True,
        }


import base64
import tempfile
from pathlib import Path

def _stl_to_png(stl_path):
    try:
        import trimesh

        mesh = trimesh.load_mesh(stl_path)
        png = mesh.scene().save_image(resolution=(512, 512))

        # 👇 STEP/STLと同じ場所に保存
        out_path = Path(stl_path).with_suffix(".png")

        with open(out_path, "wb") as f:
            f.write(png)

        return str(out_path)

    except Exception:
        return None


def _image_to_data_url(path):
    with open(path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode("utf-8")
    return f"data:image/png;base64,{b64}"


def gpt_check(step_path, openai_model):
    if not step_path or not os.path.exists(step_path):
        return "STEP file not found, so review was skipped."

    size = os.path.getsize(step_path)

    # 🔹 STEP 読み込み（軽く truncate）
    try:
        with open(step_path, "r", encoding="utf-8", errors="ignore") as f:
            step_text = f.read()
    except Exception:
        step_text = ""

    # 🔹 STL を推定して PNG 化
    stl_path = Path(step_path).with_suffix(".stl")
    preview_png = _stl_to_png(str(stl_path)) if stl_path.exists() else None

    # 🔹 GPT に渡す内容
    content = [
        {
            "type": "input_text",
            "text": f"""A CAD STEP file was generated.

File size: {size} bytes

Below is part of the STEP file:

{step_text}

Tasks:
- Check if the STEP structure looks syntactically valid
- Point out possible corruption or missing sections
- Infer possible geometry issues from structure
- If an image is provided, also check visual problems (holes, broken shapes, weird proportions)

Keep it concise. 出力は日本語で。
""",
        }
    ]

    # 🔹 画像も渡す（あれば）
    if preview_png:
        content.append(
            {
                "type": "input_image",
                "image_url": _image_to_data_url(preview_png),
            }
        )

    response = openai_client.responses.create(
        model=openai_model,
        input=[{"role": "user", "content": content}],
    )

    return response.output_text.strip()



# =========================
# Pipeline
# =========================

def pipeline(text_prompt, image, step_file, claude_model, openai_model):
    log = ""

    if text_prompt:
        log += f"🧠 Request: {text_prompt[:120]}\n\n"
    else:
        log += "🧠 Request: image-based CAD generation\n\n"

    log += "🧠 Analyzing request...\n"

    log += "🧠 Thinking...\n Please wait for 10 seconds ... \n"
    yield None, None, log, None, None



    yield None, None, log, None, None

    plan_text, code = generate_cad_artifacts(text_prompt, image, step_file, claude_model)
    plan_text = translate_to_japanese(plan_text, openai_model)

    def extract_plan_steps(plan_text: str):
        steps = []
        for line in plan_text.splitlines():
            line = line.strip()
            if not line:
                continue
            if line[0].isdigit() or line.startswith("-"):
                steps.append(line)
        return steps
    plan_steps = extract_plan_steps(plan_text)

    log += "🏗 Build Progress\n"
    if not plan_steps:
        log += "- No explicit design plan was returned.\n"
    log += "\n✅ CAD code generated\n"
    log += "⚙️ Starting CAD build...\n"
    yield None, code, log, None, None

    final_step_path = None
    final_stl_path = None
    step_idx = 0

    for update in run_cadquery_progressive(code):
        viewer_value = update["viewer"]
        step_file_value = update["step_file"]

        if update["done"]:
            log += "▶ Final model completed.\n"
        else:
            step_idx += 1
            if step_idx <= len(plan_steps):
                log += f"▶ Step {step_idx}: {plan_steps[step_idx - 1]}\n"
            else:
                log += f"▶ Executing build snapshot {step_idx}\n"

        yield viewer_value, code, log, step_file_value, None

        if update["done"]:
            final_step_path = step_file_value
            final_stl_path = viewer_value
        
        time.sleep(5)

    review = "無し"

    yield final_stl_path, code, log, final_step_path, review

# =========================
# UI
# =========================

with gr.Blocks(delete_cache=(86400, 86400)) as demo:
    gr.Markdown("## 3D CAD AIエージェント🛠")

    with gr.Row():
        with gr.Column(scale=1):
            claude_model = gr.Radio(
                choices=CLAUDE_MODELS,
                value=DEFAULT_CLAUDE_MODEL,
                label="Claude model for CAD code generation",
                visible=False
            )

            openai_model = gr.Radio(
                choices=OPENAI_MODELS,
                value=DEFAULT_OPENAI_MODEL,
                label="OpenAI model for STEP review",
                visible=False
            )

            templates = {
                "自由入力": "",
                "ボルト": "A bolt with a hexagonal socket head and a cylindrical shaft",
                "筐体": "A 60mm x 40mm x 10mm enclosure with 4 corner holes",
                "修正": "修正箇所："
            }

            text_prompt = gr.Textbox(
                label="Text CAD specification\n仕様を記載してください",
                lines=3,
                placeholder="自由入力、またはテンプレートを選択"
            )

            template = gr.Dropdown(
                choices=list(templates.keys()),
                value="自由入力",
                label="テンプレート"
            )

            template.change(
                lambda x: templates[x],
                inputs=template,
                outputs=text_prompt,
            )

            with gr.Accordion("", open=False):
                image_input = gr.Image(
                    label="2D drawing (optional)\n2D図面をアップロードすることも可能です",
                    type="filepath"
                )

                step_input = gr.File(
                    label="STEPファイルをアップロード",
                    file_types=[".step", ".stp"],
                    type="filepath"
                )

            run_btn = gr.Button("Generate CAD")

        with gr.Column(scale=4):
            viewer = gr.Model3D(label="CAD Viewer (STL)", height=560, camera_position=(20, 30, 100))

        with gr.Column(scale=1):
            status_box = gr.Textbox(label="🤖", lines=6, autoscroll=True)
            

    review_box = gr.Textbox(label="Model Check", lines=8)
    with gr.Accordion("コード", open=False):
        code_box = gr.Code(label="Generated CAD Code", language="python")

    step_file = gr.File(label="Download STEP")

    edit_button = gr.Button("編集する")
    edit_message = gr.Markdown("")

    def use_output_as_input(step_output_path):
        if step_output_path is None:
            return None
        src = Path(step_output_path)
        output_dir = Path(tempfile.mkdtemp())
        copied_path = output_dir / src.name
        shutil.copy(src, copied_path)
        return str(copied_path), "✅ 編集可能になりました"

    edit_button.click(
        fn=use_output_as_input,
        inputs=[step_file],
        outputs=[step_input, edit_message]
    )

    run_btn.click(
        fn=pipeline,
        inputs=[text_prompt, image_input, step_input, claude_model, openai_model],
        outputs=[viewer, code_box, status_box, step_file, review_box]
    )

demo.launch(
    allowed_paths=[str(ARTIFACT_DIR), str(GRADIO_TEMP_DIR)]
)