autotrain/scripts/annotate_roster_invoices.py

#!/usr/bin/env python3
"""Annotate a roster xlsx with invoice status + PDF hyperlink by matching passenger names.

Scans one or more directories for invoice PDFs whose filenames already carry the
real passenger name (produced by add_invoice_passenger_name.py or the timestamped
android flow), builds a name -> pdf map, then writes two columns into the roster:
`开票状态` (已开票 / 未开票) and `发票PDF` (filename + file:// hyperlink).
"""

from __future__ import annotations

import argparse
import os
from pathlib import Path
from urllib.parse import quote

import openpyxl

try:
    from pypdf import PdfReader
except ImportError:  # pragma: no cover
    PdfReader = None


def rel_hyperlink(pdf: Path, base: Path) -> str:
    rel = os.path.relpath(pdf, base)
    return "/".join(quote(seg) for seg in rel.split(os.sep))


def is_child_ticket(pdf: Path) -> bool:
    if PdfReader is None:
        return False
    try:
        reader = PdfReader(str(pdf))
        text = "\n".join(page.extract_text() or "" for page in reader.pages)
    except Exception:
        return False
    return "孩" in text


def extract_name(fname: str) -> str | None:
    stem = fname[:-4] if fname.lower().endswith(".pdf") else fname
    parts = stem.split("_")
    # timestamped android flow: YYYYMMDD_HHMMSS_realname_masked
    if len(parts) >= 4 and parts[0].isdigit() and len(parts[0]) == 8:
        return parts[2]
    # add_invoice_passenger_name output: <digits>-_realname_电子发票...
    if len(parts) >= 2:
        return parts[1]
    return None


def norm(name: str) -> str:
    return str(name).replace(" ", "").replace("\u3000", "").strip()


def build_map(dirs: list[Path]) -> dict[str, Path]:
    mapping: dict[str, Path] = {}
    for d in dirs:
        for p in sorted(d.glob("*.pdf")):
            nm = extract_name(p.name)
            if nm:
                mapping.setdefault(norm(nm), p)
    return mapping


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("xlsx", type=Path, help="Roster xlsx to annotate (edited in place)")
    ap.add_argument("dirs", type=Path, nargs="+", help="Directories containing named invoice PDFs")
    ap.add_argument("--sheet", default=None, help="Sheet name (default: first sheet)")
    ap.add_argument("--name-col", type=int, default=2, help="1-based column of passenger name (default 2)")
    ap.add_argument("--status-header", default="开票状态")
    ap.add_argument("--link-header", default="发票PDF")
    ap.add_argument("--done-label", default="已开票")
    ap.add_argument("--missing-label", default="未开票")
    ap.add_argument("--child-suffix", default="-孩", help="Suffix appended to done-label when the invoice is a child ticket (contains 孩).")
    args = ap.parse_args()

    xlsx = args.xlsx.expanduser().resolve()
    if not xlsx.is_file():
        raise SystemExit(f"not a file: {xlsx}")
    dirs = [d.expanduser().resolve() for d in args.dirs]
    for d in dirs:
        if not d.is_dir():
            raise SystemExit(f"not a directory: {d}")

    name_to_pdf = build_map(dirs)
    print(f"mapped PDFs: {len(name_to_pdf)}")

    wb = openpyxl.load_workbook(xlsx)
    ws = wb[args.sheet] if args.sheet else wb.worksheets[0]

    def find_or_append_col(header: str) -> int:
        for c in range(1, ws.max_column + 1):
            if ws.cell(row=1, column=c).value == header:
                return c
        col = ws.max_column + 1
        ws.cell(row=1, column=col, value=header)
        return col

    status_col = find_or_append_col(args.status_header)
    link_col = find_or_append_col(args.link_header)

    xlsx_dir = xlsx.parent
    matched = 0
    roster: set[str] = set()
    for r in range(2, ws.max_row + 1):
        name = ws.cell(row=r, column=args.name_col).value
        if not name:
            continue
        key = norm(name)
        roster.add(key)
        cell = ws.cell(row=r, column=link_col)
        cell.hyperlink = None
        pdf = name_to_pdf.get(key)
        if pdf:
            label = args.done_label
            if is_child_ticket(pdf):
                label = f"{args.done_label}{args.child_suffix}"
            ws.cell(row=r, column=status_col, value=label)
            rel = rel_hyperlink(pdf, xlsx_dir)
            cell.value = pdf.name
            cell.hyperlink = rel
            cell.style = "Hyperlink"
            matched += 1
        else:
            ws.cell(row=r, column=status_col, value=args.missing_label)
            cell.value = None

    wb.save(xlsx)

    missing = sorted(roster - set(name_to_pdf))
    extra = sorted(set(name_to_pdf) - roster)
    print(f"saved: {xlsx}")
    print(f"matched rows: {matched}")
    print(f"roster without pdf ({len(missing)}): {missing}")
    print(f"pdf without roster ({len(extra)}): {extra}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())