step2urdf/app/main.py

from __future__ import annotations

import shutil
import tempfile
from pathlib import Path

from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles

from app.config import settings
from app.models import ExportRequest, ProposeRequest
from app.services.mesh_export import geometry_backend_status
from app.services.pipeline import create_job_from_upload, export_package, load_parse, propose_draft

ROOT = Path(__file__).resolve().parents[1]
LEGACY_STATIC = ROOT / "app" / "static"
FRONTEND_DIST = ROOT / "frontend" / "dist"

app = FastAPI(title="step2urdf-tool", version="0.2.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://127.0.0.1:5678",
        "http://localhost:5678",
        "http://127.0.0.1:8787",
        "http://localhost:8787",
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/api/health")
async def health():
    return {
        "ok": True,
        "zhipu_configured": bool(settings.zhipu_api_key),
        "zhipu_model": settings.zhipu_model,
        "geometry": geometry_backend_status(),
        "primary_ui": "step2urdf frontend (Vite :5678 in dev)",
    }


@app.post("/api/parse")
async def parse_step(file: UploadFile = File(...)):
    """Optional server-side STEP text parse (legacy / debug). Geometry stays in the browser."""
    suffix = Path(file.filename or "model.STEP").suffix or ".STEP"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        tmp_path = Path(tmp.name)
        shutil.copyfileobj(file.file, tmp)
    try:
        result = create_job_from_upload(tmp_path, file.filename or "model.STEP")
    finally:
        tmp_path.unlink(missing_ok=True)
    return result


@app.get("/api/jobs/{job_id}")
async def get_job(job_id: str):
    try:
        return load_parse(job_id)
    except FileNotFoundError as e:
        raise HTTPException(404, str(e)) from e


@app.post("/api/propose")
async def propose(req: ProposeRequest):
    """Zhipu-assisted link/joint proposal.

    Preferred body (from step2urdf UI)::

        {"part_names": [...], "profile": "a7", "robot_name": "A7", "extra_hint": "..."}

    Legacy: ``job_id`` from ``/api/parse``.
    """
    names = list(req.part_names or [])
    if req.solid_names:
        # Deduplicate while preserving order
        seen: set[str] = set()
        merged: list[str] = []
        for n in names + list(req.solid_names):
            if n and n not in seen:
                seen.add(n)
                merged.append(n)
        names = merged

    try:
        draft = await propose_draft(
            req.job_id,
            profile=req.profile,
            robot_name=req.robot_name,
            extra_hint=req.extra_hint,
            part_names=names or None,
            solids_geom=list(req.solids or []),
        )
    except FileNotFoundError as e:
        raise HTTPException(404, str(e)) from e
    except ValueError as e:
        raise HTTPException(400, str(e)) from e
    except Exception as e:  # noqa: BLE001
        raise HTTPException(502, f"Zhipu/propose failed: {e}") from e
    return draft


@app.post("/api/export")
async def export(req: ExportRequest):
    """Legacy server-side URDF write. Prefer Export URDF in the step2urdf UI."""
    try:
        meta = export_package(req.job_id, req.draft, include_meshes=req.include_meshes)
    except FileNotFoundError as e:
        raise HTTPException(404, str(e)) from e
    return meta


@app.get("/api/jobs/{job_id}/download/{robot_name}")
async def download_urdf(job_id: str, robot_name: str):
    path = settings.jobs_dir / job_id / "export" / robot_name / f"{robot_name}.urdf"
    if not path.exists():
        raise HTTPException(404, "URDF not found; export first")
    return FileResponse(path, filename=f"{robot_name}.urdf", media_type="application/xml")


@app.get("/")
async def root():
    if FRONTEND_DIST.is_dir() and (FRONTEND_DIST / "index.html").exists():
        return FileResponse(FRONTEND_DIST / "index.html")
    return RedirectResponse(url="/legacy/")


# Built step2urdf UI (after `pnpm build` in frontend/)
if FRONTEND_DIST.is_dir():
    app.mount("/assets", StaticFiles(directory=str(FRONTEND_DIST / "assets")), name="frontend-assets")

# Previous MVP static UI kept under /legacy only
if LEGACY_STATIC.is_dir():
    app.mount("/legacy", StaticFiles(directory=str(LEGACY_STATIC), html=True), name="legacy")


def main() -> None:
    import uvicorn

    uvicorn.run(
        "app.main:app",
        host=settings.host,
        port=settings.port,
        reload=True,
    )


if __name__ == "__main__":
    main()