#!/usr/bin/env python3
"""IMAP email triage helper for g@germainebrown.com.

Modes:
  --collect     Print JSON with unprocessed recent inbox message summaries.
  --move UID    Move an INBOX message UID to the Suspected Spam folder and mark processed.
  --mark UID DECISION REASON
                Mark a UID processed without moving, with an audit log entry.
  --folders     Print folder list.

The script intentionally does not delete mail and does not open links or attachments.
"""
from __future__ import annotations

import argparse
import email
import imaplib
import json
import re
import ssl
from datetime import datetime, timezone
from email.header import decode_header, make_header
from email.message import Message
from email.utils import parseaddr
from pathlib import Path
from typing import Any

HOST = "mail.germainebrown.com"
PORT = 993
USER = "g@germainebrown.com"
PASSWORD_FILE = Path("/root/.config/himalaya/g-germainebrown.pass")
INBOX = "INBOX"
SUSPECTED_SPAM = "INBOX.spam"
STATE_DIR = Path("/root/.hermes/email_triage")
STATE_FILE = STATE_DIR / "state.json"
LOG_FILE = STATE_DIR / "actions.jsonl"
MAX_BODY_CHARS = 4000

BILL_KEYWORDS = re.compile(r"\b(invoice|bill|statement|amount due|payment due|due date|past due|overdue|renewal|autopay|subscription|receipt|premium|tax|utility|utilities)\b", re.I)
SPAM_KEYWORDS = re.compile(r"\b(urgent action|verify your account|password expires|account suspended|wire transfer|gift card|lottery|prize|winner|crypto|bitcoin|act now|limited time|final notice)\b", re.I)
URL_RE = re.compile(r"https?://[^\s)>'\"]+", re.I)
AMOUNT_RE = re.compile(r"(?:USD\s*)?\$\s?\d{1,3}(?:,\d{3})*(?:\.\d{2})?")
DATE_RE = re.compile(r"\b(?:Jan(?:uary)?|Feb(?:ruary)?|Mar(?:ch)?|Apr(?:il)?|May|Jun(?:e)?|Jul(?:y)?|Aug(?:ust)?|Sep(?:t(?:ember)?)?|Oct(?:ober)?|Nov(?:ember)?|Dec(?:ember)?)\s+\d{1,2},?\s+\d{4}\b|\b\d{1,2}/\d{1,2}/\d{2,4}\b", re.I)
SUSPICIOUS_TLDS = {
    "click", "xyz", "top", "site", "online", "club", "icu", "buzz", "cyou", "cam",
    "quest", "rest", "bar", "monster", "sbs", "skin", "shop", "store", "pw",
    "win", "bid", "party", "loan", "date", "download", "stream", "review", "country", "science",
}
FREEMAIL_DOMAINS = {"gmail.com", "yahoo.com", "outlook.com", "hotmail.com", "aol.com", "icloud.com", "proton.me", "protonmail.com", "mail.com"}
KNOWN_LEGIT_DOMAINS = {
    "google.com", "apple.com", "amazon.com", "stripe.com", "bankofamerica.com", "chase.com",
    "americanexpress.com", "georgiapower.com", "spectrum.net", "spectrum.com", "rocketmoney.com",
    "venmo.com", "paypal.com", "carters.com", "theepochtimes.com", "avis.com", "avisbudget.com",
    "fedex.com", "ups.com", "usps.com", "microsoft.com", "adobe.com", "intuit.com", "xfinity.com",
    "labcorp.com", "labcorpmessage.com", "gotracktownusa.com",
    "robinhood.com",
    "southernco.com",  # parent of Georgia Power (G2georgiapps@southernco.com sends usage alerts)
    "redditmail.com",  # Reddit official notification domain (via Amazon SES) — looks random to heuristic
}
USER_BLOCKED_SPAM_DOMAINS = {
    "winecountrygiftbaskets.com", "heavy.com", "moveon.org", "lendingclub.com", "upgrade.com",
    "corporateshopping.com", "hudsonridge.com", "ancestry.com", "twentytwowords.com",
    "exchange.spectrum.com",  # sales/promo flyers sent to old Gmail, not actual Spectrum billing
}
NEWS_DOMAINS = {"theepochtimes.com", "nytimes.com", "washingtonpost.com", "wsj.com", "cnn.com", "foxnews.com"}
BRAND_WORDS = {"paypal", "apple", "amazon", "microsoft", "bank", "chase", "wells", "boa", "bank of america", "netflix", "usps", "ups", "fedex", "dhl", "stripe", "venmo", "social security"}
RANDOM_LABEL_RE = re.compile(r"^(?:[a-z]*\d+[a-z\d]*|[a-z]{10,}|[a-z\d]{12,})$", re.I)


