AI_portal/dwg-counting/main.py

"""FastAPI app: upload DWG/DXF/PDF → vision-detect legend → count selected symbols → Excel."""
import asyncio
import logging
import os
import uuid
from pathlib import Path

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

from counting import count_template, debug_template
from excel_export import export_to_excel
from pdf_export import render_annotated_pdf
from renderer import crop_region, render
from vision import detect_legend

DEFAULT_FLOOR_INDEX = 0  # MVP: process the first detected floor

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

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

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

jobs: dict[str, dict] = {}
# Persistent action log so the operator can replay what the user did.
ACTION_LOG = WORK_DIR / "action.log"


def _log_action(event: str, **fields):
    import datetime, json as _json
    line = _json.dumps({
        "ts": datetime.datetime.now().isoformat(timespec="seconds"),
        "event": event,
        **fields,
    }, ensure_ascii=False)
    try:
        with open(ACTION_LOG, "a") as f:
            f.write(line + "\n")
    except Exception:
        pass


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


@app.post("/api/upload")
async def upload(file: UploadFile = File(...), auto_detect: bool = False):
    suffix = Path(file.filename or "").suffix.lower()
    if suffix != ".pdf":
        raise HTTPException(400, "Podporovaný formát: .pdf")

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

    input_path = job_dir / f"input{suffix}"
    input_path.write_bytes(await file.read())
    logger.info("Job %s: %s (%d bytes)", job_id, file.filename, input_path.stat().st_size)
    _log_action("upload", job_id=job_id, filename=file.filename,
                size=input_path.stat().st_size, auto_detect=auto_detect)

    try:
        rendered = await asyncio.to_thread(render, input_path, job_dir)
        floors = rendered["floors"]
        if not floors:
            raise RuntimeError("Z výkresu se nepodařilo nic vyrenderovat")
        floor = floors[DEFAULT_FLOOR_INDEX]
        png_path = job_dir / floor["png"]
        legend_norm = floor.get("legend_norm_bbox")
        legend_pixel_box = None
        symbols = []
        detect_path = png_path

        if legend_norm:
            from PIL import Image as _Img
            full_img = _Img.open(png_path)
            W, H = full_img.size
            nx0, ny0, nx1, ny1 = legend_norm
            # Expand box generously around the LEGENDA text
            cx = (nx0 + nx1) / 2
            cy = (ny0 + ny1) / 2
            # Legend rows extend DOWNWARD from the LEGENDA header text.
            # Crop a narrow column starting just above the header.
            half_w = 0.10
            top_pad = 0.02
            below = 0.30
            px0 = max(0, int((cx - half_w) * W))
            px1 = min(W, int((cx + half_w) * W))
            py0 = max(0, int((cy - top_pad) * H))
            py1 = min(H, int((cy + below) * H))
            legend_crop = job_dir / "legend_area.png"
            crop_img = full_img.crop((px0, py0, px1, py1))
            crop_img.save(legend_crop, "PNG")
            detect_path = legend_crop
            legend_pixel_box = (px0, py0, px1 - px0, py1 - py0)
            logger.info("Legend area %dx%d ready", px1 - px0, py1 - py0)

        if auto_detect:
            symbols = await detect_legend(detect_path)

        # Symbol bboxes returned by vision are normalized to the image vision saw
        # (detect_path), which may be the cropped legend region or the full page.
            for s in symbols:
                bbox = s.get("bbox") or {}
                if all(k in bbox for k in ("x", "y", "w", "h")):
                    crop_path = job_dir / f"sym_{s['id']}.png"
                    try:
                        crop_region(detect_path, bbox, crop_path, pad=1.5, min_px=120)
                        s["crop_file"] = crop_path.name
                    except Exception as exc:
                        logger.warning("Crop failed for %s: %s", s.get("id"), exc)

        jobs[job_id] = {
            "filename": file.filename,
            "png_path": str(png_path),
            "job_dir": str(job_dir),
            "floors": floors,
            "symbols": symbols,
            "results": [],
            "legend_pixel_box": legend_pixel_box,
            "next_user_sym_id": 1,
        }
        return {"job_id": job_id, "symbols": symbols, "floor_count": len(floors),
                "auto_detect": auto_detect}

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


@app.get("/api/preview/{job_id}")
async def preview(job_id: str):
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    return FileResponse(jobs[job_id]["png_path"], media_type="image/png")


@app.get("/api/symbol/{job_id}/{sym_id}")
async def symbol_crop(job_id: str, sym_id: str):
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    job_dir = Path(jobs[job_id]["job_dir"])
    crop_path = job_dir / f"sym_{sym_id}.png"
    if not crop_path.exists():
        raise HTTPException(404, "Crop not available")
    return FileResponse(crop_path, media_type="image/png")


