"""Gmail Takeout bulk importer — one-shot job.

Install (from repo root):
    pip install -e packages/kb-mail/
    pip install -e jobs/gmail-bulk-import/

Usage:
    # Archive only (no DB):
    gmail-bulk-import --mbox ~/takeout/allmail.mbox --archive /data/kb/archive

    # With kb-postgres on PIHA (locally on PIHA):
    gmail-bulk-import --mbox ~/takeout/allmail.mbox --archive /data/kb/archive \\
        --dsn postgresql://kb:<pw>@localhost:5433/kb

    # From another host pointing at PIHA:
    gmail-bulk-import ... --dsn postgresql://kb:<pw>@piha:5433/kb

    # Dry run — count and parse, no writes:
    gmail-bulk-import --mbox ~/takeout/allmail.mbox --archive /data/kb/archive --dry-run

    # Sample first 100 messages before full import:
    gmail-bulk-import --mbox ~/takeout/allmail.mbox --archive /data/kb/archive --limit 100

Run at low priority on PIHA:
    ionice -c 3 nice -n 19 gmail-bulk-import --mbox ... --archive ...
"""
from __future__ import annotations

import argparse
import asyncio
import email.policy
import hashlib
import json
import mailbox
import sys
from datetime import datetime, timezone
from email.utils import parsedate_to_datetime
from pathlib import Path
from typing import Optional

import asyncpg
import structlog

from kb_mail.archive import save_eml
from kb_mail.envelope import Envelope
from kb_mail.text import sanitize_surrogates as _sanitize

_log = structlog.get_logger(__name__)

BATCH_SIZE = 500
_EPOCH = datetime(1970, 1, 1, tzinfo=timezone.utc)

# Mirrors archive._UNSAFE — used to reconstruct raw_ref for already-archived messages.
_ARCHIVE_UNSAFE = str.maketrans({"/": "_", "\\": "_", ":": "_", "<": "", ">": ""})


def _eml_ref(envelope_id: str, ts: datetime) -> str:
    """Compute the archive-relative .eml path without writing (mirrors save_eml logic)."""
    safe_id = envelope_id.translate(_ARCHIVE_UNSAFE)
    return f"gmail/{ts.year:04d}/{ts.month:02d}/{safe_id}.eml"


def _message_id(msg: mailbox.mboxMessage) -> str:
    """Return a stable envelope id from Message-ID header or SHA-256 content hash.

    compat32 .get() returns an email.header.Header (not str) when the raw value
    holds 8-bit bytes — str() + _sanitize keeps the id jsonb-safe and gives
    .strip() a real string to work on. An unguarded .strip() on a Header raised
    AttributeError and, because _message_id runs before the per-message try,
    killed the whole run (the archive has proven 8-bit header bytes).
    """
    mid = _sanitize(str(msg.get("Message-ID", ""))).strip().strip("<>")
    if mid:
        return mid
    return "sha256-" + hashlib.sha256(msg.as_bytes()).hexdigest()[:32]


def _parse_date(msg: mailbox.mboxMessage) -> datetime:
    """Parse Date header to UTC-aware datetime. Returns epoch as last resort.

    str() first: a compat32 8-bit Date surfaces as an email.header.Header, and
    parsedate_to_datetime() raises on a Header (→ needless epoch_fallback) even
    when str(header) is perfectly parseable.
    """
    date_str = str(msg.get("Date", ""))
    if date_str:
        try:
            dt = parsedate_to_datetime(date_str)
            if dt.tzinfo is None:
                dt = dt.replace(tzinfo=timezone.utc)
            return dt.astimezone(timezone.utc)
        except Exception:
            pass
    return _EPOCH


def _parse_attachments(msg: mailbox.mboxMessage) -> list[dict]:
    """Return attachment descriptors from MIME structure.

    A part is an attachment if it has a filename, Content-Disposition: attachment,
    or is not inline text/plain or text/html.
    Errors on individual parts are logged and skipped — never propagated.
    """
    attachments = []
    for part in msg.walk():
        if part.get_content_maintype() == "multipart":
            continue
        content_type = part.get_content_type()
        disposition = part.get_content_disposition() or ""
        filename = part.get_filename()

        if (
            content_type in ("text/plain", "text/html")
            and disposition != "attachment"
            and not filename
        ):
            continue

        try:
            payload = part.get_payload(decode=True)
            if payload is None:
                continue
            attachments.append({
                "type": "attachment",
                "filename": filename,
                "content_type": content_type,
                "size": len(payload),
                "sha256": hashlib.sha256(payload).hexdigest(),
            })
        except Exception:
            _log.warning("attachment_parse_failed",
                         filename=filename, content_type=content_type)
    return attachments


async def _insert_batch(conn: asyncpg.Connection, envs: list[Envelope]) -> None:
    rows = [
        (
            env.id,
            env.source,
            env.ts,
            json.dumps(env.geo) if env.geo is not None else None,
            env.raw_ref,
            json.dumps(env.entities),
        )
        for env in envs
    ]
    await conn.executemany(
        """
        INSERT INTO envelope (id, source, ts, geo, raw_ref, entities)
        VALUES ($1, $2, $3, $4::jsonb, $5, $6::jsonb)
        ON CONFLICT (id) DO NOTHING
        """,
        rows,
    )
    _log.debug("batch_flushed", count=len(envs))


