Model/scripts/build_pashub_accuracy_fixtures.py

"""Build the pashub SAP-accuracy test fixtures from S3 + hubspot_deal_data.

One-time / rebuild-able provenance script for
`backend/documents_parser/tests/test_pashub_sap_accuracy.py`.

For the Guinness GMCA project it pulls every deal that has a `pre_sap`
rating and a pashub `rd_sap_site_note` PDF, downloads the PDF from S3,
strips the embedded survey photos (the extractor only ever reads the PDF
*text layer* via `page.get_text()`, so removing raster images leaves a
byte-identical token stream while shrinking each file ~25x), keeps only
pashub-format PDFs (Elmhurst-format notes route to a different extractor),
and writes them plus a `manifest.json` into the fixtures directory.

Requires DB + AWS credentials (reads `backend/.env`); NOT run in CI. The
committed fixtures + manifest are what the test consumes.

    python scripts/build_pashub_accuracy_fixtures.py [--limit N] [--force]
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Optional

import boto3
import psycopg2
import pymupdf
from dotenv import dotenv_values

REPO_ROOT = Path(__file__).resolve().parents[1]
ENV_PATH = REPO_ROOT / "backend" / ".env"
_FIXTURES_ROOT = (
    REPO_ROOT / "backend" / "documents_parser" / "tests" / "fixtures"
)
# Default cohort (Guinness GMCA); override per-portfolio with --project-code / --out-dir.
DEFAULT_FIXTURES_SUBDIR = "pashub_accuracy"
DEFAULT_PROJECT_CODE = (
    "[The Guinness Partnership GMCA Bid - WH:SF - 052026 - 1219680020683]"
)
ELMHURST_MARKER = "Elmhurst Energy Systems"

# One row per deal: the newest pashub RdSAP site-note that has a pre_sap.
_DEAL_QUERY = """
    select distinct on (h.deal_id)
        h.deal_id, h.uprn, h.pre_sap, u.s3_file_bucket, u.s3_file_key
    from hubspot_deal_data h
    join uploaded_files u on u.hubspot_deal_id = h.deal_id
    where h.project_code = %s
      and h.pre_sap is not null and h.pre_sap <> ''
      and u.file_type = 'rd_sap_site_note'
    order by h.deal_id, u.s3_upload_timestamp desc