def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def load_password() -> str:
    if not PASSWORD_FILE.exists():
        raise SystemExit(f"Missing password file: {PASSWORD_FILE}")
    return PASSWORD_FILE.read_text().strip()


def connect() -> imaplib.IMAP4_SSL:
    ctx = ssl.create_default_context()
    m = imaplib.IMAP4_SSL(HOST, PORT, ssl_context=ctx)
    m.login(USER, load_password())
    return m


def q(name: str) -> str:
    return '"' + name.replace('\\', '\\\\').replace('"', '\\"') + '"'


def load_state() -> dict[str, Any]:
    STATE_DIR.mkdir(parents=True, exist_ok=True)
    if not STATE_FILE.exists():
        return {"processed_uids": {}, "created_at": now_iso()}
    try:
        return json.loads(STATE_FILE.read_text())
    except Exception:
        return {"processed_uids": {}, "created_at": now_iso(), "state_read_error": True}


def save_state(state: dict[str, Any]) -> None:
    STATE_DIR.mkdir(parents=True, exist_ok=True)
    tmp = STATE_FILE.with_suffix(".tmp")
    tmp.write_text(json.dumps(state, indent=2, sort_keys=True))
    tmp.replace(STATE_FILE)


def log_action(entry: dict[str, Any]) -> None:
    STATE_DIR.mkdir(parents=True, exist_ok=True)
    entry = {"timestamp": now_iso(), **entry}
    with LOG_FILE.open("a") as f:
        f.write(json.dumps(entry, sort_keys=True) + "\n")


def decode_mime(value: str | None) -> str:
    if not value:
        return ""
    try:
        return str(make_header(decode_header(value)))
    except Exception:
        return value


def parse_addr(value: str | None) -> str:
    return decode_mime(value).strip()


def sender_domain(value: str | None) -> str:
    addr = parseaddr(decode_mime(value))[1].lower()
    return addr.rsplit("@", 1)[-1] if "@" in addr else ""


def base_domain(domain: str) -> str:
    parts = domain.lower().split(".")
    return ".".join(parts[-2:]) if len(parts) >= 2 else domain.lower()


def sender_reputation_hints(sender: str, subject: str) -> list[str]:
    hints: list[str] = []
    domain = sender_domain(sender)
    bdom = base_domain(domain)
    display = parseaddr(sender)[0].lower()
    combined = f"{display} {subject}".lower()
    if not domain:
        return ["sender_missing_domain"]
    if bdom in USER_BLOCKED_SPAM_DOMAINS:
        hints.append("user_blocked_spam_domain")
    if domain in KNOWN_LEGIT_DOMAINS or bdom in KNOWN_LEGIT_DOMAINS:
        hints.append("known_legitimate_sender_domain")
    tld = domain.rsplit(".", 1)[-1] if "." in domain else ""
    if tld in SUSPICIOUS_TLDS and bdom not in KNOWN_LEGIT_DOMAINS:
        hints.append(f"low_reputation_tld_.{tld}")
    if any(RANDOM_LABEL_RE.match(label) for label in domain.split(".")[:-1]) and bdom not in KNOWN_LEGIT_DOMAINS:
        hints.append("random_looking_sender_domain")
    for brand in BRAND_WORDS:
        token = brand.replace(" ", "")
        if brand in combined and token not in domain.replace("-", "") and bdom not in NEWS_DOMAINS and bdom not in KNOWN_LEGIT_DOMAINS:
            hints.append(f"possible_{token}_spoof_from_unrelated_domain")
            break
    if bdom in FREEMAIL_DOMAINS and re.search(r"\b(invoice|payment|security|account|order submitted|receipt|bank|paypal|apple|amazon)\b", combined, re.I):
        hints.append("sensitive_claim_from_freemail_sender")
    return hints


