ImpTune/imptune/api/icons.py

"""Printer icon API — upload, fetch-from-web, and serve."""
from __future__ import annotations

from fastapi import APIRouter, Form, Request, UploadFile
from fastapi.responses import FileResponse, HTMLResponse, Response

from imptune.db.models import Icon, Printer
from imptune.services import websearch
from imptune.services.icons import (
    MAX_ICON_BYTES,
    ImageError,
    IconTooLarge,
    icon_path,
    set_printer_icon,
)

router = APIRouter(prefix="/printers")


def _status_fragment(printer_id: int, sha256: str) -> HTMLResponse:
    """The `#icon-status` block, re-rendered after a successful save.

    `?v=` busts the 5-minute private cache the GET route sets — without it the
    preview keeps showing the icon that was just replaced.
    """
    return HTMLResponse(
        content=(
            # Swapped-in fragments are normally English-only; this one is on the
            # everyday path, so it goes through the i18n store with the English
            # text as its fallback body.
            '<p class="ok-note" x-data x-text="$store.i18n.t(\'icon_uploaded\')">'
            "Icon uploaded successfully</p>"
            f'<div class="icon-preview"><img src="/printers/{printer_id}/icon?v={sha256[:8]}"'
            ' width="56" height="56" alt="">'
            '<span class="meta">256&times;256 PNG</span></div>'
        ),
        status_code=200,
    )


def _error(message: str, status_code: int = 422) -> HTMLResponse:
    return HTMLResponse(content=f"<p class='error-note'>{message}</p>", status_code=status_code)


def _owned_printer(request: Request, printer_id: int) -> Printer | None:
    return Printer.get_or_none(
        (Printer.id == printer_id) & (Printer.owner == request.state.owner)
    )


@router.post("/{printer_id}/icon", response_class=HTMLResponse)
def upload_icon(request: Request, printer_id: int, file: UploadFile) -> HTMLResponse:
    """Accept a printer icon, normalize it to 256x256 PNG, and store it.

    Any raster Pillow can decode is accepted and letterboxed into the box —
    only an unreadable file or one over 750 KB is rejected. Replaces the
    printer's previous icon.
    """
    if _owned_printer(request, printer_id) is None:
        return HTMLResponse(content="<p>Printer not found.</p>", status_code=404)

    data = file.file.read(MAX_ICON_BYTES + 1)
    try:
        icon = set_printer_icon(printer_id, data, file.filename or "icon.png")
    except IconTooLarge as exc:
        return _error(str(exc))
    except ImageError:
        return _error("That file is not a readable image.")

    return _status_fragment(printer_id, icon.sha256)


@router.post("/{printer_id}/icon/from-web", response_class=HTMLResponse)
def upload_icon_from_web(
    request: Request, printer_id: int, url: str = Form(...)
) -> HTMLResponse:
    """Download a search-result image server-side and use it as the icon."""
    if _owned_printer(request, printer_id) is None:
        return HTMLResponse(content="<p>Printer not found.</p>", status_code=404)

    try:
        data = websearch.fetch_image(url)
    except websearch.WebSearchError as exc:
        return _error(str(exc), status_code=400)

    try:
        icon = set_printer_icon(printer_id, data, url.rsplit("/", 1)[-1][:120] or "web.png")
    except IconTooLarge as exc:
        return _error(str(exc))
    except ImageError:
        return _error("That URL did not return a readable image.")

    return _status_fragment(printer_id, icon.sha256)


@router.get("/{printer_id}/icon")
def get_icon(request: Request, printer_id: int) -> Response:
    """Serve the stored 256x256 PNG so the UI can show what was uploaded.

    Owner-scoped: a printer belonging to another owner reads as missing.
    """
    if _owned_printer(request, printer_id) is None:
        return Response(status_code=404)

    icon = Icon.get_or_none(Icon.printer == printer_id)
    if icon is None:
        return Response(status_code=404)

    path = icon_path(icon.sha256)
    if not path.exists():
        return Response(status_code=404)

    return FileResponse(
        path,
        media_type="image/png",
        headers={"Cache-Control": "private, max-age=300"},
    )