"""


@dataclass(frozen=True)
class ManifestEntry:
    deal_id: str
    uprn: Optional[int]
    pre_sap_raw: str
    pre_sap_score: int
    pdf: str  # filename relative to the fixtures dir


def _parse_pre_sap(raw: str) -> Optional[int]:
    """`'C73'` / `'73'` / `'SAP 73'` -> `73`; unparseable -> None."""
    match = re.search(r"(\d{1,3})", raw)
    if match is None:
        return None
    score = int(match.group(1))
    return score if 1 <= score <= 100 else None


def _strip_images(pdf_bytes: bytes) -> bytes:
    """Remove embedded raster images; the text layer is untouched."""
    doc: Any = pymupdf.open(stream=pdf_bytes, filetype="pdf")
    try:
        for page in doc:
            for image in page.get_images(full=True):
                try:
                    page.delete_image(image[0])
                except Exception:  # noqa: BLE001 - best-effort per image
                    pass
        out: bytes = doc.tobytes(garbage=4, deflate=True, clean=True)
        return out
    finally:
        doc.close()


def _is_elmhurst(pdf_bytes: bytes) -> bool:
    doc: Any = pymupdf.open(stream=pdf_bytes, filetype="pdf")
    try:
        pages: list[str] = [page.get_text() for page in doc]
        return ELMHURST_MARKER in "\n".join(pages)
    finally:
        doc.close()


def _fetch_deals(limit: Optional[int], project_code: str) -> list[tuple[Any, ...]]:
    env = dotenv_values(ENV_PATH)
    conn: Any = psycopg2.connect(
        host=env["DB_HOST"],
        port=env.get("DB_PORT", "5432"),
        dbname=env["DB_NAME"],
        user=env["DB_USERNAME"],
        password=env["DB_PASSWORD"],
        connect_timeout=10,
    )
    try:
        cur: Any = conn.cursor()
        cur.execute(_DEAL_QUERY, (project_code,))
        rows: list[tuple[Any, ...]] = cur.fetchall()
    finally:
        conn.close()
    return rows[:limit] if limit is not None else rows


def build(
    limit: Optional[int], force: bool, project_code: str, fixtures_dir: Path
) -> None:
    manifest_path = fixtures_dir / "manifest.json"
    fixtures_dir.mkdir(parents=True, exist_ok=True)
    rows = _fetch_deals(limit, project_code)
    print(f"{len(rows)} candidate deals (pre_sap + rd_sap_site_note)")

    s3: Any = boto3.client("s3")  # pyright: ignore[reportUnknownMemberType]
    entries: list[ManifestEntry] = []
    skipped: dict[str, int] = {"no_pre_sap": 0, "elmhurst": 0, "s3_error": 0}

    for deal_id, uprn, pre_sap_raw, bucket, key in rows:
        score = _parse_pre_sap(pre_sap_raw)
        if score is None:
            skipped["no_pre_sap"] += 1
            continue

        pdf_name = f"{deal_id}.pdf"
        pdf_path = fixtures_dir / pdf_name

        if pdf_path.exists() and not force:
            entries.append(
                ManifestEntry(deal_id, uprn, pre_sap_raw, score, pdf_name)
            )
            continue

        try:
            raw: bytes = s3.get_object(Bucket=bucket, Key=key)["Body"].read()
        except Exception as exc:  # noqa: BLE001
            print(f"  {deal_id}: S3 error {exc}", file=sys.stderr)
            skipped["s3_error"] += 1
            continue

        if _is_elmhurst(raw):
            skipped["elmhurst"] += 1
            continue

        stripped = _strip_images(raw)
        pdf_path.write_bytes(stripped)
        entries.append(
            ManifestEntry(deal_id, uprn, pre_sap_raw, score, pdf_name)
        )
        print(
            f"  {deal_id}: pre_sap={pre_sap_raw}->{score} "
            f"{len(raw) / 1e6:.1f}MB -> {len(stripped) / 1e6:.2f}MB"
        )

    entries.sort(key=lambda e: e.deal_id)
    manifest_path.write_text(
        json.dumps(
            {
                "project_code": project_code,
                "source": "hubspot_deal_data + uploaded_files (rd_sap_site_note) + S3",
                "note": (
                    "pre_sap is pashub's own RdSAP rating; PDFs are image-"
                    "stripped (text layer preserved). Rebuild with "
                    "scripts/build_pashub_accuracy_fixtures.py"
                ),
                "fixtures": [asdict(e) for e in entries],
            },
            indent=2,
        )
        + "\n"
    )
    print(
        f"\nwrote {len(entries)} fixtures + manifest to {fixtures_dir}\n"
        f"skipped: {skipped}"
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--limit", type=int, default=None)
    parser.add_argument("--force", action="store_true")
    parser.add_argument(
        "--project-code",
        default=DEFAULT_PROJECT_CODE,
        help="hubspot_deal_data.project_code to select (default: Guinness GMCA)",
    )
    parser.add_argument(
        "--out-dir",
        default=DEFAULT_FIXTURES_SUBDIR,
        help=(
            "fixtures subdirectory name under tests/fixtures/ (default: "
            f"{DEFAULT_FIXTURES_SUBDIR}); use a per-portfolio name to avoid "
            "clobbering another cohort's manifest"
        ),
    )
    args = parser.parse_args()
    build(
        limit=args.limit,
        force=args.force,
        project_code=args.project_code,
        fixtures_dir=_FIXTURES_ROOT / args.out_dir,
    )


if __name__ == "__main__":
    main()