PacketStation/server/packetstation/pdf/prints.py

"""PRINTED COPY PDFs: the packet print kinds and the ITP print, each page stamped with the old print stamp box.

Kinds: traveler, map-h (drawing with heat map), map-w (drawing with weld map), clean (the drawing pages as attached,
void pages stamped VOID), void (the weld-map raster of every void page), databook (the filed book when there is one,
else built now) and doc (one document of the packet: an MTR, a packet document, a Q.C. folder document or the filed
data book). The person prints the PDF from the viewer.
"""

from __future__ import annotations

import io
from collections.abc import Callable, Mapping
from datetime import datetime
from typing import Any

from pypdf import PdfReader, PdfWriter

from . import databook, maps, parts, stamps, traveler
from . import itp as itp_pdf
from .style import Built, NothingToPrint, PdfError, WhoRef, clean, utc_now

KINDS: dict[str, str] = {
    "traveler": "Job traveler",
    "map-h": "Drawing with heat map",
    "map-w": "Drawing with weld map",
    "clean": "Drawing (no markup)",
    "void": "Superseded drawings (VOID)",
    "databook": "Data book",
    "doc": "Document",
}
SHEET_KINDS = ("map-h", "map-w", "clean")


def filed_book_asset(conn: Any, packet_id: int) -> Any:
    """The data book stored on the packet's unbroken seal, or None."""
    r = conn.execute(
        "SELECT databook_asset_id FROM packet_seals WHERE packet_id = %s AND broken_at IS NULL"
        " AND databook_asset_id IS NOT NULL ORDER BY id DESC LIMIT 1",
        (int(packet_id),),
    ).fetchone()
    return None if r is None else r[0]


def doc_label(conn: Any, p: Mapping[str, Any], asset_id: Any) -> str:
    """The label of one of the packet's documents, or NothingToPrint when the asset is not one of them."""
    want = str(asset_id or "")
    if not want:
        raise NothingToPrint("No document was named.")
    for b in p.get("bom") or []:
        if str(b.get("mtr_asset_id") or "") == want:
            heat = str(b.get("heat") or "").strip()
            return f"MTR - item {b.get('item')} {heat}".rstrip()
    for d in p.get("docs") or []:
        if str(d.get("asset_id") or "") == want:
            return str(d.get("title") or d.get("kind_label") or d.get("file_name") or "Document")
    for f in p.get("qc_folder_docs") or []:
        if str(f.get("url") or "").rsplit("/", 1)[-1] == want:
            return f"{f.get('type')} - {f.get('title') or ''}".rstrip(" -")
    r = conn.execute(
        "SELECT 1 FROM packet_seals WHERE packet_id = %s AND databook_asset_id::text = %s", (int(p["id"]), want)
    ).fetchone()
    if r is not None:
        return "Data book"
    raise NothingToPrint("The document is not one of this packet's.")


def _sheets(p: Mapping[str, Any], kind: str, sheet: int | None) -> list[dict[str, Any]]:
    drawing = [pg for pg in p.get("pages") or [] if pg.get("type") == "drawing"]
    if sheet is not None:
        found = [pg for pg in drawing if int(pg["n"]) == int(sheet)]
        if not found:
            raise NothingToPrint(f"The packet has no drawing sheet {sheet}.")
        return found
    if kind == "clean":
        return drawing
    return [pg for pg in drawing if not pg.get("void")]


def _add_pdf(writer: PdfWriter, data: bytes) -> None:
    for pg in PdfReader(io.BytesIO(data)).pages:
        writer.add_page(pg)


def _rasters(
    conn: Any, settings: Any, p: Mapping[str, Any], pages: list[dict[str, Any]], fn: Callable[..., Any]
) -> PdfWriter:
    writer = PdfWriter()
    for pg in pages:
        try:
            jpeg, w, h = fn(conn, settings, p, pg)
        except Exception as e:
            raise PdfError(f"Drawing sheet {pg['n']} could not be read ({type(e).__name__}).") from e
        _add_pdf(writer, parts.raster_page_pdf(jpeg, w, h))
    return writer


def _clean(
    conn: Any, settings: Any, p: Mapping[str, Any], pages: list[dict[str, Any]]
) -> tuple[PdfWriter, dict[int, str]]:
    writer = PdfWriter()
    voids: dict[int, str] = {}
    readers: dict[str, PdfReader] = {}
    tz = str(getattr(settings, "site_tz", "") or "America/Denver")
    for pg in pages:
        try:
            key = str(pg.get("asset_id"))
            if key not in readers:
                _row, path = maps.asset_file(conn, settings, pg.get("asset_id"))
                readers[key] = parts.open_pdf(path)
            writer.add_page(readers[key].pages[int(pg.get("pn") or 1) - 1])
        except Exception as e:
            raise PdfError(f"Drawing sheet {pg['n']} could not be read ({type(e).__name__}).") from e
        if pg.get("void"):
            voids[len(writer.pages) - 1] = maps.void_text(pg, tz)
    return writer, voids


