PacketStation/tests/server/closehelpers.py

"""Helpers for the phase 7 tests: a packet at complete on either route, the feed state written by SQL, a rich
packet for the data book (two drawing sheets, a void sheet, a heat table, two welds, MTRs, a short with a receipt,
a packet document and a Q.C. folder document), and PDF text helpers.

Built on shophelpers and qchelpers (imported, never edited). Nothing here reads VISUAL or M-Files.
"""

from __future__ import annotations

import io
import secrets
from collections.abc import Iterator
from typing import Any

from fakedir import FakeDirectory
from pkthelpers import step
from pshelpers import app_conn
from psycopg.types.json import Jsonb
from qchelpers import accept, add_packet_doc, at_qc, finish_inspection, make_qc_doc, upload_doc
from shophelpers import (
    SIGN,
    Shop,
    draft,
    finish_stage,
    im,
    line_of,
    make_unit,
    ok,
    open_shop,
    release,
    start,
    submit,
    to_station,
    to_warehouse,
)

from packetstation.settings import Settings

PDF = "application/pdf"
HEATS = {1: "HT-10", 2: "HT-20", 3: "HT-30"}
RECEIPT_HEAT = "HT-31"


def open_close_shop(settings: Settings, fake_dir: FakeDirectory) -> Iterator[Shop]:
    yield from open_shop(settings, fake_dir)


# ---------------------------------------------------------------- PDFs and files


def pdf_reader(data: bytes) -> Any:
    from pypdf import PdfReader

    return PdfReader(io.BytesIO(data))


def pdf_pages(data: bytes) -> int:
    return len(pdf_reader(data).pages)


def page_texts(data: bytes) -> list[str]:
    return [p.extract_text() or "" for p in pdf_reader(data).pages]


def pdf_text(data: bytes) -> str:
    """The text of every page, joined with newlines."""
    return "\n".join(page_texts(data))


def flat(text: str) -> str:
    """Whitespace folded to single spaces (text extraction breaks lines where the layout wraps)."""
    return " ".join(text.split())


def tiny_png(color: tuple[int, int, int] | None = None) -> bytes:
    from PIL import Image

    buf = io.BytesIO()
    c = color or (secrets.randbelow(256), secrets.randbelow(256), secrets.randbelow(256))
    Image.new("RGB", (4, 4), c).save(buf, "PNG")
    return buf.getvalue()


def is_pdf(data: bytes) -> bool:
    return data[:5] == b"%PDF-"


# ---------------------------------------------------------------- SQL


def q1(settings: Settings, sql: str, params: Any = ()) -> Any:
    with app_conn(settings) as conn:
        return conn.execute(sql, params).fetchone()


def qall(settings: Settings, sql: str, params: Any = ()) -> list[Any]:
    with app_conn(settings) as conn:
        return list(conn.execute(sql, params).fetchall())


def set_feed_legs(
    settings: Settings,
    packet_id: int,
    legs: list[dict[str, Any]] | None,
    *,
    closed: bool = False,
    observed: bool = True,
) -> None:
    """Simulate the feed: the packet becomes a VISUAL packet (source_kind visual, a visual_key), packet_visual_state
    holds the legs as the feed writes them ({key: "base/lot/split/sub", seq, st}), and when `closed` the feed's
    close is set: visual_closed_at, and vclosed_at with no vclosed_by (C5: the feed is the writer of vclosed on a
    VISUAL packet). `legs` None writes no state row (the feed has not looked yet)."""
    with app_conn(settings) as conn:
        conn.execute(
            "UPDATE packets SET source_kind = 'visual', visual_key = coalesce(visual_key, %s),"
            " visual_closed_at = CASE WHEN %s THEN now() END, vclosed_at = CASE WHEN %s THEN now() END,"
            " vclosed_by = NULL WHERE id = %s",
            (f"test-{packet_id}-{secrets.token_hex(3)}", closed, closed, packet_id),
        )
        if legs is None:
            conn.execute("DELETE FROM packet_visual_state WHERE packet_id = %s", (packet_id,))
            return
        conn.execute(
            "INSERT INTO packet_visual_state (packet_id, legs, observed_at)"
            " VALUES (%s, %s, CASE WHEN %s THEN now() END)"
            " ON CONFLICT (packet_id) DO UPDATE SET legs = EXCLUDED.legs, observed_at = EXCLUDED.observed_at",
            (packet_id, Jsonb(legs), observed),
        )


def leg(base: str, lot: str = "001", st: str = "R", seq: int = 10, split: str = "0", sub: str = "0") -> dict[str, Any]:
    """One leg as the feed stores it, with `status` beside `st` (the contract names the code `status`)."""
    return {"key": f"{base}/{lot}/{split}/{sub}", "seq": seq, "st": st, "status": st}


