ImpTune/imptune/api/drivers.py

"""Driver API — upload, rename, and driver-library icons."""
from __future__ import annotations

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

import imptune.config as _cfg
from imptune.api.pages import build_driver_data, printer_counts_by_driver
from imptune.db.models import Driver, DriverIcon
from imptune.services import websearch
from imptune.services.icons import (
    MAX_ICON_BYTES,
    ImageError,
    IconTooLarge,
    icon_path,
    set_driver_icon,
)
from imptune.services.inf_parser import _detect_encoding, parse_inf
from imptune.storage.driver_store import DriverStore
from imptune.templating import templates

router = APIRouter(prefix="/drivers")

MAX_UPLOAD_BYTES = 100 * 1024 * 1024  # 100 MB
MAX_DISPLAY_NAME = 120


def _error_response(message: str, status_code: int = 400) -> HTMLResponse:
    """Return an HTMX-friendly error fragment swapped into #driver-list."""
    return HTMLResponse(
        content=f"<div id='driver-list' class='error'><p>{message}</p></div>",
        status_code=status_code,
    )


def _render_driver_list(request: Request, **extra) -> HTMLResponse:
    """Re-render the whole `#driver-list` table — the target every form swaps."""
    drivers = Driver.select().order_by(Driver.uploaded_at.desc())
    context = {
        "driver_data": build_driver_data(drivers),
        "usage": printer_counts_by_driver(request.state.owner),
    }
    context.update(extra)
    return templates.TemplateResponse(
        request=request, name="partials/driver_list.html", context=context
    )


@router.post("/upload", response_class=HTMLResponse)
def upload_driver(
    request: Request,
    file: UploadFile,
    caller: str = Form(""),
) -> HTMLResponse:
    """Accept a driver ZIP, parse its INF, persist via DriverStore + Peewee ORM.

    Returns an HTMX partial (partials/driver_list.html) on success, or an
    inline error fragment with HTTP 400 on validation failure.
    """
    data = file.file.read(MAX_UPLOAD_BYTES + 1)
    if len(data) > MAX_UPLOAD_BYTES:
        return _error_response("File exceeds 100 MB limit.")

    # Must end with .zip
    filename = file.filename or ""
    if not filename.lower().endswith(".zip"):
        return _error_response("Only .zip files are accepted.")

    # Must be a valid ZIP archive
    if not zipfile.is_zipfile(io.BytesIO(data)):
        return _error_response("Uploaded file is not a valid ZIP archive.")

    with zipfile.ZipFile(io.BytesIO(data)) as zf:
        zip_names = zf.namelist()

        # Reject zip-slip paths
        for name in zip_names:
            if ".." in name or name.startswith("/"):
                return _error_response("ZIP contains unsafe paths.")

        # Find .inf files
        inf_names = [n for n in zip_names if n.lower().endswith(".inf")]
        if not inf_names:
            return _error_response("No .inf file found in the uploaded ZIP.")

        # Prefer amd64/x64 INF when multiple exist; fall back to alphabetical first
        preferred = [
            n for n in inf_names if "amd64" in n.lower() or "x64" in n.lower()
        ]
        chosen_inf = preferred[0] if preferred else sorted(inf_names)[0]

        raw_inf = zf.read(chosen_inf)

    # Decode INF
    encoding = _detect_encoding(raw_inf)
    inf_text = raw_inf.decode(encoding)

    # Parse INF
    parsed = parse_inf(inf_text, inf_filename=chosen_inf, zip_names=zip_names)

    # Persist file (content-addressed, dedup automatic)
    # Read DRIVERS_DIR at call time so tests can monkeypatch imptune.config.DRIVERS_DIR
    store = DriverStore(_cfg.DRIVERS_DIR)
    sha256 = store.save(data)

    # Upsert Driver record (no duplicate if same SHA256)
    new_driver, _created = Driver.get_or_create(
        sha256=sha256,
        defaults={
            "original_filename": filename,
            "size_bytes": len(data),
            "driver_desc": json.dumps(parsed.driver_names),
            "inf_filename": parsed.inf_filename,
            "architecture": parsed.architecture,
            "has_cat_file": parsed.has_cat_file,
        },
    )

    # Build driver_data for template
    driver_data = build_driver_data(Driver.select().order_by(Driver.uploaded_at.desc()))

    # When called from the printer form, emit primary fragment + OOB select refresh
    if caller == "printer_form":
        return templates.TemplateResponse(
            request=request,
            name="partials/driver_upload_with_oob.html",
            context={
                "driver_data": driver_data,
                "new_driver_id": new_driver.id,
                "parsed": parsed,
            },
        )

    # Default: existing behavior — driver list fragment only
    return templates.TemplateResponse(
        request=request,
        name="partials/driver_list.html",
        context={
            "driver_data": driver_data,
            "parsed": parsed,
            "usage": printer_counts_by_driver(request.state.owner),
        },
    )


