ImpTune/imptune/api/pages.py

import json
from collections import defaultdict
from pathlib import Path

from fastapi import APIRouter, Request
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates
from peewee import JOIN, fn


router = APIRouter()

templates = Jinja2Templates(directory=str(Path(__file__).parent.parent / "templates"))

# Group label for printers with no client — the template translates it.
UNASSIGNED_GROUP = "Unassigned"


def printer_counts_by_driver(owner) -> dict[int, int]:
    """How many of this owner's printers use each (globally shared) driver."""
    from imptune.db.models import Printer

    rows = (
        Printer.select(Printer.driver, fn.COUNT(Printer.id).alias("n"))
        .where((Printer.owner == owner) & Printer.driver.is_null(False))
        .group_by(Printer.driver)
    )
    return {row.driver_id: row.n for row in rows}


def printer_counts_by_client(owner) -> dict[int, int]:
    """How many printers each client groups."""
    from imptune.db.models import Printer

    rows = (
        Printer.select(Printer.client, fn.COUNT(Printer.id).alias("n"))
        .where((Printer.owner == owner) & Printer.client.is_null(False))
        .group_by(Printer.client)
    )
    return {row.client_id: row.n for row in rows}


def group_printers_by_client(printers) -> dict[str, list]:
    """Group printers under their client name, unassigned ones last.

    The template renders groups in insertion order, and "a printer nobody has
    filed yet" belongs at the bottom of the page, not the top.
    """
    grouped: dict[str, list] = defaultdict(list)
    for p in printers:
        grouped[p.client.name if p.client_id else UNASSIGNED_GROUP].append(p)

    unassigned = grouped.pop(UNASSIGNED_GROUP, None)
    ordered = {name: grouped[name] for name in sorted(grouped)}
    if unassigned:
        ordered[UNASSIGNED_GROUP] = unassigned
    return ordered


def _driver_data(drivers) -> list[dict]:
    """Attach the parsed driver names to each Driver row for template use."""
    return [
        {"driver": d, "names": json.loads(d.driver_desc) if d.driver_desc else []}
        for d in drivers
    ]


@router.get("/", response_class=HTMLResponse)
def dashboard(request: Request):
    from imptune.db.models import Client, Driver, Printer

    owner = request.state.owner
    recent_printers = list(
        Printer.select(Printer, Client)
        .join(Client, JOIN.LEFT_OUTER)
        .where(Printer.owner == owner)
        .order_by(Printer.created_at.desc())
        .limit(5)
    )
    recent_packages = list(
        Printer.select(Printer, Client)
        .join(Client, JOIN.LEFT_OUTER)
        .switch(Printer)
        .where((Printer.owner == owner) & Printer.driver.is_null(False))
        .order_by(Printer.created_at.desc())
        .limit(5)
    )
    printer_count = Printer.select().where(Printer.owner == owner).count()
    ready_count = (
        Printer.select()
        .where((Printer.owner == owner) & Printer.driver.is_null(False))
        .count()
    )
    return templates.TemplateResponse(
        request=request,
        name="dashboard.html",
        context={
            "recent_printers": recent_printers,
            "recent_packages": recent_packages,
            "driver_count": Driver.select().count(),
            "printer_count": printer_count,
            "client_count": Client.select().where(Client.owner == owner).count(),
            "ready_count": ready_count,
        },
    )


@router.get("/drivers", response_class=HTMLResponse)
def drivers_page(request: Request):
    from imptune.db.models import Driver

    drivers = list(Driver.select().order_by(Driver.uploaded_at.desc()))
    return templates.TemplateResponse(
        request=request,
        name="drivers.html",
        context={
            "driver_data": _driver_data(drivers),
            "usage": printer_counts_by_driver(request.state.owner),
        },
    )


@router.get("/printers", response_class=HTMLResponse)
def printers_page(request: Request):
    from imptune.db.models import Client, Driver, Printer

    owner = request.state.owner
    query = (
        Printer.select(Printer, Client)
        .join(Client, JOIN.LEFT_OUTER)
        .where(Printer.owner == owner)
        .order_by(Client.name, Printer.name)
    )
    grouped = group_printers_by_client(query)

    clients = list(Client.select().where(Client.owner == owner).order_by(Client.name))
    all_drivers = list(Driver.select().order_by(Driver.uploaded_at.desc()))

    printer_count = sum(len(v) for v in grouped.values())
    ready_count = sum(1 for v in grouped.values() for p in v if p.driver_id)

    return templates.TemplateResponse(
        request=request,
        name="printers.html",
        context={
            "grouped": grouped,
            "clients": clients,
            "driver_data": _driver_data(all_drivers),
            "printer_count": printer_count,
            "ready_count": ready_count,
        },
    )


