homelab-codex-ws/jobs/documents-ingest/src/documents_ingest/extractor.py

"""PDF attachment extractor — kb-postgres mail archive -> Paperless consume/.

Module 5 ("documents-ingest"), Phase 1 (kb/phases/kb-m5-documents-ingest.md,
section "Domkniecie dlugu z maili"): pull a sample of PDF attachments out of the
Gmail .eml archive and drop them into Paperless' consume/ dir so Paperless does
the OCR + correspondent-detection. This is NOT the Paperless/Nextcloud envelope
adapter (that is a later phase of module 5) — this job only produces files for
Paperless to ingest; it does not write to the `envelope` table itself.

Runs on PIHA (the only host with the mail archive, the Paperless consume dir,
and kb-postgres all reachable without extra hops):

Install (from repo root):
    pip install -e jobs/documents-ingest/

Usage:
    # Dry run (default) — preview what would be extracted, no writes:
    documents-ingest --dsn postgresql://kb:<pw>@localhost:5433/kb

    # Real run — write files into consume/ and update the registry:
    documents-ingest --dsn postgresql://kb:<pw>@localhost:5433/kb --apply

    # Smaller/larger sample, different window:
    documents-ingest --dsn ... --limit 50 --since-days 180 --min-size 100000

DSN can also come from the KB_DSN env var instead of --dsn.
"""
from __future__ import annotations

import argparse
import asyncio
import email
import email.policy
import hashlib
import json
import os
import re
import sys
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Optional

import asyncpg
import structlog

_log = structlog.get_logger(__name__)

DEFAULT_ARCHIVE_ROOT = Path("/home/oskar/kb/mail/archive")
DEFAULT_CONSUME_DIR = Path("/opt/homelab/data/paperless/consume")
DEFAULT_REGISTRY = Path("/opt/homelab/data/documents-ingest/registry.json")
DEFAULT_LIMIT = 150
DEFAULT_MIN_SIZE = 50_000
DEFAULT_SINCE_DAYS = 365

CONSUME_UID = 1000
CONSUME_GID = 1000

_UNSAFE_NAME = re.compile(r"[^A-Za-z0-9._-]+")
_MAX_NAME_LEN = 150


@dataclass
class Candidate:
    """One PDF attachment manifest entry that passed the size/type filter."""

    envelope_id: str
    ts: datetime
    filename: str
    size: int
    sha256: str


def sanitize_filename(filename: str, fallback: str = "attachment.pdf") -> str:
    """Reduce a manifest filename to a safe, bounded basename ending in .pdf."""
    name = Path(filename or "").name  # drop any path components
    name = _UNSAFE_NAME.sub("_", name).strip("._")
    if not name:
        name = fallback
    if not name.lower().endswith(".pdf"):
        name = f"{name}.pdf"
    if len(name) > _MAX_NAME_LEN:
        name = f"{name[: _MAX_NAME_LEN - 4]}.pdf"
    return name


def build_consume_name(ts: datetime, filename: str, sha256: str, used: set[str]) -> str:
    """Build a unique, readable consume/ filename: <date>_<sanitized>[_<hash8>].pdf.

    `used` tracks names already claimed in this run (and pre-seeded with names
    already present in consume/) so repeated filenames don't collide.
    """
    safe = sanitize_filename(filename)
    stem = safe[: -len(".pdf")]
    base = f"{ts:%Y-%m-%d}_{stem}.pdf"
    if base not in used:
        used.add(base)
        return base
    disambiguated = f"{ts:%Y-%m-%d}_{stem}_{sha256[:8]}.pdf"
    used.add(disambiguated)
    return disambiguated


def pdf_candidates_from_entities(
    envelope_id: str, ts: datetime, entities: list, min_size: int
) -> list[Candidate]:
    """Filter an envelope's attachment manifest down to qualifying PDF entries."""
    candidates = []
    for entity in entities:
        if entity.get("type") != "attachment":
            continue
        if entity.get("content_type") != "application/pdf":
            continue
        size = entity.get("size") or 0
        if size <= min_size:
            continue
        sha256 = entity.get("sha256")
        if not sha256:
            continue
        candidates.append(
            Candidate(
                envelope_id=envelope_id,
                ts=ts,
                filename=entity.get("filename") or "",
                size=size,
                sha256=sha256,
            )
        )
    return candidates


def find_pdf_parts(raw: bytes) -> list[tuple[Optional[str], bytes]]:
    """Return (filename, payload) for every application/pdf MIME part in the message.

    filename is whatever `email.policy.default` decodes it to (RFC 2047 aware) —
    it will NOT byte-for-byte equal a raw manifest filename that still carries
    encoded-word/folding artifacts (the manifest was built by a different parser
    at import time). sha256 is the authoritative match; see process_candidate().
    """
    msg = email.message_from_bytes(raw, policy=email.policy.default)
    parts = []
    for part in msg.walk():
        if part.get_content_maintype() == "multipart":
            continue
        if part.get_content_type() != "application/pdf":
            continue
        payload = part.get_payload(decode=True)
        if payload is not None:
            parts.append((part.get_filename(), payload))
    return parts