@router.patch("/{driver_id}", response_class=HTMLResponse)
def rename_driver(
    request: Request, driver_id: int, display_name: str = Form("")
) -> HTMLResponse:
    """Set or clear a driver's display name.

    Drivers are global/shared, so this rename is what every Owner sees — the
    same as the driver row itself. An empty value clears the rename and the
    listing falls back to the ZIP filename.
    """
    driver = Driver.get_or_none(Driver.id == driver_id)
    if driver is None:
        return _error_response(f"Driver {driver_id} not found.", status_code=404)

    name = display_name.strip()
    if len(name) > MAX_DISPLAY_NAME:
        return _error_response(
            f"Name must be at most {MAX_DISPLAY_NAME} characters."
        )

    driver.display_name = name or None
    driver.save()
    return _render_driver_list(request)


def _driver_icon_status(driver_id: int, sha256: str) -> HTMLResponse:
    """Inline confirmation for the rename/icon dialog, plus an OOB row refresh.

    The dialog stays open after picking an icon, so the thumbnail in the table
    row behind it is swapped out of band instead of re-rendering the table and
    tearing the open `<dialog>` out of the DOM.
    """
    src = f"/drivers/{driver_id}/icon?v={sha256[:8]}"
    return HTMLResponse(
        content=(
            '<p class="ok-note" x-data x-text="$store.i18n.t(\'icon_uploaded\')">'
            "Icon saved</p>"
            f'<div class="icon-preview"><img src="{src}" width="56" height="56" alt="">'
            '<span class="meta">256&times;256 PNG</span></div>'
            f'<span id="driver-thumb-{driver_id}" class="driver-thumb" hx-swap-oob="true">'
            f'<img src="{src}" width="24" height="24" alt=""></span>'
        ),
        status_code=200,
    )


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


@router.post("/{driver_id}/icon", response_class=HTMLResponse)
def upload_driver_icon(
    request: Request, driver_id: int, file: UploadFile
) -> HTMLResponse:
    """Attach an icon to a driver. Normalized to 256x256 PNG like printer icons."""
    if Driver.get_or_none(Driver.id == driver_id) is None:
        return _icon_error("Driver not found.", status_code=404)

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

    return _driver_icon_status(driver_id, icon.sha256)


@router.post("/{driver_id}/icon/from-web", response_class=HTMLResponse)
def driver_icon_from_web(
    request: Request, driver_id: int, url: str = Form(...)
) -> HTMLResponse:
    """Download a search-result image server-side and use it as the driver icon."""
    if Driver.get_or_none(Driver.id == driver_id) is None:
        return _icon_error("Driver not found.", status_code=404)

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

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

    return _driver_icon_status(driver_id, icon.sha256)


@router.get("/{driver_id}/icon")
def get_driver_icon(driver_id: int) -> Response:
    """Serve a driver icon. Not owner-scoped — the driver library is shared."""
    icon = DriverIcon.get_or_none(DriverIcon.driver == driver_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"},
    )