@router.get("/printers/new", response_class=HTMLResponse)
def printers_new_page(request: Request):
    from imptune.db.models import Client, Driver

    clients = list(Client.select().where(Client.owner == request.state.owner).order_by(Client.name))
    all_drivers = list(Driver.select().order_by(Driver.uploaded_at.desc()))
    return templates.TemplateResponse(
        request=request,
        name="printers_new.html",
        context={"clients": clients, "driver_data": _driver_data(all_drivers)},
    )


@router.get("/printers/{printer_id}", response_class=HTMLResponse)
def printer_detail(request: Request, printer_id: int):
    from imptune.db.models import Client, Driver, Icon, Printer

    printer = (
        Printer.select(Printer, Client, Driver)
        .join(Client, JOIN.LEFT_OUTER)
        .switch(Printer)
        .join(Driver, JOIN.LEFT_OUTER)
        .where((Printer.id == printer_id) & (Printer.owner == request.state.owner))
        .first()
    )
    if printer is None:
        return HTMLResponse(
            content="<h1>404 Not Found</h1><p>Printer not found.</p>",
            status_code=404,
        )

    driver_names: list[str] = []
    if printer.driver_id and printer.driver.driver_desc:
        driver_names = json.loads(printer.driver.driver_desc)

    has_driver = printer.driver_id is not None and bool(driver_names)
    icon = Icon.get_or_none(Icon.printer == printer_id)

    install_cmd = "powershell.exe -ExecutionPolicy Bypass -File install.ps1"
    uninstall_cmd = "powershell.exe -ExecutionPolicy Bypass -File uninstall.ps1"

    return templates.TemplateResponse(
        request=request,
        name="printer_detail.html",
        context={
            "printer": printer,
            "driver_names": driver_names,
            "has_driver": has_driver,
            "has_icon": icon is not None,
            "install_cmd": install_cmd,
            "uninstall_cmd": uninstall_cmd,
        },
    )


@router.get("/clients", response_class=HTMLResponse)
def clients_page(request: Request):
    from imptune.db.models import Client

    owner = request.state.owner
    clients = list(Client.select().where(Client.owner == owner).order_by(Client.name))
    return templates.TemplateResponse(
        request=request,
        name="clients.html",
        context={"clients": clients, "counts": printer_counts_by_client(owner)},
    )


@router.get("/clients/{client_id}", response_class=HTMLResponse)
def client_detail(request: Request, client_id: int):
    from imptune.db.models import Client, Driver, Printer

    owner = request.state.owner
    client = Client.get_or_none((Client.id == client_id) & (Client.owner == owner))
    if client is None:
        return HTMLResponse(
            content="<h1>404 Not Found</h1><p>Client not found.</p>",
            status_code=404,
        )

    query = (
        Printer.select(Printer, Client)
        .join(Client, JOIN.LEFT_OUTER)
        .where((Printer.client == client_id) & (Printer.owner == owner))
        .order_by(Printer.name)
    )
    printers = list(query)
    grouped = {client.name: printers}

    clients = list(Client.select().where(Client.owner == owner).order_by(Client.name))
    all_drivers = list(Driver.select().order_by(Driver.uploaded_at.desc()))

    return templates.TemplateResponse(
        request=request,
        name="client_detail.html",
        context={
            "client": client,
            "grouped": grouped,
            "clients": clients,
            "driver_data": _driver_data(all_drivers),
            "printer_count": len(printers),
            "ready_count": sum(1 for p in printers if p.driver_id),
        },
    )


@router.get("/packages", response_class=HTMLResponse)
def packages_page(request: Request):
    from imptune.db.models import Client, Driver, Icon, Printer

    owner = request.state.owner
    printers = list(
        Printer.select(Printer, Client, Driver)
        .join(Client, JOIN.LEFT_OUTER)
        .switch(Printer)
        .join(Driver, JOIN.LEFT_OUTER)
        .where((Printer.owner == owner) & Printer.driver.is_null(False))
        .order_by(Printer.name)
    )
    pending_count = (
        Printer.select()
        .where((Printer.owner == owner) & Printer.driver.is_null(True))
        .count()
    )
    icon_printer_ids = {
        row.printer_id
        for row in Icon.select(Icon.printer).where(
            Icon.printer.in_([p.id for p in printers] or [0])
        )
    }
    return templates.TemplateResponse(
        request=request,
        name="packages.html",
        context={
            "printers": printers,
            "pending_count": pending_count,
            "icon_printer_ids": icon_printer_ids,
        },
    )