async def run_import(
    mbox_path: Path,
    archive_root: Path,
    dsn: Optional[str] = None,
    dry_run: bool = False,
    limit: Optional[int] = None,
) -> dict[str, int]:
    """Import all messages from an mbox file.

    Returns stats: {processed, imported, skipped, errors, db_inserted,
    db_insert_failed, epoch_fallback, msgs_with_attachments, total_attachments,
    total_attachment_bytes}.

    Two invariants always hold (checked by main() for the exit code):
      * processed == imported + skipped + errors
      * with a DB connection: db_inserted + db_insert_failed == imported + skipped

    `imported`/`skipped` count archive writes; `db_inserted`/`db_insert_failed`
    count DB rows. They are tracked separately because a failed batch insert
    leaves the archive ahead of the DB, and nothing else would catch the drift.

    dry_run=True parses and counts everything without writing.
    limit=N stops after processing N messages (for sampling).
    """
    stats: dict[str, int] = {
        "processed": 0,
        "imported": 0,
        "skipped": 0,
        "errors": 0,
        "db_inserted": 0,
        "db_insert_failed": 0,
        "epoch_fallback": 0,
        "msgs_with_attachments": 0,
        "total_attachments": 0,
        "total_attachment_bytes": 0,
    }

    conn: Optional[asyncpg.Connection] = None
    if not dry_run and dsn:
        conn = await asyncpg.connect(dsn)

    pending: list[Envelope] = []

    async def _flush() -> None:
        # Always clear pending, even when the insert fails: a poisoned batch
        # (DB down, one bad row) must never be re-flushed by every subsequent
        # message. A failed insert leaves the archive ahead of the DB, counted
        # as db_insert_failed so the drift is visible instead of silent.
        if conn is not None and pending:
            try:
                await _insert_batch(conn, pending)
                stats["db_inserted"] += len(pending)
            except Exception:
                _log.exception("db_insert_failed", count=len(pending))
                stats["db_insert_failed"] += len(pending)
        pending.clear()

    try:
        mbox = mailbox.mbox(str(mbox_path), create=False)
        for msg in mbox:
            if limit is not None and stats["processed"] >= limit:
                break

            stats["processed"] += 1

            # The ENTIRE per-message body is guarded: a single malformed
            # message (8-bit header, unwalkable MIME, bad save) counts as one
            # error and never aborts the run.
            envelope_id = "<unknown>"
            try:
                envelope_id = _message_id(msg)
                ts = _parse_date(msg)

                if ts == _EPOCH:
                    stats["epoch_fallback"] += 1
                    _log.warning("epoch_fallback", id=envelope_id)

                attachments = _parse_attachments(msg)
                stats["total_attachments"] += len(attachments)
                stats["total_attachment_bytes"] += sum(a["size"] for a in attachments)
                if attachments:
                    stats["msgs_with_attachments"] += 1

                if dry_run:
                    stats["imported"] += 1
                    continue

                raw = msg.as_bytes(policy=email.policy.compat32)
                try:
                    raw_ref = await save_eml(archive_root, envelope_id, "gmail", ts, raw)
                    stats["imported"] += 1
                except FileExistsError:
                    raw_ref = _eml_ref(envelope_id, ts)
                    stats["skipped"] += 1

                pending.append(
                    Envelope(id=envelope_id, source="gmail", ts=ts,
                             raw_ref=raw_ref, entities=attachments)
                )
                if len(pending) >= BATCH_SIZE:
                    await _flush()

            except Exception:
                _log.exception("message_failed", id=envelope_id)
                stats["errors"] += 1

            if stats["processed"] % 500 == 0:
                _log.info("progress", **stats)

        # Final flush swallows its own insert error (see _flush) so the run
        # always reaches import_complete with balanced stats instead of dying
        # here with the counters lost.
        await _flush()

    finally:
        if conn:
            await conn.close()

    _log.info("import_complete", dry_run=dry_run, **stats)
    return stats


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Import Gmail Takeout .mbox into KB archive and envelope DB."
    )
    parser.add_argument("--mbox", required=True, type=Path, metavar="PATH",
                        help="Path to Gmail Takeout .mbox file")
    parser.add_argument("--archive", required=True, type=Path, metavar="DIR",
                        help="Archive root directory")
    parser.add_argument("--dsn", metavar="DSN",
                        help="asyncpg DSN for kb-postgres "
                             "(e.g. postgresql://kb:<pw>@localhost:5433/kb on PIHA)")
    parser.add_argument("--dry-run", action="store_true",
                        help="Parse and count only; no writes to archive or DB")
    parser.add_argument("--limit", type=int, metavar="N",
                        help="Process at most N messages (sample before full import)")
    args = parser.parse_args()

    if not args.mbox.is_file():
        _log.error("mbox_not_found", path=str(args.mbox))
        sys.exit(1)

    stats = asyncio.run(
        run_import(args.mbox, args.archive, args.dsn or None, args.dry_run, args.limit)
    )

    balanced = (
        stats["processed"] == stats["imported"] + stats["skipped"] + stats["errors"]
    )
    if not balanced:
        _log.error("stats_unbalanced", **stats)

    # Non-zero exit on any message failure, any DB drift (archive ahead of DB),
    # or a stats imbalance — the run must not report clean when it isn't.
    failed = stats["errors"] > 0 or stats["db_insert_failed"] > 0 or not balanced
    sys.exit(1 if failed else 0)


if __name__ == "__main__":
    main()
