"""Gmail header backfill — one-shot job, module 5 phase 2 (kb/phases/kb-m5-faza2.md, §5).

Backfills `{"type": "headers", ...}` (§4.1) onto the 225 030 existing `source='gmail'`
envelope rows in kb-postgres, which today carry only an attachment manifest. This is an
UPDATE against production data (not an INSERT like gmail-bulk-import) — a deliberately
separate job from gmail-bulk-import, see plan §2 decision 7.

Runs on PIHA (needs local access to the .eml archive):

Install (from repo root):
    pip install -e jobs/gmail-header-backfill/

Usage:
    # Dry run (default) — parse and count only, no DB writes:
    gmail-header-backfill --dsn postgresql://kb:<pw>@localhost:5433/kb --limit 100

    # Real run — apply the UPDATE for this slice:
    gmail-header-backfill --dsn ... --limit 1000 --offset 0 --apply

    # Next slice (stable, deterministic partitioning by id order):
    gmail-header-backfill --dsn ... --limit 1000 --offset 1000 --apply

DSN can also come from the KB_DSN env var instead of --dsn.

Idempotency: the UPDATE only touches rows that don't already carry a `headers` entity
(§5.2 `WHERE NOT EXISTS`), so re-running any slice (including after a crash) is safe —
already-backfilled rows are skipped, not double-appended.

Messages the typed (policy.default) parse rejects fall back to a degraded
compat32/raw-string parse (parse_headers_fallback) — same §4.1 entity shape, counted
separately as `parsed_fallback`, never silently mixed into ordinary successes.
"""
from __future__ import annotations

import argparse
import asyncio
import email
import email.policy
import email.utils
import json
import os
import sys
from pathlib import Path
from typing import Optional

import asyncpg
import structlog

from kb_mail.text import sanitize_surrogates as _sanitize

_log = structlog.get_logger(__name__)

DEFAULT_ARCHIVE_ROOT = Path("/home/oskar/kb/mail/archive")
DEFAULT_LIMIT = 1000
UPDATE_BATCH_SIZE = 500

_UPDATE_SQL = """
UPDATE envelope
SET entities = entities || $2::jsonb
WHERE id = $1
  AND NOT EXISTS (
      SELECT 1 FROM jsonb_array_elements(entities) e WHERE e->>'type' = 'headers'
  )
"""


def parse_headers(raw: bytes) -> dict:
    """Parse Gmail .eml headers into the `{"type": "headers", ...}` shape (plan §4.1).

    `from`/`to`/`cc` are decoded via `email.policy.default` (resolves RFC 2047
    encoded-words to real Unicode) and then split into {name, address} pairs with
    `email.utils.getaddresses`. `delivered_to` is kept as a list of raw (decoded but
    unparsed) strings — `Delivered-To` may repeat per hop, and every occurrence is
    kept (see plan §1.8/§4.1). `date_raw` comes from a separate compat32 parse: the
    default policy's DateHeader reformats the header (e.g. corrects the weekday name,
    zero-pads the day) instead of preserving the literal original text.
    """
    msg = email.message_from_bytes(raw, policy=email.policy.default)
    msg_compat = email.message_from_bytes(raw, policy=email.policy.compat32)

    from_headers = msg.get_all("From") or []
    if len(from_headers) > 1:
        _log.warning("headers.multiple_from", count=len(from_headers))

    from_obj: Optional[dict] = None
    if from_headers:
        from_addrs = email.utils.getaddresses([str(from_headers[0])])
        if from_addrs and from_addrs[0][1]:
            name, address = from_addrs[0]
            from_obj = {"name": name or None, "address": address}

    to_list = [
        {"name": name or None, "address": address}
        for name, address in email.utils.getaddresses(
            [str(h) for h in (msg.get_all("To") or [])]
        )
        if address
    ]

    cc_list = [
        {"name": name or None, "address": address}
        for name, address in email.utils.getaddresses(
            [str(h) for h in (msg.get_all("Cc") or [])]
        )
        if address
    ]

    delivered_to = [str(h) for h in (msg.get_all("Delivered-To") or [])]

    subject_header = msg.get("Subject")
    subject = str(subject_header) if subject_header is not None else None

    # compat32 .get() returns an email.header.Header (not str) when the raw
    # value contains 8-bit bytes — str() + _sanitize keeps it JSON/jsonb-safe.
    # An unguarded Header here crashed the original full run mid-slice
    # (TypeError at json.dumps), silently losing the rest of the slice.
    date_header = msg_compat.get("Date")
    date_raw = _sanitize(str(date_header)) if date_header is not None else None

    return {
        "type": "headers",
        "from": from_obj,
        "to": to_list,
        "cc": cc_list,
        "delivered_to": delivered_to,
        "subject": subject,
        "date_raw": date_raw,
    }