def load_registry(path: Path) -> dict[str, dict]:
    """Load the sha256 -> ingested-record idempotency registry. Missing file = empty."""
    if not path.is_file():
        return {}
    return json.loads(path.read_text())


def save_registry(path: Path, registry: dict[str, dict]) -> None:
    """Persist the registry atomically (write-tmp + rename)."""
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(registry, indent=2, sort_keys=True))
    tmp.replace(path)


def _decode_jsonb(value: object) -> object:
    """asyncpg may return jsonb as a str or an already-decoded object."""
    if value is None:
        return None
    if isinstance(value, str):
        return json.loads(value)
    return value


def _existing_consume_names(consume_dir: Path) -> set[str]:
    try:
        return {p.name for p in consume_dir.iterdir() if p.is_file()}
    except OSError:
        return set()


async def fetch_candidate_envelopes(
    conn: asyncpg.Connection, since: datetime, min_size: int, limit: int
) -> list[asyncpg.Record]:
    """Envelopes carrying at least one qualifying PDF attachment, newest first."""
    return await conn.fetch(
        """
        SELECT id, raw_ref, ts, entities
        FROM envelope
        WHERE source = 'gmail'
          AND ts > $1
          AND EXISTS (
              SELECT 1 FROM jsonb_array_elements(entities) AS att
              WHERE att->>'content_type' = 'application/pdf'
                AND (att->>'size')::numeric > $2
          )
        ORDER BY ts DESC
        LIMIT $3
        """,
        since,
        min_size,
        limit,
    )


def process_candidate(
    raw: bytes,
    candidate: Candidate,
    registry: dict[str, dict],
    used_names: set[str],
    apply: bool,
) -> tuple[str, Optional[dict]]:
    """Resolve one candidate against the .eml bytes and the registry.

    Returns (status, example) where status is one of:
    "duplicate", "sha_mismatch", "parse_error", "extracted".
    `example` is a dict describing the (would-be) consume/ file, or None.
    Does not touch the filesystem — writing is the caller's job (see run()).
    """
    if candidate.sha256 in registry:
        _log.info("skip.duplicate", envelope_id=candidate.envelope_id, filename=candidate.filename)
        return "duplicate", None

    try:
        pdf_parts = find_pdf_parts(raw)
    except Exception:
        _log.warning(
            "skip.parse_error", envelope_id=candidate.envelope_id, filename=candidate.filename, exc_info=True
        )
        return "parse_error", None

    if not pdf_parts:
        _log.warning(
            "skip.parse_error",
            envelope_id=candidate.envelope_id,
            filename=candidate.filename,
            reason="no_pdf_parts_in_mime",
        )
        return "parse_error", None

    # sha256 is the proof of identity (manifest filenames may carry raw RFC 2047
    # encoded-word/folding artifacts from the original import parser, so they
    # won't always byte-match what email.policy.default decodes today — the hash
    # doesn't have that problem). Filename match is used only to distinguish
    # "found the named part but its bytes changed" from "not present at all".
    matched = next(
        (
            (fname, payload)
            for fname, payload in pdf_parts
            if hashlib.sha256(payload).hexdigest() == candidate.sha256
        ),
        None,
    )
    if matched is None:
        if any(fname == candidate.filename for fname, _ in pdf_parts):
            _log.warning(
                "skip.sha_mismatch",
                envelope_id=candidate.envelope_id,
                filename=candidate.filename,
                expected_sha256=candidate.sha256,
            )
            return "sha_mismatch", None
        _log.warning(
            "skip.parse_error",
            envelope_id=candidate.envelope_id,
            filename=candidate.filename,
            reason="attachment_not_found_in_mime",
        )
        return "parse_error", None

    # Prefer the decoded filename from the MIME part itself for the human-facing
    # consume/ name — the manifest filename may still carry encoded-word artifacts.
    matched_filename, matched_payload = matched
    display_filename = matched_filename or candidate.filename
    target_name = build_consume_name(candidate.ts, display_filename, candidate.sha256, used_names)
    example = {
        "envelope_id": candidate.envelope_id,
        "filename": display_filename,
        "target_name": target_name,
        "size": candidate.size,
        "sha256": candidate.sha256,
        "payload": matched_payload if apply else None,
    }
    return "extracted", example


