ImpTune/imptune/api/icons.py

"""Icon upload API — POST /printers/{printer_id}/icon."""
from __future__ import annotations

import hashlib
import io
from pathlib import Path

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

import imptune.config as cfg
from imptune.db.models import Icon, Printer

router = APIRouter(prefix="/printers")

MAX_ICON_BYTES = 750 * 1024  # 750 KB


@router.post("/{printer_id}/icon", response_class=HTMLResponse)
def upload_icon(request: Request, printer_id: int, file: UploadFile) -> HTMLResponse:
    """Accept a printer icon PNG, validate it, store it, and update the Icon record.

    Validation rules:
    - Format: PNG only
    - Dimensions: exactly 256x256 pixels
    - Size: at most 750 KB

    Replaces any previously uploaded icon for this printer.
    Returns an HTMX-friendly HTML fragment.
    """
    # Check printer exists and belongs to this owner
    printer = Printer.get_or_none(
        (Printer.id == printer_id) & (Printer.owner == request.state.owner)
    )
    if printer is None:
        return HTMLResponse(
            content="<p>Printer not found.</p>",
            status_code=404,
        )

    # Read file (read one byte extra to detect oversized files)
    data = file.file.read(MAX_ICON_BYTES + 1)
    if len(data) > MAX_ICON_BYTES:
        return HTMLResponse(
            content="<p>Icon exceeds 750 KB limit.</p>",
            status_code=422,
        )

    # Validate with Pillow
    try:
        img = Image.open(io.BytesIO(data))
    except Exception:
        return HTMLResponse(
            content="<p>Icon must be PNG format.</p>",
            status_code=422,
        )

    if img.format != "PNG":
        return HTMLResponse(
            content="<p>Icon must be PNG format.</p>",
            status_code=422,
        )

    if img.size != (256, 256):
        return HTMLResponse(
            content=f"<p>Icon must be 256x256 pixels, got {img.size}.</p>",
            status_code=422,
        )

    # Store SHA256-addressed on disk. cfg.ICONS_DIR is the same path the
    # .intunewin export reads from — deriving it a second time from DATA_DIR
    # would silently diverge if the layout ever changes.
    sha256 = hashlib.sha256(data).hexdigest()
    icons_dir = Path(cfg.ICONS_DIR)
    icons_dir.mkdir(parents=True, exist_ok=True)
    icon_path = icons_dir / sha256
    icon_path.write_bytes(data)

    # Replace existing Icon record for this printer
    Icon.delete().where(Icon.printer == printer_id).execute()
    Icon.create(
        printer=printer_id,
        sha256=sha256,
        original_filename=file.filename or "icon.png",
        size_bytes=len(data),
    )

    return HTMLResponse(
        content=(
            '<p class="ok-note">Icon uploaded successfully</p>'
            f'<div class="icon-preview"><img src="/printers/{printer_id}/icon?v={sha256[:8]}"'
            f' width="56" height="56" alt=""><span class="meta">{img.size[0]}&times;{img.size[1]} PNG</span></div>'
        ),
        status_code=200,
    )


@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.
    """
    printer = Printer.get_or_none(
        (Printer.id == printer_id) & (Printer.owner == request.state.owner)
    )
    if printer 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)

    icon_path = Path(cfg.ICONS_DIR) / icon.sha256
    if not icon_path.exists():
        return Response(status_code=404)

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