def parse_headers_fallback(raw: bytes) -> dict:
    """Degraded parse for messages the typed (policy.default) path rejects.

    Real-world triggers found in the archive (diagnosis 2026-07-14, 9 of
    225 030): RFC 2047 encoded-words that decode to text with CR/LF in a
    display name (ValueError in headerregistry), group syntax in To:
    ("unlisted-recipients:;"), and a _header_value_parser bug on malformed
    display names (fixed in newer CPython, present on PIHA's 3.11).

    Everything comes from a compat32 parse: address fields are split into
    {name, address} by email.utils.getaddresses over the RAW header text — no
    RFC 2047 decoding, so names may keep literal =?...?= encoded-words.
    Subject and delivered_to stay raw strings too. The returned shape is
    exactly parse_headers()'s §4.1 shape — only parse QUALITY degrades, never
    the schema. Callers count these separately (stats["parsed_fallback"]).
    """
    msg = email.message_from_bytes(raw, policy=email.policy.compat32)

    from_headers = [str(h) for h in (msg.get_all("From") or [])]
    if len(from_headers) > 1:
        _log.warning("headers.multiple_from", count=len(from_headers))

    from_obj: Optional[dict] = None
    if from_headers:
        from_addrs = email.utils.getaddresses([from_headers[0]])
        if from_addrs and from_addrs[0][1]:
            name, address = from_addrs[0]
            from_obj = {"name": _sanitize(name) or None, "address": _sanitize(address)}

    to_list = [
        {"name": _sanitize(name) or None, "address": _sanitize(address)}
        for name, address in email.utils.getaddresses(
            [str(h) for h in (msg.get_all("To") or [])]
        )
        if address
    ]

    cc_list = [
        {"name": _sanitize(name) or None, "address": _sanitize(address)}
        for name, address in email.utils.getaddresses(
            [str(h) for h in (msg.get_all("Cc") or [])]
        )
        if address
    ]

    delivered_to = [_sanitize(str(h)) for h in (msg.get_all("Delivered-To") or [])]

    subject_header = msg.get("Subject")
    subject = _sanitize(str(subject_header)) if subject_header is not None else None

    date_raw = _sanitize(msg.get("Date"))

    return {
        "type": "headers",
        "from": from_obj,
        "to": to_list,
        "cc": cc_list,
        "delivered_to": delivered_to,
        "subject": subject,
        "date_raw": date_raw,
    }


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 _has_headers(entities: list) -> bool:
    return any(isinstance(e, dict) and e.get("type") == "headers" for e in entities)


async def fetch_batch(conn: asyncpg.Connection, limit: int, offset: int) -> list:
    """One deterministic slice of `source='gmail'` envelopes, ordered by id.

    Ordering by id (not filtering out already-backfilled rows here) keeps --offset
    stable across runs — the same --limit/--offset always names the same slice of the
    225 030-row table, so partial runs can be resumed or re-verified predictably.
    Idempotency itself is enforced by the UPDATE's WHERE NOT EXISTS (see _UPDATE_SQL).
    """
    return await conn.fetch(
        """
        SELECT id, raw_ref, entities
        FROM envelope
        WHERE source = 'gmail'
        ORDER BY id
        LIMIT $1 OFFSET $2
        """,
        limit,
        offset,
    )


async def _update_batch(conn: asyncpg.Connection, rows: list[tuple[str, str]]) -> None:
    """Multi-row UPDATE for one batch — a single executemany round, not one transaction per row."""
    if not rows:
        return
    await conn.executemany(_UPDATE_SQL, rows)
    _log.debug("update_batch_flushed", count=len(rows))