@app.get("/api/legend/{job_id}")
async def legend_image(job_id: str):
    """Return the cropped legend area image (what vision saw)."""
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    legend_path = Path(jobs[job_id]["job_dir"]) / "legend_area.png"
    if not legend_path.exists():
        # Fall back to full page if no legend crop
        legend_path = Path(jobs[job_id]["png_path"])
    return FileResponse(legend_path, media_type="image/png")


class RecropRequest(BaseModel):
    bbox: dict  # {x, y, w, h} normalized 0-1 relative to the legend image


class CreateSymbolRequest(BaseModel):
    bbox: dict  # normalized 0-1 relative to the FULL drawing image
    description: str
    source: str = "drawing"  # "drawing" or "legend"


@app.post("/api/symbols/{job_id}")
async def create_user_symbol(job_id: str, req: CreateSymbolRequest):
    """User drew a rectangle on the drawing → create a symbol from that crop."""
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    job = jobs[job_id]
    job_dir = Path(job["job_dir"])
    # Pick source image
    if req.source == "legend":
        src = job_dir / "legend_area.png"
        if not src.exists():
            src = Path(job["png_path"])
    else:
        src = Path(job["png_path"])

    sym_id = f"user_{job['next_user_sym_id']}"
    job["next_user_sym_id"] += 1
    crop_path = job_dir / f"sym_{sym_id}.png"
    # User's rectangle is exact — no padding, no min size enforcement.
    crop_region(src, req.bbox, crop_path, pad=0.0, min_px=0)
    sym = {
        "id": sym_id,
        "description": req.description or sym_id,
        "bbox": req.bbox,
        "crop_file": crop_path.name,
        "user_defined": True,
    }
    job["symbols"].append(sym)
    _log_action("create_symbol", job_id=job_id, sym_id=sym_id,
                description=req.description, bbox=req.bbox, source=req.source)
    return sym


@app.post("/api/symbols/{job_id}/upload")
async def upload_symbol_image(
    job_id: str,
    file: UploadFile = File(...),
    description: str = "",
):
    """Accept a pre-cropped symbol image as a template, bypassing the
    rectangle-drawing UI. Useful when the user has a clean PNG from elsewhere.
    """
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    job = jobs[job_id]
    job_dir = Path(job["job_dir"])
    sym_id = f"user_{job['next_user_sym_id']}"
    job["next_user_sym_id"] += 1
    crop_path = job_dir / f"sym_{sym_id}.png"
    raw = await file.read()
    crop_path.write_bytes(raw)
    # Normalize: ensure RGB on white background (drop alpha so processing
    # doesn't see transparency as "not white")
    from PIL import Image as _Img
    img = _Img.open(crop_path)
    if img.mode in ("RGBA", "LA"):
        bg = _Img.new("RGB", img.size, (255, 255, 255))
        bg.paste(img, mask=img.split()[-1])
        bg.save(crop_path, "PNG")
    sym = {
        "id": sym_id,
        "description": (description or file.filename or sym_id).strip(),
        "crop_file": crop_path.name,
        "user_defined": True,
        "uploaded": True,
    }
    job["symbols"].append(sym)
    _log_action("upload_symbol", job_id=job_id, sym_id=sym_id,
                description=sym["description"], filename=file.filename,
                size=len(raw))
    return sym


@app.delete("/api/symbols/{job_id}/{sym_id}")
async def delete_symbol(job_id: str, sym_id: str):
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    job = jobs[job_id]
    job["symbols"] = [s for s in job["symbols"] if s["id"] != sym_id]
    crop = Path(job["job_dir"]) / f"sym_{sym_id}.png"
    if crop.exists():
        crop.unlink()
    return {"ok": True}


@app.post("/api/auto-detect/{job_id}")
async def trigger_auto_detect(job_id: str):
    """Run the vision legend detection on demand (optional shortcut)."""
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    job = jobs[job_id]
    job_dir = Path(job["job_dir"])
    legend_path = job_dir / "legend_area.png"
    detect_path = legend_path if legend_path.exists() else Path(job["png_path"])
    found = await detect_legend(detect_path)
    for s in found:
        bbox = s.get("bbox") or {}
        if all(k in bbox for k in ("x", "y", "w", "h")):
            crop_path = job_dir / f"sym_{s['id']}.png"
            try:
                crop_region(detect_path, bbox, crop_path, pad=1.5, min_px=120)
                s["crop_file"] = crop_path.name
            except Exception as exc:
                logger.warning("Crop failed for %s: %s", s.get("id"), exc)
    # Replace vision-detected ones (keep user-defined)
    job["symbols"] = [s for s in job["symbols"] if s.get("user_defined")] + found
    return {"symbols": job["symbols"]}


@app.get("/api/drawing/{job_id}")
async def drawing_image(job_id: str):
    """Return the rendered full drawing image (for the user to crop on)."""
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    return FileResponse(jobs[job_id]["png_path"], media_type="image/png")