def strip_html(text: str) -> str:
    text = re.sub(r"(?is)<(script|style).*?>.*?</\1>", " ", text)
    text = re.sub(r"(?s)<[^>]+>", " ", text)
    text = re.sub(r"&nbsp;", " ", text)
    text = re.sub(r"&amp;", "&", text)
    text = re.sub(r"&lt;", "<", text)
    text = re.sub(r"&gt;", ">", text)
    return re.sub(r"\s+", " ", text).strip()


def body_text(msg: Message) -> tuple[str, list[str]]:
    attachments: list[str] = []
    parts: list[str] = []
    if msg.is_multipart():
        for part in msg.walk():
            disp = (part.get_content_disposition() or "").lower()
            ctype = part.get_content_type().lower()
            filename = decode_mime(part.get_filename()) if part.get_filename() else ""
            if filename or disp == "attachment":
                attachments.append(filename or ctype)
                continue
            if ctype in ("text/plain", "text/html"):
                try:
                    payload = part.get_payload(decode=True) or b""
                    charset = part.get_content_charset() or "utf-8"
                    txt = payload.decode(charset, errors="replace")
                    if ctype == "text/html":
                        txt = strip_html(txt)
                    parts.append(txt)
                except Exception:
                    pass
    else:
        ctype = msg.get_content_type().lower()
        try:
            payload = msg.get_payload(decode=True) or b""
            charset = msg.get_content_charset() or "utf-8"
            txt = payload.decode(charset, errors="replace")
            if ctype == "text/html":
                txt = strip_html(txt)
            parts.append(txt)
        except Exception:
            pass
    text = re.sub(r"\s+", " ", "\n".join(parts)).strip()
    return text[:MAX_BODY_CHARS], attachments


def ensure_folder(m: imaplib.IMAP4_SSL, folder: str) -> None:
    typ, data = m.list()
    folders = []
    for raw in data or []:
        line = raw.decode(errors="replace") if isinstance(raw, bytes) else str(raw)
        folders.append(line)
    if any((f'"{folder}"' in line) or line.endswith(" " + folder) for line in folders):
        return
    typ, data = m._simple_command("CREATE", q(folder))
    if typ != "OK":
        raise RuntimeError(f"Could not create folder {folder}: {data}")


def list_folders() -> None:
    with connect() as m:
        ensure_folder(m, SUSPECTED_SPAM)
        typ, data = m.list()
        print(json.dumps({"status": typ, "folders": [x.decode(errors="replace") if isinstance(x, bytes) else str(x) for x in (data or [])]}, indent=2))
        m.logout()


def fetch_message(m: imaplib.IMAP4_SSL, uid: str) -> dict[str, Any] | None:
    typ, data = m.uid("FETCH", uid, "(RFC822)")
    if typ != "OK" or not data:
        return None
    raw_msg = None
    for item in data:
        if isinstance(item, tuple) and item[1]:
            raw_msg = item[1]
            break
    if raw_msg is None:
        return None
    msg = email.message_from_bytes(raw_msg)
    text, attachments = body_text(msg)
    subject = decode_mime(msg.get("Subject"))
    sender = parse_addr(msg.get("From"))
    to = parse_addr(msg.get("To"))
    date = decode_mime(msg.get("Date"))
    message_id = (msg.get("Message-ID") or "").strip()
    urls = URL_RE.findall(text)[:12]
    combined = f"{subject}\n{sender}\n{text}"
    hints = []
    if BILL_KEYWORDS.search(combined):
        hints.append("bill_keywords")
    if SPAM_KEYWORDS.search(combined):
        hints.append("spam_keywords")
    if attachments:
        hints.append("has_attachments")
    if urls:
        hints.append("has_links")
    hints.extend(sender_reputation_hints(sender, subject))
    amounts = AMOUNT_RE.findall(combined)[:5]
    dates = DATE_RE.findall(combined)[:5]
    return {
        "uid": uid,
        "message_id": message_id,
        "from": sender,
        "to": to,
        "subject": subject,
        "date": date,
        "snippet": text[:1200],
        "attachments": attachments[:10],
        "url_count": len(urls),
        "sample_urls": urls[:5],
        "detected_amounts": amounts,
        "detected_dates": dates,
        "heuristic_hints": hints,
    }