async def run(
    dsn: str,
    archive_root: Path = DEFAULT_ARCHIVE_ROOT,
    limit: int = DEFAULT_LIMIT,
    offset: int = 0,
    apply: bool = False,
) -> dict[str, int]:
    """Backfill headers for one --limit/--offset slice of `source='gmail'` envelopes.

    dry-run (apply=False) parses and counts everything without writing to the DB.
    Returns stats that must always balance:

        scanned = updated + already_has_headers + parse_errors + read_errors
                  + missing_file

    `parsed_fallback` is a labeled SUBSET of `updated` (rows recovered by
    parse_headers_fallback() after the typed parse raised), not a separate
    outcome — it exists so degraded parses are never silently mixed into
    ordinary successes. Anomalies that aren't errors (e.g. more than one
    From: header) are logged (`headers.multiple_from`) rather than counted.
    """
    stats = {
        "scanned": 0,
        "already_has_headers": 0,
        "updated": 0,
        "parsed_fallback": 0,
        "read_errors": 0,
        "parse_errors": 0,
        "missing_file": 0,
    }

    conn = await asyncpg.connect(dsn)
    try:
        rows = await fetch_batch(conn, limit, offset)

        pending: list[tuple[str, str]] = []

        async def _flush() -> None:
            if apply:
                await _update_batch(conn, pending)
            pending.clear()

        for row in rows:
            stats["scanned"] += 1
            envelope_id = row["id"]
            entities = _decode_jsonb(row["entities"]) or []

            if _has_headers(entities):
                stats["already_has_headers"] += 1
                continue

            eml_path = archive_root / row["raw_ref"]
            try:
                raw = eml_path.read_bytes()
            except FileNotFoundError:
                _log.info("skip.missing_file", envelope_id=envelope_id,
                          expected_path=str(eml_path))
                stats["missing_file"] += 1
                continue
            except OSError:
                _log.warning("skip.read_error", envelope_id=envelope_id, path=str(eml_path))
                stats["read_errors"] += 1
                continue

            # json.dumps stays INSIDE the try: a non-serializable header value
            # must count as this row's parse_error, not crash the whole slice
            # (that's exactly how the original full run lost 4999 rows).
            try:
                patch = json.dumps([parse_headers(raw)])
            except Exception as typed_exc:
                try:
                    patch = json.dumps([parse_headers_fallback(raw)])
                except Exception:
                    _log.warning("skip.parse_error", envelope_id=envelope_id, exc_info=True)
                    stats["parse_errors"] += 1
                    continue
                _log.info("headers.parsed_fallback", envelope_id=envelope_id,
                          typed_error=repr(typed_exc))
                stats["parsed_fallback"] += 1

            stats["updated"] += 1
            pending.append((envelope_id, patch))
            if len(pending) >= UPDATE_BATCH_SIZE:
                await _flush()

        await _flush()
    finally:
        await conn.close()

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


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Backfill {'type': 'headers', ...} onto existing source='gmail' "
        "envelope rows in kb-postgres (module 5, phase 2 backfill — plan §5)."
    )
    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("--limit", type=int, default=DEFAULT_LIMIT,
                        help=f"Envelopes in this slice (default: {DEFAULT_LIMIT})")
    parser.add_argument("--offset", type=int, default=0,
                        help="Slice offset, ordered by envelope id (default: 0)")
    parser.add_argument("--apply", action="store_true",
                        help="Actually write the UPDATE. Default is dry-run (parse + count only).")
    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)

    stats = asyncio.run(
        run(
            dsn=args.dsn,
            archive_root=args.archive_root,
            limit=args.limit,
            offset=args.offset,
            apply=args.apply,
        )
    )

    mode = "APPLY" if args.apply else "DRY-RUN"
    _log.info("summary", mode=mode, **stats)
    failures = stats["read_errors"] + stats["parse_errors"] + stats["missing_file"]
    sys.exit(1 if failures > 0 else 0)


if __name__ == "__main__":
    main()