TieMeasureFlow/src/vision_worker/main.py

"""The heavy container: the runner behind an internal API.

Separate from the FastAPI server on purpose. The API image stays light, a
VisionSuite upgrade does not restart production traffic, and an execution that
crashes does not take the other tablets' requests down with it.
"""
from __future__ import annotations

import asyncio
import io
import json
import os
import re
from pathlib import Path

import numpy as np
from fastapi import FastAPI, File, Form, HTTPException, UploadFile, status
from PIL import Image, UnidentifiedImageError

from src.vision.costruttori import CostruttoreError, build_dxf_compare  # #83
from src.vision.runner import (
    GraphBuildError,
    RunOutcome,
    engine_version,
    run_graph,
    validate_graph,
)

app = FastAPI(title="TieMeasureFlow Vision Worker", version="0.1.0")

# Below the backend's own 120s budget on purpose (#34): the worker must give
# up and answer 504 before the backend's httpx client times out on its own,
# so the operator sees a reason instead of a bare timeout.
_RUN_TIMEOUT_S = float(os.environ.get("VISION_RUN_TIMEOUT_S", "100"))
# Two slow graphs at once must not also stall a third tablet's request past
# any reasonable wait - reject it instead, with a status the backend can turn
# into a 503 the operator understands ("try again"), not a queued request
# that quietly outlasts anyone's patience.
_MAX_CONCURRENCY = int(os.environ.get("VISION_MAX_CONCURRENCY", "2"))
_run_slots = asyncio.Semaphore(_MAX_CONCURRENCY)


# #46: radice sotto cui vivono le cartelle delle ricette (volume uploads,
# montato in sola lettura nel worker). `recipe_dir` e' relativo a questa radice.
_RECIPES_ROOT_DEFAULT = "/app/uploads"


def _unprocessable(detail: str) -> HTTPException:
    return HTTPException(
        status_code=status.HTTP_422_UNPROCESSABLE_CONTENT, detail=detail
    )


def _resolve_recipe_dir(recipe_dir: str | None) -> Path | None:
    """Percorso RELATIVO sotto VISION_RECIPES_ROOT -> cartella assoluta.

    Assoluto, `..` o fuori radice (anche via symlink) -> 422. Vuoto -> None.
    """
    if recipe_dir is None or not recipe_dir.strip():
        return None
    rel = recipe_dir.strip()
    # Letto a ogni richiesta (non a import) cosi' i test possono cambiarlo.
    root = Path(os.environ.get("VISION_RECIPES_ROOT", _RECIPES_ROOT_DEFAULT)).resolve()
    if (
        Path(rel).is_absolute()
        # #46: lettera di unita (C:/x) rifiutata su ogni OS; UNC coperto da startswith("\\")
        or re.match(r"^[A-Za-z]:", rel)
        or rel.startswith(("/", "\\"))
        or ".." in Path(rel.replace("\\", "/")).parts
    ):
        raise _unprocessable("recipe_dir must be a relative path under the recipes root")
    resolved = (root / rel).resolve()
    if resolved != root and root not in resolved.parents:
        raise _unprocessable("recipe_dir is outside the recipes root")
    return resolved


def _parse_graph(graph: str) -> dict:
    try:
        parsed = json.loads(graph)
    except json.JSONDecodeError as exc:
        raise _unprocessable(f"graph is not valid JSON: {exc}") from exc
    if not isinstance(parsed, dict):
        raise _unprocessable(
            f"graph must be a JSON object, got {type(parsed).__name__}"
        )
    return parsed


def _parse_calibration(calibration: str | None) -> dict | None:
    if calibration is None or not calibration.strip():
        return None
    try:
        parsed = json.loads(calibration)
    except json.JSONDecodeError as exc:
        raise _unprocessable(f"calibration is not valid JSON: {exc}") from exc
    if not isinstance(parsed, dict) or not (
        parsed.get("mm_per_px") or parsed.get("px_per_mm")
    ):
        raise _unprocessable(
            "calibration must be {'mm_per_px': float} or {'px_per_mm': float}"
        )
    return parsed


class _ImageDecodeError(Exception):
    """The upload is not a decodable image - raised inside the worker thread
    so it survives the trip through `asyncio.to_thread`."""


@app.get("/health")
async def health() -> dict:
    """Deliberately just this: reading `engine_version()` (cached after the
    first call) is the only work here, so this answers even while `/run` is
    busy decoding an image or running a graph in its own thread."""
    return {"status": "ok", "engine_version": engine_version()}


def _decode_and_run(
    content: bytes,
    graph: dict,
    calibration: dict | None = None,
    recipe_dir: Path | None = None,
) -> RunOutcome:
    """The synchronous, CPU/IO-bound half of a request: PIL decode plus
    `run_graph`. Runs in a worker thread (see `run` below) so the event loop -
    and therefore `/health` - stays responsive while it executes. Once this
    thread is past its budget it is not interruptible: it keeps computing to
    completion in the background even after `run` has already answered 504
    for it (see the timeout handling there).
    """
    try:
        frame = np.array(Image.open(io.BytesIO(content)).convert("L"))
    except (UnidentifiedImageError, OSError) as exc:
        raise _ImageDecodeError(str(exc)) from exc
    return run_graph(frame, graph, calibration=calibration, recipe_dir=recipe_dir)