async def run(
    dsn: str,
    archive_root: Path = DEFAULT_ARCHIVE_ROOT,
    consume_dir: Path = DEFAULT_CONSUME_DIR,
    registry_path: Path = DEFAULT_REGISTRY,
    limit: int = DEFAULT_LIMIT,
    min_size: int = DEFAULT_MIN_SIZE,
    since_days: int = DEFAULT_SINCE_DAYS,
    apply: bool = False,
) -> tuple[dict[str, int], list[dict]]:
    """Scan candidate envelopes and (if apply) extract PDFs into consume/.

    Returns (stats, examples). stats always reflects what was found/matched;
    in dry-run mode "extracted" counts what *would* be written.
    """
    conn = await asyncpg.connect(dsn)
    try:
        since = datetime.now(timezone.utc) - timedelta(days=since_days)
        rows = await fetch_candidate_envelopes(conn, since, min_size, limit)
    finally:
        await conn.close()

    registry = load_registry(registry_path)
    used_names = _existing_consume_names(consume_dir)

    stats = {
        "envelopes_scanned": 0,
        "pdf_candidates": 0,
        "extracted": 0,
        "skipped_duplicate": 0,
        "skipped_sha_mismatch": 0,
        "skipped_parse_error": 0,
        "errors": 0,
    }
    examples: list[dict] = []

    for row in rows:
        stats["envelopes_scanned"] += 1
        entities = _decode_jsonb(row["entities"]) or []
        ts = row["ts"]
        if ts.tzinfo is None:
            ts = ts.replace(tzinfo=timezone.utc)

        candidates = pdf_candidates_from_entities(row["id"], ts, entities, min_size)
        if not candidates:
            continue

        eml_path = archive_root / row["raw_ref"]
        try:
            raw = eml_path.read_bytes()
        except OSError:
            _log.error("eml_read_failed", envelope_id=row["id"], path=str(eml_path))
            stats["errors"] += len(candidates)
            continue

        for candidate in candidates:
            stats["pdf_candidates"] += 1
            status, example = process_candidate(raw, candidate, registry, used_names, apply)
            stats[
                {
                    "duplicate": "skipped_duplicate",
                    "sha_mismatch": "skipped_sha_mismatch",
                    "parse_error": "skipped_parse_error",
                    "extracted": "extracted",
                }[status]
            ] += 1

            if status != "extracted" or example is None:
                continue
            examples.append({k: v for k, v in example.items() if k != "payload"})

            if not apply:
                continue

            target_path = consume_dir / example["target_name"]
            target_path.write_bytes(example["payload"])
            try:
                os.chown(target_path, CONSUME_UID, CONSUME_GID)
            except PermissionError:
                _log.warning("chown_failed", path=str(target_path))

            registry[candidate.sha256] = {
                "envelope_id": candidate.envelope_id,
                "filename": example["filename"],
                "consume_name": example["target_name"],
                "size": candidate.size,
                "ingested_at": datetime.now(timezone.utc).isoformat(),
            }
            _log.info("written", envelope_id=candidate.envelope_id, target=example["target_name"])

    if apply:
        save_registry(registry_path, registry)

    _log.info("run_complete", apply=apply, **stats)
    return stats, examples


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Extract PDF attachments from the mail archive into Paperless consume/ "
        "(module 5, phase 1 — sample, not bulk)."
    )
    parser.add_argument("--dsn", default=os.environ.get("KB_DSN"),
                        help="asyncpg DSN for kb-postgres (or set KB_DSN env var)")
    parser.add_argument("--archive-root", type=Path, default=DEFAULT_ARCHIVE_ROOT,
                        help=f"Mail .eml archive root (default: {DEFAULT_ARCHIVE_ROOT})")
    parser.add_argument("--consume-dir", type=Path, default=DEFAULT_CONSUME_DIR,
                        help=f"Paperless consume/ directory (default: {DEFAULT_CONSUME_DIR})")
    parser.add_argument("--registry", type=Path, default=DEFAULT_REGISTRY,
                        help=f"Idempotency registry JSON (default: {DEFAULT_REGISTRY})")
    parser.add_argument("--limit", type=int, default=DEFAULT_LIMIT,
                        help=f"Max envelopes to consider (default: {DEFAULT_LIMIT} — a sample, not bulk)")
    parser.add_argument("--min-size", type=int, default=DEFAULT_MIN_SIZE,
                        help=f"Minimum PDF attachment size in bytes (default: {DEFAULT_MIN_SIZE})")
    parser.add_argument("--since-days", type=int, default=DEFAULT_SINCE_DAYS,
                        help=f"Only envelopes newer than N days (default: {DEFAULT_SINCE_DAYS})")
    parser.add_argument("--apply", action="store_true",
                        help="Actually write files and update the registry. Default is dry-run.")
    args = parser.parse_args()

    if not args.dsn:
        _log.error("missing_dsn", hint="pass --dsn or set KB_DSN")
        sys.exit(1)
    if not args.archive_root.is_dir():
        _log.error("archive_root_not_found", path=str(args.archive_root))
        sys.exit(1)
    if args.apply and not args.consume_dir.is_dir():
        _log.error("consume_dir_not_found", path=str(args.consume_dir))
        sys.exit(1)

    stats, examples = asyncio.run(
        run(
            dsn=args.dsn,
            archive_root=args.archive_root,
            consume_dir=args.consume_dir,
            registry_path=args.registry,
            limit=args.limit,
            min_size=args.min_size,
            since_days=args.since_days,
            apply=args.apply,
        )
    )

    mode = "APPLY" if args.apply else "DRY-RUN"
    _log.info("summary", mode=mode, **stats)
    for ex in examples[:20]:
        _log.info("example", target_name=ex["target_name"], size=ex["size"], envelope_id=ex["envelope_id"])
    if len(examples) > 20:
        _log.info("examples_truncated", shown=20, total=len(examples))

    sys.exit(1 if stats["errors"] > 0 else 0)


if __name__ == "__main__":
    main()