def _doc(conn: Any, settings: Any, asset_id: Any) -> PdfWriter:
    try:
        _row, path = maps.asset_file(conn, settings, asset_id)
        kind = parts.sniff(path)
        writer = PdfWriter()
        if kind == "pdf":
            rd = parts.open_pdf(path)
            for pg in rd.pages:
                writer.add_page(pg)
        elif kind in ("png", "jpeg"):
            _add_pdf(writer, parts.image_page_pdf(path))
        else:
            raise PdfError("The file is not a PDF, PNG or JPEG.")
        return writer
    except PdfError:
        raise
    except Exception as e:
        raise PdfError(f"The document could not be read ({type(e).__name__}).") from e


def _stamped(
    writer: PdfWriter,
    settings: Any,
    *,
    who: WhoRef,
    now: datetime,
    what: str,
    kind: str,
    packet: Mapping[str, Any] | None,
    packet_id: Any,
    voids: Mapping[int, str] | None = None,
    title: str,
    warnings: list[str] | None = None,
) -> Built:
    if len(writer.pages) == 0:
        raise NothingToPrint("Nothing to print.")
    warnings = [] if warnings is None else warnings
    text = stamps.print_stamp(who, now, what, packet)
    stamp, sig = stamps.signed_stamp(
        settings,
        text,
        packet_id=packet_id,
        kind=f"print:{kind}",
        seal_hash_hex=None,
        who=who,
        now=now,
        warnings=warnings,
    )
    voids = voids or {}

    def draw(c: Any, i: int, w: float, h: float) -> None:
        if i in voids:
            stamps.draw_void(c, w, h, voids[i])
        stamps.draw_box_stamp(c, w, stamp)

    stamps.overlay_writer(writer, draw)
    keywords = ", ".join(x for x in ("printed copy", kind, f"stamp {sig}" if sig else "") if x)
    meta = stamps.metadata(title, f"Printed copy - {what}", keywords, "Packet Station - printed copy", now=now)
    data = stamps.write(writer, meta)
    return Built(
        bytes=data,
        pages=len(PdfReader(io.BytesIO(data)).pages),
        warnings=warnings,
        stamp=stamp,
        stamp_sig=sig,
        what=what,
    )


def build_print(
    conn: Any,
    settings: Any,
    packet_id: int,
    kind: str,
    *,
    who: WhoRef,
    now: datetime | None = None,
    sheet: int | None = None,
    asset_id: Any = None,
) -> Built:
    """A PRINTED COPY of one packet print kind. NothingToPrint when the kind has nothing for this packet (no void
    pages, an unknown sheet, a document that is not the packet's); PdfError when it cannot be built."""
    if kind not in KINDS:
        raise NothingToPrint(f"Unknown print kind {kind}.")
    now = utc_now(now)
    p = parts.load_packet(conn, packet_id)
    project, lot = parts.tag_of(p)
    what = KINDS[kind]
    voids: dict[int, str] = {}
    warnings: list[str] = []
    if kind == "traveler":
        data, _n = traveler.pages_of(conn, settings, p)
        # Page numbers as on the exported traveler, then the print box.
        writer = PdfWriter(clone_from=PdfReader(io.BytesIO(traveler.finish(data, p, "", now, box=False))))
    elif kind in ("map-h", "map-w"):
        pages = _sheets(p, kind, sheet)
        fn = maps.heat_map_page if kind == "map-h" else maps.weld_map_page
        writer = _rasters(conn, settings, p, pages, fn)
    elif kind == "clean":
        writer, voids = _clean(conn, settings, p, _sheets(p, kind, sheet))
    elif kind == "void":
        pages = parts.void_pages(p)
        if not pages:
            raise NothingToPrint("The packet has no void pages.")
        writer = _rasters(conn, settings, p, pages, maps.weld_map_page)
    elif kind == "databook":
        stored = asset_id or (filed_book_asset(conn, packet_id) if p.get("status") == "filed" else None)
        if stored is not None:
            writer = _doc(conn, settings, stored)
        else:
            if p.get("status") not in ("complete", "filed"):
                raise NothingToPrint("The data book is built once the packet is back from Q.C.")
            book = databook.build(conn, settings, packet_id, who=who, now=now, mode="print")
            warnings += book.warnings
            writer = PdfWriter(clone_from=PdfReader(io.BytesIO(book.bytes)))
    else:
        what = doc_label(conn, p, asset_id)
        writer = _doc(conn, settings, asset_id)
    if sheet is not None and kind in SHEET_KINDS:
        what = f"{what} (sheet {sheet})"
    return _stamped(
        writer,
        settings,
        who=who,
        now=now,
        what=clean(what),
        kind=kind,
        packet=p,
        packet_id=p["id"],
        voids=voids,
        title=f"{project} Lot {lot} {clean(what)}",
        warnings=warnings,
    )


def print_itp(
    conn: Any, settings: Any, itp_row: Mapping[str, Any], *, who: WhoRef, now: datetime | None = None
) -> Built:
    """The ITP PDF stamped PRINTED COPY, with its page numbers."""
    now = utc_now(now)
    data, _n, r = itp_pdf.pages_of(conn, settings, itp_row)
    num = itp_pdf.finish(data, r, "", now, box=False)
    writer = PdfWriter(clone_from=PdfReader(io.BytesIO(num)))
    what = f"ITP {clean(r.get('no'))}"
    return _stamped(
        writer,
        settings,
        who=who,
        now=now,
        what=what,
        kind="itp",
        packet=None,
        packet_id=f"itp-{itp_row['id']}",
        title=what,
    )