@app.get("/api/debug/{job_id}/{sym_id}")
async def debug_symbol(job_id: str, sym_id: str):
    """Diagnostics about a symbol's template: size, ink, match scores."""
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    job = jobs[job_id]
    crop = Path(job["job_dir"]) / f"sym_{sym_id}.png"
    if not crop.exists():
        raise HTTPException(404, "Crop not available")
    drawing = Path(job["png_path"])
    info = await asyncio.to_thread(debug_template, crop, drawing)
    return info


@app.get("/api/debug-template/{job_id}/{sym_id}")
async def debug_template_image(job_id: str, sym_id: str):
    """Return the *processed* template (what the matcher actually sees)."""
    import cv2
    from counting import _prep, _crop_to_content
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    job = jobs[job_id]
    crop = Path(job["job_dir"]) / f"sym_{sym_id}.png"
    if not crop.exists():
        raise HTTPException(404, "Crop not available")
    tmpl = _prep(crop)
    tmpl = _crop_to_content(tmpl)
    out = Path(job["job_dir"]) / f"sym_{sym_id}_processed.png"
    cv2.imwrite(str(out), tmpl)
    return FileResponse(str(out), media_type="image/png")


@app.post("/api/symbol/{job_id}/{sym_id}/recrop")
async def recrop_symbol(job_id: str, sym_id: str, req: RecropRequest):
    """Replace a symbol's crop. bbox is normalized 0-1 relative to the FULL
    drawing image (which is what the frontend shows in the editor)."""
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    job = jobs[job_id]
    src = Path(job["png_path"])
    crop_path = Path(job["job_dir"]) / f"sym_{sym_id}.png"
    crop_region(src, req.bbox, crop_path, pad=0.0, min_px=0)
    return {"ok": True, "crop_file": crop_path.name}


class CountRequest(BaseModel):
    symbol_ids: list[str]
    threshold: float | None = None  # Override default (0.7); lower = more matches


@app.post("/api/count/{job_id}")
async def count(job_id: str, req: CountRequest):
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    job = jobs[job_id]
    png = Path(job["png_path"])
    job_dir = Path(job["job_dir"])
    selected = [s for s in job["symbols"] if s["id"] in req.symbol_ids]

    # Determine the legend mask box (avoid matching the legend itself).
    legend_box = job.get("legend_pixel_box")  # set during upload if available

    thr = req.threshold if req.threshold is not None else None

    def _count_one(sym):
        crop = job_dir / f"sym_{sym['id']}.png"
        if not crop.exists():
            return {"id": sym["id"], "description": sym["description"],
                    "count": 0, "matches": [], "notes": "no crop"}
        try:
            kwargs = {"exclude_box": legend_box}
            if thr is not None:
                kwargs["threshold"] = thr
            res = count_template(crop, png, **kwargs)
        except Exception as exc:
            logger.exception("Counting failed for %s", sym["id"])
            return {"id": sym["id"], "description": sym["description"],
                    "count": 0, "matches": [], "notes": f"error: {exc}"}
        return {
            "id": sym["id"],
            "description": sym["description"],
            "count": res["count"],
            "matches": res["matches"],
            "notes": "" if res["count"] else "žádné shody nenalezeny",
        }

    # Serialize OpenCV calls — parallel matchTemplate on a 4000px drawing
    # blows past 4GB peak memory and OOM-kills the container.
    results = []
    for s in selected:
        r = await asyncio.to_thread(_count_one, s)
        results.append(r)
        _log_action("count_one", job_id=job_id, sym_id=s["id"],
                    description=s.get("description"), count=r.get("count"))
    job["results"] = list(results)
    # Trim matches from API response (keep them server-side for PDF export)
    response_results = [{k: v for k, v in r.items() if k != "matches"} | {"count": r["count"]}
                       for r in job["results"]]
    return {"results": response_results}


@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]
    if not job["results"]:
        raise HTTPException(400, "Nejprve spočítejte symboly")
    out_path = Path(job["job_dir"]) / "counts.xlsx"
    export_to_excel(job["results"], job["filename"] or "drawing", str(out_path))
    stem = Path(job["filename"]).stem if job["filename"] else "drawing"
    return FileResponse(
        str(out_path),
        media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        filename=f"symboly_{stem}.xlsx",
    )


@app.get("/api/export-pdf/{job_id}")
async def export_pdf(job_id: str):
    if job_id not in jobs:
        raise HTTPException(404, "Not found")
    job = jobs[job_id]
    if not job["results"]:
        raise HTTPException(400, "Nejprve spočítejte symboly")
    png_path = Path(job["png_path"])
    stem = Path(job["filename"]).stem if job["filename"] else "drawing"
    out_path = Path(job["job_dir"]) / "annotated.pdf"
    await asyncio.to_thread(
        render_annotated_pdf, png_path, job["results"], out_path, stem,
    )
    return FileResponse(
        str(out_path),
        media_type="application/pdf",
        filename=f"vyznaceno_{stem}.pdf",
    )


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


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