"""Gmail header backfill — one-shot job, module 5 phase 2 (docs/kb/modules/05-faza2-plan.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.
"""
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

_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

    date_raw = msg_compat.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: scanned, already_has_headers, updated, read_errors, parse_errors.
    Anomalies that aren't errors (e.g. more than one From: header) are logged
    (`headers.multiple_from`) rather than counted here — see parse_headers().
    """
    stats = {
        "scanned": 0,
        "already_has_headers": 0,
        "updated": 0,
        "read_errors": 0,
        "parse_errors": 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 OSError:
                _log.warning("skip.read_error", envelope_id=envelope_id, path=str(eml_path))
                stats["read_errors"] += 1
                continue

            try:
                headers = parse_headers(raw)
            except Exception:
                _log.warning("skip.parse_error", envelope_id=envelope_id, exc_info=True)
                stats["parse_errors"] += 1
                continue

            stats["updated"] += 1
            pending.append((envelope_id, json.dumps([headers])))
            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)
    sys.exit(1 if stats["read_errors"] + stats["parse_errors"] > 0 else 0)


if __name__ == "__main__":
    main()