"""Read personal mail through a deliberately read-only IMAP command surface."""

from __future__ import annotations

import argparse
import hashlib
import imaplib
import json
import re
import tomllib
from dataclasses import dataclass
from datetime import date, datetime, timedelta, timezone
from email import policy
from email.header import decode_header, make_header
from email.parser import BytesParser
from email.utils import getaddresses, parsedate_to_datetime
from html import unescape
from html.parser import HTMLParser
from pathlib import Path
from typing import Any

from error_reporting import record_failure
from notify import ROOT, read_env


STATE_DIR = ROOT / "state"
MESSAGES_DIR = ROOT / "data" / "messages"
CURSORS_DIR = STATE_DIR / "cursors"
EXCLUDED_IDS_PATH = STATE_DIR / "excluded-message-ids.jsonl"
WORK_DOMAINS_PATH = ROOT / "config" / "work-domains.txt"
ACCOUNTS_PATH = ROOT / "config" / "accounts.toml"


class _TextExtractor(HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self.parts: list[str] = []

    def handle_data(self, data: str) -> None:
        self.parts.append(data)

    def text(self) -> str:
        return re.sub(r"\s+", " ", unescape(" ".join(self.parts))).strip()


def _html_to_text(value: str) -> str:
    # Real mail contains Outlook conditional comments (<![endif]-->) and other
    # malformed markup that makes html.parser raise AssertionError mid-parse.
    # Fall back to a plain tag strip rather than losing the message.
    parser = _TextExtractor()
    try:
        parser.feed(value)
        text = parser.text()
    except (AssertionError, ValueError):
        text = ""
    if text:
        return text
    stripped = re.sub(r"(?is)<(script|style).*?</\1>", " ", value)
    stripped = re.sub(r"(?s)<[^>]*>", " ", stripped)
    return re.sub(r"\s+", " ", unescape(stripped)).strip()


def _decode_header(value: str | None) -> str:
    if not value:
        return ""
    try:
        return str(make_header(decode_header(value)))
    except (LookupError, UnicodeError):
        return value


def _read_list(path: Path) -> set[str]:
    if not path.exists():
        return set()
    return {
        line.strip().lower().lstrip("@")
        for line in path.read_text(encoding="utf-8").splitlines()
        if line.strip() and not line.lstrip().startswith("#")
    }


def _matches_work(headers: list[str], exclusions: set[str]) -> bool:
    addresses = {
        address.lower()
        for _, address in getaddresses(headers)
        if address
    }
    for address in addresses:
        domain = address.rsplit("@", 1)[-1]
        for excluded in exclusions:
            if "@" in excluded and address == excluded:
                return True
            if "@" not in excluded and (
                domain == excluded or domain.endswith("." + excluded)
            ):
                return True
    return False


class _BodyStructureParser:
    TOKEN = re.compile(
        r'\s*(?:(\()|(\))|("(?:\\.|[^"\\])*")|(\{(\d+)\}\r?\n)|([^\s()]+))',
        re.DOTALL,
    )

    def __init__(self, raw: bytes) -> None:
        self.raw = raw.decode("latin-1", errors="replace")
        self.position = 0

    def parse(self) -> Any:
        start = self.raw.find("(")
        if start < 0:
            raise ValueError("BODYSTRUCTURE response did not contain a structure")
        self.position = start
        return self._value()

    def _value(self) -> Any:
        match = self.TOKEN.match(self.raw, self.position)
        if not match:
            raise ValueError("Invalid BODYSTRUCTURE response")
        self.position = match.end()
        if match.group(1):
            items = []
            while True:
                close = self.TOKEN.match(self.raw, self.position)
                if close and close.group(2):
                    self.position = close.end()
                    return items
                items.append(self._value())
        if match.group(3):
            return bytes(match.group(3)[1:-1], "latin-1").decode(
                "unicode_escape", errors="replace"
            )
        if match.group(4):
            size = int(match.group(5))
            value = self.raw[self.position : self.position + size]
            self.position += size
            return value
        atom = match.group(6)
        return None if atom and atom.upper() == "NIL" else atom


def _parameter_map(value: Any) -> dict[str, str]:
    if not isinstance(value, list):
        return {}
    result: dict[str, str] = {}
    for index in range(0, len(value) - 1, 2):
        result[str(value[index]).lower()] = str(value[index + 1])
    return result


def _parts_from_structure(
    node: Any, prefix: str = ""
) -> tuple[list[tuple[str, str, str]], list[str]]:
    text_parts: list[tuple[str, str, str]] = []
    attachments: list[str] = []
    if not isinstance(node, list) or not node:
        return text_parts, attachments

    if isinstance(node[0], list):
        child_number = 1
        for child in node:
            if not isinstance(child, list):
                break
            number = f"{prefix}.{child_number}" if prefix else str(child_number)
            child_text, child_files = _parts_from_structure(child, number)
            text_parts.extend(child_text)
            attachments.extend(child_files)
            child_number += 1
        return text_parts, attachments

    media_type = str(node[0] or "").lower()
    subtype = str(node[1] or "").lower() if len(node) > 1 else ""
    params = _parameter_map(node[2] if len(node) > 2 else None)
    disposition: Any = None
    for candidate in node[8:]:
        if (
            isinstance(candidate, list)
            and candidate
            and str(candidate[0]).lower() in {"attachment", "inline"}
        ):
            disposition = candidate
            break
    disposition_params = _parameter_map(
        disposition[1] if isinstance(disposition, list) and len(disposition) > 1 else None
    )
    filename = disposition_params.get("filename") or params.get("name")
    if filename:
        attachments.append(_decode_header(filename))

    is_attachment = bool(
        filename
        or (
            isinstance(disposition, list)
            and disposition
            and str(disposition[0]).lower() == "attachment"
        )
    )
    if media_type == "text" and subtype in {"plain", "html"} and not is_attachment:
        charset = params.get("charset", "utf-8")
        text_parts.append((prefix or "1", subtype, charset))
    return text_parts, attachments


def _fetch_literal(
    connection: imaplib.IMAP4_SSL, uid: bytes, query: str
) -> tuple[bytes, bytes]:
    result, data = connection.uid("FETCH", uid, query)
    if result != "OK":
        raise RuntimeError("IMAP fetch failed")
    metadata = b" ".join(
        item if isinstance(item, bytes) else item[0]
        for item in data
        if item and (isinstance(item, bytes) or isinstance(item, tuple))
    )
    content = b"".join(
        item[1] for item in data if isinstance(item, tuple) and len(item) == 2
    )
    return metadata, content


def _fetch_message(
    connection: imaplib.IMAP4_SSL, uid: bytes, account: str
) -> dict[str, Any]:
    _, header_bytes = _fetch_literal(connection, uid, "(BODY.PEEK[HEADER])")
    structure_meta, _ = _fetch_literal(connection, uid, "(BODYSTRUCTURE)")
    header = BytesParser(policy=policy.default).parsebytes(header_bytes)
    structure_text = structure_meta.upper().find(b"BODYSTRUCTURE")
    if structure_text < 0:
        raise RuntimeError("IMAP did not return BODYSTRUCTURE")
    structure = _BodyStructureParser(
        structure_meta[structure_text + len(b"BODYSTRUCTURE") :]
    ).parse()
    text_parts, attachments = _parts_from_structure(structure)

    plain: list[str] = []
    html: list[str] = []
    for part_number, subtype, charset in text_parts:
        _, body_bytes = _fetch_literal(
            connection, uid, f"(BODY.PEEK[{part_number}])"
        )
        try:
            decoded = body_bytes.decode(charset, errors="replace")
        except LookupError:
            decoded = body_bytes.decode("utf-8", errors="replace")
        (plain if subtype == "plain" else html).append(decoded)

    body = "\n".join(plain).strip()
    if not body and html:
        body = _html_to_text("\n".join(html))

    message_id = str(header.get("Message-ID", "")).strip()
    if not message_id:
        fallback = hashlib.sha256(
            header_bytes + body.encode("utf-8", errors="replace")
        ).hexdigest()
        message_id = f"missing:{fallback}"

    try:
        source_date = parsedate_to_datetime(str(header.get("Date", "")))
        if source_date and source_date.tzinfo is None:
            source_date = source_date.replace(tzinfo=timezone.utc)
        date_value = source_date.astimezone(timezone.utc).isoformat() if source_date else None
    except (TypeError, ValueError, OverflowError):
        date_value = None

    return {
        "account": account,
        "message_id": message_id,
        "uid": int(uid),
        "from": _decode_header(str(header.get("From", ""))),
        "to": _decode_header(str(header.get("To", ""))),
        "reply_to": _decode_header(str(header.get("Reply-To", ""))),
        "subject": _decode_header(str(header.get("Subject", ""))),
        "date": date_value,
        "body": body,
        "body_sha256": hashlib.sha256(body.encode("utf-8")).hexdigest(),
        "attachments": attachments,
    }


def _cursor_path(alias: str) -> Path:
    return CURSORS_DIR / f"{alias}.json"


def _load_cursor(alias: str) -> dict[str, Any] | None:
    path = _cursor_path(alias)
    if not path.exists():
        return None
    return json.loads(path.read_text(encoding="utf-8"))


def _known_message_ids() -> set[tuple[str, str]]:
    known: set[tuple[str, str]] = set()
    if not MESSAGES_DIR.exists():
        return known
    for path in MESSAGES_DIR.glob("*.jsonl"):
        for line in path.read_text(encoding="utf-8").splitlines():
            if line.strip():
                record = json.loads(line)
                # Records created by the original single-account collector belong
                # to signup, whose cursor is migrated with them.
                known.add((record.get("account", "signup"), record["message_id"]))
    return known


def _known_excluded_hashes() -> set[str]:
    if not EXCLUDED_IDS_PATH.exists():
        return set()
    return {
        line.strip()
        for line in EXCLUDED_IDS_PATH.read_text(encoding="ascii").splitlines()
        if line.strip()
    }


def _atomic_json(path: Path, value: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"
    )
    temporary.replace(path)


@dataclass
class Collection:
    messages: list[dict[str, Any]]
    excluded_work_count: int
    coverage_rebaselined: bool
    covered_to_utc: str
    notes: list[str] | None = None


# The single-account collector that lived here was deleted 2026-07-29.
# It had been dead since the multi-account rewrite and could not have run if called:
# it referenced load_env() (never imported), CURSOR_PATH (never defined), and called
# _load_cursor() and _fetch_message() with the wrong arity. 118 lines that read like
# the live collector and would have NameErrored on the first line of the try block.
# Recoverable at baseline commit a6bf776.


@dataclass(frozen=True)
class Account:
    alias: str
    imap_host: str
    username: str
    password_env: str


def _load_accounts() -> list[Account]:
    raw = tomllib.loads(ACCOUNTS_PATH.read_text(encoding="utf-8"))
    accounts: list[Account] = []
    aliases: set[str] = set()
    for entry in raw.get("accounts", []):
        account = Account(
            alias=str(entry["alias"]),
            imap_host=str(entry["imap_host"]),
            username=str(entry["username"]),
            password_env=str(entry["password_env"]),
        )
        if not re.fullmatch(r"[a-z0-9][a-z0-9_-]*", account.alias):
            raise RuntimeError(f"Invalid account alias: {account.alias}")
        if account.alias in aliases:
            raise RuntimeError(f"Duplicate account alias: {account.alias}")
        if not (
            account.password_env == "PCOS_IMAP_APP_PASSWORD"
            or account.password_env.startswith("PCOS_IMAP_APP_PASSWORD_")
        ):
            raise RuntimeError(
                f"Account {account.alias} has an invalid password environment name."
            )
        aliases.add(account.alias)
        accounts.append(account)
    if not accounts:
        raise RuntimeError("No accounts configured in config/accounts.toml.")
    return accounts


def _build_search_criteria(search_terms: list[str], since: date | None) -> str:
    text_criteria = []
    for term in search_terms:
        escaped = term.replace("\\", "\\\\").replace('"', '\\"')
        text_criteria.append(f'TEXT "{escaped}"')
    while len(text_criteria) > 1:
        right = text_criteria.pop()
        left = text_criteria.pop()
        text_criteria.append(f"OR {left} {right}")
    criteria = text_criteria[0]
    if since is not None:
        criteria = f"SINCE {since.strftime('%d-%b-%Y')} {criteria}"
    return criteria


def _collect_account(
    account: Account,
    password: str,
    since: date | None,
    search_terms: list[str] | None,
    exclusions: set[str],
    known_ids: set[tuple[str, str]],
    excluded_hashes: set[str],
    now: datetime,
) -> tuple[list[dict[str, Any]], int, bool, str | None]:
    cursor = _load_cursor(account.alias)
    default_since = (now - timedelta(days=7)).strftime("%d-%b-%Y")
    overlap_since = (now - timedelta(hours=48)).strftime("%d-%b-%Y")
    connection: imaplib.IMAP4_SSL | None = None
    try:
        try:
            connection = imaplib.IMAP4_SSL(account.imap_host, 993, timeout=30)
        except (OSError, imaplib.IMAP4.error) as error:
            summary = record_failure(f"{account.alias} IMAP TLS connection failed", error)
            return [], 0, False, f"{account.alias}: IMAP TLS connection failed: {summary}"
        try:
            connection.login(account.username, password.replace(" ", ""))
        except imaplib.IMAP4.error as error:
            summary = record_failure(f"{account.alias} IMAP authentication failed", error)
            return [], 0, False, f"{account.alias}: IMAP authentication failed: {summary}"
        try:
            # readonly=True makes imaplib issue EXAMINE and maintain its state.
            result, _ = connection.select("INBOX", readonly=True)
        except imaplib.IMAP4.error as error:
            summary = record_failure(f"{account.alias} IMAP EXAMINE failed", error)
            return [], 0, False, f"{account.alias}: could not examine INBOX read-only: {summary}"
        if result != "OK":
            return [], 0, False, f"{account.alias}: could not examine INBOX read-only."

        uidvalidity_values = connection.response("UIDVALIDITY")[1]
        if not uidvalidity_values:
            return [], 0, False, f"{account.alias}: UIDVALIDITY was unavailable."
        uidvalidity = int(uidvalidity_values[0])
        rebaselined = bool(cursor and cursor.get("uidvalidity") != uidvalidity)

        if search_terms:
            criteria = _build_search_criteria(search_terms, since)
            result, data = connection.uid("SEARCH", None, criteria)
        elif since is not None:
            result, data = connection.uid(
                "SEARCH", None, "SINCE", since.strftime("%d-%b-%Y")
            )
        elif cursor and not rebaselined:
            last_uid = int(cursor.get("last_uid", 0))
            result, data = connection.uid("SEARCH", None, f"UID {last_uid + 1}:*")
        else:
            result, data = connection.uid("SEARCH", None, "SINCE", default_since)
        if result != "OK":
            return [], 0, rebaselined, f"{account.alias}: primary search failed."
        uids = set(data[0].split()) if data and data[0] else set()

        if search_terms:
            print(f"{account.alias}: {len(uids)} search hits.", flush=True)
        else:
            result, overlap_data = connection.uid("SEARCH", None, "SINCE", overlap_since)
            if result != "OK":
                return [], 0, rebaselined, f"{account.alias}: overlap search failed."
            if overlap_data and overlap_data[0]:
                uids.update(overlap_data[0].split())

        new_messages: list[dict[str, Any]] = []
        excluded = 0
        max_uid = int(cursor.get("last_uid", 0)) if cursor and not rebaselined else 0
        skipped_bad = 0

        # An account cannot be excluded from itself. rateright.com.au is on the work
        # list, so checking that mailbox against the full list would empty it — which
        # is why there used to be a blanket "skip the check for rateright" carve-out.
        #
        # That carve-out was too wide: it also stopped LFCS mail forwarded to
        # admin@rateright.com.au from being excluded, and one such message from
        # 2026-02-08 ("Fwd: Fw: Eastern Creek LTP drawings for quotation",
        # admin@lfcs.com.au -> admin@rateright.com.au) is in the store because of it.
        # Dropping only the account's own domain keeps every other work domain live.
        own_domain = account.username.rsplit("@", 1)[-1].lower()
        account_exclusions = {value for value in exclusions if value != own_domain}

        for uid in sorted(uids, key=int):
            max_uid = max(max_uid, int(uid))
            try:
                message = _fetch_message(connection, uid, account.alias)
            except Exception as error:
                record_failure(f"{account.alias} message parse failed", error)
                skipped_bad += 1
                continue
            if _matches_work(
                [message["from"], message["to"], message["reply_to"]], account_exclusions
            ):
                excluded_hash = hashlib.sha256(
                    f"{account.alias}\0{message['message_id']}".encode("utf-8")
                ).hexdigest()
                if excluded_hash not in excluded_hashes:
                    excluded += 1
                    excluded_hashes.add(excluded_hash)
                    EXCLUDED_IDS_PATH.parent.mkdir(parents=True, exist_ok=True)
                    with EXCLUDED_IDS_PATH.open(
                        "a", encoding="ascii", newline="\n"
                    ) as excluded_file:
                        excluded_file.write(excluded_hash + "\n")
                continue
            identity = (account.alias, message["message_id"])
            if identity not in known_ids:
                new_messages.append(message)
                known_ids.add(identity)

        MESSAGES_DIR.mkdir(parents=True, exist_ok=True)
        output_path = MESSAGES_DIR / f"{now.date().isoformat()}.jsonl"
        with output_path.open("a", encoding="utf-8", newline="\n") as handle:
            for message in new_messages:
                handle.write(json.dumps(message, ensure_ascii=False) + "\n")
        for line in output_path.read_text(encoding="utf-8").splitlines():
            if line.strip():
                json.loads(line)

        # A historical sweep must not disturb the daily cursor.
        if since is None and not search_terms:
            _atomic_json(
                _cursor_path(account.alias),
                {
                    "uidvalidity": uidvalidity,
                    "last_uid": max_uid,
                    "last_run_utc": now.isoformat(),
                },
            )
        note = (
            f"{account.alias}: {skipped_bad} messages could not be parsed."
            if skipped_bad
            else None
        )
        return new_messages, excluded, rebaselined, note
    finally:
        if connection is not None:
            try:
                connection.logout()
            except imaplib.IMAP4.error:
                pass


def collect(
    since: date | None = None, search_terms: list[str] | None = None
) -> Collection:
    environment = read_env()
    now = datetime.now(timezone.utc)
    exclusions = _read_list(WORK_DOMAINS_PATH)
    known_ids = _known_message_ids()
    excluded_hashes = _known_excluded_hashes()
    messages: list[dict[str, Any]] = []
    excluded = 0
    rebaselined = False
    notes: list[str] = []

    for account in _load_accounts():
        password = environment.get(account.password_env, "")
        if not password:
            notes.append(f"{account.alias}: skipped; {account.password_env} is not set.")
            continue
        account_messages, account_excluded, account_rebaselined, note = (
            _collect_account(
                account,
                password,
                since,
                search_terms,
                exclusions,
                known_ids,
                excluded_hashes,
                now,
            )
        )
        messages.extend(account_messages)
        excluded += account_excluded
        rebaselined = rebaselined or account_rebaselined
        if note:
            notes.append(note)
    return Collection(messages, excluded, rebaselined, now.isoformat(), notes)


def _parse_since(value: str) -> date:
    try:
        parsed = date.fromisoformat(value)
    except ValueError as error:
        raise argparse.ArgumentTypeError("--since must be YYYY-MM-DD") from error
    if parsed > datetime.now(timezone.utc).date():
        raise argparse.ArgumentTypeError("--since cannot be in the future")
    return parsed


def main() -> int:
    parser = argparse.ArgumentParser(description="Collect configured IMAP accounts.")
    parser.add_argument(
        "--since",
        type=_parse_since,
        help="one-off historical sweep from YYYY-MM-DD; live cursors are unchanged",
    )
    parser.add_argument(
        "--search",
        action="append",
        dest="search_terms",
        metavar="TERM",
        help="server-side text search; repeat to OR multiple terms",
    )
    arguments = parser.parse_args()
    result = collect(arguments.since, arguments.search_terms)
    print(f"Collected {len(result.messages)} new messages.")
    for note in result.notes or []:
        print(note)
    return 0


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except Exception as error:
        summary = record_failure("Collection failed", error)
        print(f"Collection failed: {summary}")
        raise SystemExit(1)
