AI_portal/dwg-rooms/main.py

"""FastAPI app: upload DWG/DXF → extract rooms → export Excel."""
import json
import logging
import os
import subprocess
import uuid
from pathlib import Path

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

from excel_export import export_to_excel
from extractor import DEFAULT_EXAMPLES, extract_rooms
from llm_helper import enhance_with_llm

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="DWG Room Extractor")
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

WORK_DIR = Path(os.getenv("WORK_DIR", "/tmp/dwg-rooms"))
WORK_DIR.mkdir(parents=True, exist_ok=True)

jobs: dict[str, dict] = {}


def _dwg_to_dxf(dwg_path: Path, out_dir: Path) -> Path:
    """Convert DWG → DXF using LibreDWG dwgread. Returns DXF path."""
    dxf_path = out_dir / f"{dwg_path.stem}.dxf"
    r = subprocess.run(
        ["dwgread", "-O", "DXF", "-o", str(dxf_path), str(dwg_path)],
        capture_output=True,
        text=True,
        timeout=120,
    )
    if not dxf_path.exists():
        raise RuntimeError(
            f"DWG conversion failed (exit {r.returncode}): {r.stderr or r.stdout or 'no output'}"
        )
    return dxf_path


@app.get("/")
async def root():
    return FileResponse("static/index.html")


@app.get("/api/defaults")
async def defaults():
    return {"examples": DEFAULT_EXAMPLES}


@app.post("/api/upload")
async def upload(
    file: UploadFile = File(...),
    examples: str = Form(default=""),  # JSON array of strings
):
    suffix = Path(file.filename or "input.dxf").suffix.lower()
    if suffix not in (".dwg", ".dxf"):
        raise HTTPException(400, "Supported formats: .dwg, .dxf")

    try:
        ex_list = json.loads(examples) if examples else None
        if ex_list is not None and not isinstance(ex_list, list):
            ex_list = None
    except json.JSONDecodeError:
        ex_list = None

    job_id = str(uuid.uuid4())
    job_dir = WORK_DIR / job_id
    job_dir.mkdir()

    input_path = job_dir / f"input{suffix}"
    content = await file.read()
    input_path.write_bytes(content)
    logger.info("Job %s: %s (%d bytes), examples=%r", job_id, file.filename, len(content), ex_list)

    steps: list[str] = []

    try:
        if suffix == ".dwg":
            steps.append("Konverze DWG → DXF")
            dxf_path = _dwg_to_dxf(input_path, job_dir)
        else:
            dxf_path = input_path

        steps.append("Extrakce místností (pravidla)")
        rooms, unmatched = extract_rooms(str(dxf_path), examples=ex_list)
        logger.info("Job %s: %d rooms, %d unmatched texts", job_id, len(rooms), len(unmatched))

        if unmatched and os.getenv("LITELLM_API_KEY", ""):
            steps.append("AI rozšíření")
            llm_rooms = await enhance_with_llm(unmatched)
            logger.info("Job %s: LLM added %d rooms", job_id, len(llm_rooms))
            rooms.extend(llm_rooms)

        jobs[job_id] = {"filename": file.filename, "rooms": rooms}
        return {"job_id": job_id, "room_count": len(rooms), "steps": steps}

    except Exception as exc:
        logger.error("Job %s failed: %s", job_id, exc)
        raise HTTPException(500, str(exc))


@app.get("/api/rooms/{job_id}")
async def get_rooms(job_id: str):
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    return jobs[job_id]["rooms"]


@app.put("/api/rooms/{job_id}")
async def update_rooms(job_id: str, rooms: list[dict]):
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    jobs[job_id]["rooms"] = rooms
    return {"ok": True, "count": len(rooms)}


@app.get("/api/export/{job_id}")
async def export(job_id: str):
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    job = jobs[job_id]
    excel_path = WORK_DIR / job_id / "rooms.xlsx"
    export_to_excel(job["rooms"], str(excel_path))
    stem = Path(job["filename"]).stem
    return FileResponse(
        str(excel_path),
        media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        filename=f"mistnosti_{stem}.xlsx",
    )


@app.get("/health")
async def health():
    return {"status": "ok"}


app.mount("/static", StaticFiles(directory="static"), name="static")