TieMeasureFlow/src/backend/api/routers/vision.py

"""The client asks for a measurement and gets an outcome.

It does not know, and must not know, whether the server or the station computed
the numbers: that is what lets a station's configuration change without touching
the frontend or the recipes.
"""
from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile, status
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from src.backend.api.middleware.api_key import require_maker, require_measurement_tec
from src.backend.api.routers.files import ALLOWED_IMAGE_TYPES, validate_file_size
from src.backend.config import settings
from src.backend.database import get_db
from src.backend.models.api.vision import VisionExecuteResponse, VisionMeasurementResult
from src.backend.models.orm.task import RecipeTask
from src.backend.models.orm.user import User
from src.backend.services import vision_service

router = APIRouter(prefix="/api/vision", tags=["vision"])

# A vision graph measures pixels, not a PDF: stricter than files.py's
# ALLOWED_TYPES, which also accepts documents for technical drawings.
_IMAGE_EXTENSIONS = {
    "image/jpeg": "jpg",
    "image/png": "png",
    "image/gif": "gif",
    "image/webp": "webp",
}

# Sniffed from the first bytes, independent of the declared Content-Type: a
# text file uploaded as "image/png" passes the header check above and must
# still be caught before it reaches the worker as a bare 500 (#15). BMP and
# TIFF are included even though `_IMAGE_EXTENSIONS` does not accept them,
# because PIL - and so the worker - can decode both; this check only refuses
# what is not a real image at all, it does not narrow the allow-list above.
_IMAGE_SIGNATURES: tuple[tuple[bytes, ...], ...] = (
    (b"\x89PNG\r\n\x1a\n",),         # PNG
    (b"\xff\xd8\xff",),              # JPEG
    (b"BM",),                        # BMP
    (b"II*\x00", b"MM\x00*"),        # TIFF, little- and big-endian
    (b"GIF87a", b"GIF89a"),          # GIF
)


def _looks_like_an_image(content: bytes) -> bool:
    """True if `content` starts with a known image format's magic bytes."""
    for signatures in _IMAGE_SIGNATURES:
        if any(content.startswith(sig) for sig in signatures):
            return True
    # WebP: "RIFF" + 4-byte size + "WEBP", the size in between is not fixed.
    return len(content) >= 12 and content[:4] == b"RIFF" and content[8:12] == b"WEBP"


async def _task_or_404(db: AsyncSession, task_id: int) -> RecipeTask:
    """Shared by the endpoints of this router."""
    task = (await db.execute(
        select(RecipeTask).where(RecipeTask.id == task_id)
    )).scalar_one_or_none()
    if task is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="Task not found"
        )
    return task


async def _read_image(image: UploadFile) -> bytes:
    """Type checked from the header before reading, size checked once the body
    is in hand - the same two-step convention `files.py::upload_file` already
    uses, reused rather than reinvented here.
    """
    if image.content_type not in ALLOWED_IMAGE_TYPES:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"File type {image.content_type} not allowed. Must be an image.",
        )
    content = await image.read()
    if not validate_file_size(len(content)):
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=(
                f"File size {len(content)} bytes exceeds maximum "
                f"{settings.max_upload_size_mb}MB"
            ),
        )
    if not _looks_like_an_image(content):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
            detail="the file's content does not match a supported image format",
        )
    return content


@router.post("/execute", response_model=VisionExecuteResponse)
async def execute(
    task_id: int = Form(...),
    image: UploadFile = File(...),
    lot_number: str | None = Form(None),
    serial_number: str | None = Form(None),
    production_run_id: int | None = Form(None),
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(require_measurement_tec),
) -> VisionExecuteResponse:
    task = await _task_or_404(db, task_id)

    result, saved = await vision_service.execute_task(
        db,
        task,
        await _read_image(image),
        current_user.id,
        version_id=task.version_id,
        lot_number=lot_number,
        serial_number=serial_number,
        production_run_id=production_run_id,
    )
    await db.commit()

    return VisionExecuteResponse(
        vision_result_id=result.id,
        engine_version=result.engine_version,
        measurements=[
            VisionMeasurementResult(
                subtask_id=m.subtask_id, value=float(m.value), pass_fail=m.pass_fail,
            )
            for m in saved
        ],
    )


# Composing a recipe and keeping a reference image are Maker operations, not
# measuring: both endpoints below gate on `require_maker`, not the
# `require_measurement_tec` that `execute` above uses.


@router.post("/reference-images")
async def add_reference_image(
    task_id: int = Form(...),
    image: UploadFile = File(...),
    note: str | None = Form(None),
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(require_maker),
) -> dict:
    task = await _task_or_404(db, task_id)
    content = await _read_image(image)
    reference = await vision_service.save_reference_image(
        db, task.id, content, note=note,
        extension=_IMAGE_EXTENSIONS[image.content_type],
    )
    await db.commit()
    return {"id": reference.id, "path": reference.path, "note": reference.note}


@router.post("/preview")
async def preview(
    task_id: int = Form(...),
    image: UploadFile = File(...),
    db: AsyncSession = Depends(get_db),
    current_user: User = Depends(require_maker),
) -> dict:
    task = await _task_or_404(db, task_id)
    return await vision_service.preview(db, task, await _read_image(image))