@app.post("/run")
async def run(
    image: UploadFile = File(...),
    graph: str = Form(...),
    recipe_dir: str | None = Form(None),
    calibration: str | None = Form(None),
) -> dict:
    parsed_graph = _parse_graph(graph)
    resolved_dir = _resolve_recipe_dir(recipe_dir)
    parsed_calibration = _parse_calibration(calibration)

    content = await image.read()

    if _run_slots.locked():
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail=f"worker busy: already running {_MAX_CONCURRENCY} job(s)",
        )

    async with _run_slots:
        try:
            outcome = await asyncio.wait_for(
                asyncio.to_thread(
                    _decode_and_run,
                    content,
                    parsed_graph,
                    parsed_calibration,
                    resolved_dir,
                ),
                timeout=_RUN_TIMEOUT_S,
            )
        except asyncio.TimeoutError as exc:
            raise HTTPException(
                status_code=status.HTTP_504_GATEWAY_TIMEOUT,
                detail=f"vision run exceeded {_RUN_TIMEOUT_S}s",
            ) from exc
        except _ImageDecodeError as exc:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail=f"image is not decodable: {exc}",
            ) from exc
        except GraphBuildError as exc:
            # The graph itself is bad (schema version, a tool missing
            # "kind", ...) - a bad request, not a server fault.
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail=str(exc),
            ) from exc
        except Exception as exc:
            # The graph built fine; something failed while actually running
            # it - an engine fault, not the caller's. Still named, not a
            # naked 500.
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"vision run failed: {type(exc).__name__}: {exc}",
            ) from exc

    return {
        # `outcome.outputs` never carries NaN/inf/bool (see runner.run_graph):
        # a JSON encoder that accepts non-finite floats (Python's stdlib json
        # does, by default) must never get the chance to emit them here.
        "outputs": outcome.outputs,
        "failures": [
            {"tool_id": f.tool_id, "tool_name": f.tool_name, "error": f.error}
            for f in outcome.failures
        ],
        "engine_version": outcome.engine_version,
        "duration_ms": outcome.duration_ms,
        "invalid_outputs": outcome.invalid_outputs,
        # #46: contorni per l'overlay sull'immagine e esito VisionSuite.
        "overlay": outcome.overlay,
        "outcome": outcome.outcome,
    }


@app.post("/validate")
async def validate(
    graph: str = Form(...),
    recipe_dir: str | None = Form(None),
) -> dict:
    """Validates a v2 recipe (not a draft: a missing meta file is an error)
    and lists its numeric outputs. No image, so it is cheap: no run slot."""
    parsed_graph = _parse_graph(graph)
    resolved_dir = _resolve_recipe_dir(recipe_dir)
    ok, errors, outputs = await asyncio.to_thread(
        validate_graph, parsed_graph, resolved_dir
    )
    return {"ok": ok, "errors": errors, "outputs": outputs}


def _decode_and_build_dxf(content: bytes, dxf: bytes, px_per_mm: float,
                          band_mm: float, min_points_pct: float,
                          polarity: str) -> dict:
    """#83: meta' sincrona di /recipes/dxf-compare (decodifica + costruzione),
    in un thread come `_decode_and_run`."""
    try:
        frame = np.array(Image.open(io.BytesIO(content)).convert("L"))
    except (UnidentifiedImageError, OSError) as exc:
        raise _ImageDecodeError(str(exc)) from exc
    return build_dxf_compare(frame, dxf, px_per_mm, band_mm, min_points_pct, polarity)


@app.post("/recipes/dxf-compare")
async def recipes_dxf_compare(
    image: UploadFile = File(...),
    dxf: UploadFile = File(...),
    px_per_mm: float = Form(...),
    band_mm: float = Form(0.1),
    min_points_pct: float = Form(95.0),
    polarity: str = Form("dark"),
) -> dict:
    """#83: costruisce una ricetta DXF compare dall'immagine di riferimento, il
    disegno e la scala. Stessi semaforo, timeout e mappa errori di /run."""
    content = await image.read()
    dxf_bytes = await dxf.read()

    if _run_slots.locked():
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail=f"worker busy: already running {_MAX_CONCURRENCY} job(s)",
        )

    async with _run_slots:
        try:
            return await asyncio.wait_for(
                asyncio.to_thread(
                    _decode_and_build_dxf, content, dxf_bytes, px_per_mm,
                    band_mm, min_points_pct, polarity,
                ),
                timeout=_RUN_TIMEOUT_S,
            )
        except asyncio.TimeoutError as exc:
            raise HTTPException(
                status_code=status.HTTP_504_GATEWAY_TIMEOUT,
                detail=f"vision run exceeded {_RUN_TIMEOUT_S}s",
            ) from exc
        except _ImageDecodeError as exc:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail=f"image is not decodable: {exc}",
            ) from exc
        except (CostruttoreError, GraphBuildError) as exc:
            # pezzo non trovato, DXF illeggibile, parametri invalidi: input, non guasto
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
                detail=str(exc),
            ) from exc
        except Exception as exc:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"recipe build failed: {type(exc).__name__}: {exc}",
            ) from exc