def collect(max_messages: int, unread_only: bool) -> None:
    state = load_state()
    processed = set(state.get("processed_uids", {}).keys())
    with connect() as m:
        ensure_folder(m, SUSPECTED_SPAM)
        typ, _ = m.select(INBOX, readonly=True)
        if typ != "OK":
            raise RuntimeError("Could not select INBOX")
        criteria = "UNSEEN" if unread_only else "ALL"
        typ, data = m.uid("SEARCH", None, criteria)
        if typ != "OK":
            raise RuntimeError(f"Search failed: {data}")
        uids = (data[0].decode().split() if data and data[0] else [])
        uids = [u for u in uids if u not in processed]
        selected = list(reversed(uids)) if max_messages <= 0 else list(reversed(uids))[:max_messages]
        messages = []
        for uid in selected:
            item = fetch_message(m, uid)
            if item:
                messages.append(item)
        m.logout()
    print(json.dumps({
        "account": USER,
        "host": HOST,
        "inbox": INBOX,
        "suspected_spam_folder": SUSPECTED_SPAM,
        "mode": "collect",
        "unread_only": unread_only,
        "unprocessed_available": len(uids),
        "returned": len(messages),
        "messages": messages,
        "instruction": "Classify each message as legit, suspected_spam, bill_or_invoice, or uncertain. Move only high-confidence suspected_spam using this script's --move UID. For bills, notify the user. For legit/uncertain, use --mark UID DECISION REASON after deciding."
    }, indent=2))


def mark(uid: str, decision: str, reason: str, moved_to: str | None = None) -> None:
    state = load_state()
    state.setdefault("processed_uids", {})[uid] = {
        "decision": decision,
        "reason": reason,
        "moved_to": moved_to,
        "processed_at": now_iso(),
    }
    save_state(state)
    log_action({"uid": uid, "decision": decision, "reason": reason, "moved_to": moved_to})


def move(uid: str, reason: str) -> None:
    with connect() as m:
        ensure_folder(m, SUSPECTED_SPAM)
        typ, _ = m.select(INBOX, readonly=False)
        if typ != "OK":
            raise RuntimeError("Could not select INBOX")
        typ, data = m.uid("COPY", uid, q(SUSPECTED_SPAM))
        if typ != "OK":
            raise RuntimeError(f"COPY failed for UID {uid}: {data}")
        typ, data = m.uid("STORE", uid, "+FLAGS", "(\\Deleted)")
        if typ != "OK":
            raise RuntimeError(f"STORE Deleted failed for UID {uid}: {data}")
        m.expunge()
        m.logout()
    mark(uid, "suspected_spam", reason, SUSPECTED_SPAM)
    print(json.dumps({"status": "moved", "uid": uid, "folder": SUSPECTED_SPAM, "reason": reason}))


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--collect", action="store_true")
    ap.add_argument("--folders", action="store_true")
    ap.add_argument("--move")
    ap.add_argument("--mark", nargs=3, metavar=("UID", "DECISION", "REASON"))
    ap.add_argument("--max", type=int, default=0, help="Max messages to collect; 0 means no limit")
    ap.add_argument("--all", action="store_true", help="Collect all recent unprocessed messages, not just unread")
    ap.add_argument("--reason", default="High-confidence suspected spam/phishing")
    args = ap.parse_args()
    if args.folders:
        list_folders(); return
    if args.collect:
        collect(max_messages=args.max, unread_only=not args.all); return
    if args.move:
        move(args.move, args.reason); return
    if args.mark:
        uid, decision, reason = args.mark
        mark(uid, decision, reason)
        print(json.dumps({"status": "marked", "uid": uid, "decision": decision, "reason": reason}))
        return
    ap.print_help()


if __name__ == "__main__":
    main()