def void_page_by_sql(settings: Settings, packet_id: int, of_n: int = 1, co_no: int = 1) -> int:
    """A change order (approved, minor, no hold) and a void copy of sheet `of_n` as a new page row, the shape a
    change order approval leaves (void_co_id, void_at, void_by, void_new_n). Returns the void page's n."""
    with app_conn(settings) as conn:
        by = conn.execute("SELECT id FROM users WHERE sam_account = 'clerk.clara'").fetchone()[0]
        co = conn.execute(
            "INSERT INTO change_orders (packet_id, co_no, origin, review, minor, hold, reason, issued_by,"
            " approved_at, approved_by) VALUES (%s, %s, 'manual', 'approved', true, false, 'Nozzle moved', %s,"
            " now(), %s) RETURNING id",
            (packet_id, co_no, by, by),
        ).fetchone()[0]
        src = conn.execute(
            "SELECT asset_id, pn, dwg, rev FROM packet_pages WHERE packet_id = %s AND n = %s", (packet_id, of_n)
        ).fetchone()
        n = conn.execute("SELECT max(n) + 1 FROM packet_pages WHERE packet_id = %s", (packet_id,)).fetchone()[0]
        conn.execute(
            "INSERT INTO packet_pages (packet_id, n, type, asset_id, pn, dwg, rev, void_co_id, void_at, void_by,"
            " void_new_n) VALUES (%s, %s, 'drawing', %s, %s, %s, %s, %s, now(), %s, %s)",
            (packet_id, n, src[0], src[1], src[2], "A", co, by, of_n),
        )
    return int(n)


def remove_by_sql(settings: Settings, packet_id: int) -> None:
    with app_conn(settings) as conn:
        conn.execute(
            "UPDATE packets SET removed_at = now(), removed_reason = 'test' WHERE id = %s",
            (packet_id,),
        )


def delete_asset_file(settings: Settings, asset_id: str) -> None:
    """Remove a stored file from disk (the asset row stays): a part that cannot be read."""
    from packetstation.domain import assets

    row = q1(settings, "SELECT storage_key FROM assets WHERE id = %s", (asset_id,))
    assets.path_of(settings, str(row[0])).unlink()


# ---------------------------------------------------------------- packets


def at_complete(shop: Shop, unit_id: int | None = None, route: str = "databook", welds: int = 2) -> dict[str, Any]:
    """A shop packet taken through Q.C. and accepted with `route` (databook or shipping): status complete."""
    pkt = finish_inspection(shop, at_qc(shop, unit_id, welds=welds))
    out = ok(accept(shop, pkt, route))
    assert out["status"] == "complete" and out["routed_to"] == route, out["status"]
    return out


def vclose(shop: Shop, pkt: dict[str, Any], by: str = "clerk.clara") -> dict[str, Any]:
    return ok(shop(by).post(f"/api/packets/{pkt['id']}/vclosed", json={"closed": True}))


def file_packet(shop: Shop, pkt: dict[str, Any], by: str = "clerk.clara", **body: Any) -> Any:
    return step(shop(by), pkt, "file", sign=SIGN, **body)


def filed(shop: Shop, unit_id: int | None = None) -> dict[str, Any]:
    pkt = vclose(shop, at_complete(shop, unit_id))
    out = ok(file_packet(shop, pkt))
    assert out["status"] == "filed"
    return out


def attach_mtr(shop: Shop, pkt: dict[str, Any], item: int, data: bytes, name: str, mime: str, by: str) -> Any:
    """Upload an MTR (purpose mtr) and set it on a BOM line. Returns the packet answer."""
    a = upload_doc(shop(by), name, data, mime, purpose="mtr")
    cur = line_of(pkt, item)
    return ok(
        shop(by).patch(
            f"/api/packets/{pkt['id']}/bom/{item}", json={"mtr_asset_id": a["id"]}, headers=im(cur["row_version"])
        )
    )


def add_doc(shop: Shop, pkt: dict[str, Any], title: str = "RT report W1-W2", by: str = "qc.quinn", **kw: Any) -> Any:
    return ok(add_packet_doc(shop(by), pkt["id"], title=title, **kw))


