ImpTune/imptune/api/pages.py

import json
from collections import defaultdict

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

from imptune.templating import templates

router = APIRouter()

# 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 build_driver_data(drivers) -> list[dict]:
    """Attach parsed driver names + icon presence to each Driver row.

    Shared by every template that lists drivers (library table, printer form
    select, edit modal) so a rename or a new icon shows up in all of them.
    """
    from imptune.services.icons import driver_icon_ids, driver_search_text
    from imptune.services.websearch import image_query

    drivers = list(drivers)
    with_icons = driver_icon_ids(drivers)
    data = []
    for d in drivers:
        names = json.loads(d.driver_desc) if d.driver_desc else []
        search_text = driver_search_text(d, names)
        data.append(
            {
                "driver": d,
                "names": names,
                "has_icon": d.id in with_icons,
                "search_text": search_text,
                "image_query": image_query(search_text),
            }
        )
    return data


@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

    from imptune.services.websearch import BRAND_SUGGESTIONS

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


@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": build_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": build_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
    from imptune.services.websearch import image_query

    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,
            "image_query": image_query(printer.name),
            "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": build_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,
        },
    )