def rich_complete(shop: Shop, *, itp_doc: bool = False, mtrs: bool = True) -> dict[str, Any]:
    """The data book packet, complete on the data book route:

    - drawing D-5001 rev B, two sheets, BOM items 1 to 3, two work orders;
    - in the warehouse: items 1 and 2 pulled with heats HT-10 and HT-20; item 3 short 2 of 4 (On order) with heat
      HT-30, then 1 pc received with heat HT-31 (one still open); an MTR PDF on item 1 and an MTR PNG on item 2
      (when `mtrs`); the heat table on sheet 1;
    - the shop: two welds, the checklist and the filler, submitted, sent to Q.C., inspected and accepted;
    - a packet document (RT report) and a Q.C. folder document on the unit (WPS / PQR, or ITP when `itp_doc`);
    - sheet 1 voided by change order 1 as a new sheet 3.

    Returns the fresh packet read. `pkt["_unit_id"]` and `pkt["_project_id"]` carry the ids."""
    project_id, unit_id = make_unit(shop.settings)
    pkt = to_warehouse(shop, draft(shop, unit_id, pages=2))
    wade = shop("wh.wade")
    pid = pkt["id"]
    for item in (1, 2):
        cur = line_of(pkt, item)
        pkt = ok(wade.patch(f"/api/packets/{pid}/bom/{item}", json={"pulled": True}, headers=im(cur["row_version"])))
        cur = line_of(pkt, item)
        pkt = ok(
            wade.patch(f"/api/packets/{pid}/bom/{item}", json={"heat": HEATS[item]}, headers=im(cur["row_version"]))
        )
    pkt = ok(
        wade.post(
            f"/api/packets/{pid}/bom/3/short",
            json={"qty": 2, "reason": "On order", "heat": HEATS[3], "expect_open": False},
        )
    )
    pkt = ok(wade.post(f"/api/packets/{pid}/bom/3/short/receipts", json={"qty": 1, "heat": RECEIPT_HEAT}))
    if mtrs:
        pkt = attach_mtr(shop, pkt, 1, upload_bytes_pdf(), "mtr-1.pdf", PDF, "wh.wade")
        pkt = attach_mtr(shop, pkt, 2, tiny_png(), "mtr-2.png", "image/png", "wh.wade")
    pkt = ok(wade.put(f"/api/packets/{pid}/heat-table/1", json={"x": 0.6, "y": 0.6, "s": 1.0}))
    pkt = release(shop, pkt)
    pkt = start(shop, to_station(shop, pkt))
    pkt = submit(shop, finish_stage(shop, pkt, welds=2))
    pkt = ok(step(shop("sup.wes"), pkt, "send_to_qc", book=True, spool=True, location="Bay 2", sign=SIGN))
    pkt = finish_inspection(shop, pkt)
    pkt = accept_with_open_short(shop, pkt, 3)
    add_doc(shop, pkt)
    r = make_qc_doc(shop, unit_id=unit_id, type="ITP" if itp_doc else "WPS / PQR", title="Unit folder doc")
    assert r.status_code in (200, 201), r.text
    void_page_by_sql(shop.settings, pid, of_n=1)
    out = ok(shop("clerk.clara").get(f"/api/packets/{pid}"))
    out["_unit_id"] = unit_id
    out["_project_id"] = project_id
    return out


def approve_receipts(shop: Shop, pkt: dict[str, Any], by: str = "sup.wes") -> dict[str, Any]:
    """The department supervisor approves every received short not yet approved (C7). Returns the packet read."""
    for sh in pkt.get("shorts") or []:
        for r in sh["receipts"]:
            if r["approval"] is None:
                url = f"/api/packets/{pkt['id']}/shorts/receipts/{r['id']}/approve"
                pkt = ok(shop(by).post(url, json={"sign": SIGN}))
    return pkt


def accept_with_open_short(shop: Shop, pkt: dict[str, Any], item: int) -> dict[str, Any]:
    """Accept to the data book while item `item` keeps its open short: a packet completed before change C7, which
    the data book and the trace CSV still show. The receipts are approved through the API; the open short is closed
    by SQL for the accept only (C7 refuses an accept with an open short) and opened again after it."""
    pkt = approve_receipts(shop, pkt)
    sid = line_of(pkt, item)["short"]["id"]
    with app_conn(shop.settings) as conn:
        conn.execute("UPDATE bom_shorts SET closed_at = now() WHERE id = %s", (sid,))
    pkt = ok(accept(shop, fresh_read(shop, pkt["id"]), "databook"))
    with app_conn(shop.settings) as conn:
        conn.execute("UPDATE bom_shorts SET closed_at = NULL WHERE id = %s", (sid,))
    return fresh_read(shop, pkt["id"])


def fresh_read(shop: Shop, pid: int) -> dict[str, Any]:
    return ok(shop("clerk.clara").get(f"/api/packets/{pid}"))


def upload_bytes_pdf() -> bytes:
    """A one-page PDF with its own bytes (the store dedupes equal bytes)."""
    from qchelpers import unique_pdf

    return unique_pdf()


# ---------------------------------------------------------------- the builders' inputs


def who(shop: Shop, sam: str = "clerk.clara") -> dict[str, Any]:
    """The WhoRef of a test user: {id, display_name, emp_no}."""
    row = q1(shop.settings, "SELECT id, display_name, emp_no FROM users WHERE sam_account = %s", (sam,))
    return {"id": str(row[0]), "display_name": row[1], "emp_no": row[2]}


def builder_conn(settings: Settings) -> Any:
    """A connection for calling a builder directly (autocommit, app role)."""
    return app_conn(settings)