#!/usr/bin/env python3
"""Lightweight Fastmail JMAP client: email search, calendar search, drafts.

The headless sibling of the Fastmail MCP server, in the same mold as
tools/twenty_client.py: stdlib only, token auth, small surface. The
/gather-recent-context command uses the MCP tools when a session is
interactive; this client covers cron/CI runs and ad-hoc shell use.

Auth: set FASTMAIL_API_TOKEN (Fastmail Settings -> Privacy & Security ->
API tokens; scopes: Mail read + write for drafting, Calendar read).
Optionally FASTMAIL_SESSION_URL (default: https://api.fastmail.com/jmap/session).

CLI:
  python3 tools/fastmail_client.py search --query "interview" --days 14
  python3 tools/fastmail_client.py search --from recruiter@acme.com --limit 20
  python3 tools/fastmail_client.py thread  --id <threadId>
  python3 tools/fastmail_client.py events --days 7 [--query "onsite"]
  python3 tools/fastmail_client.py draft  --to a@b.com --subject "…" --body "…"
  python3 tools/fastmail_client.py reply  --email-id <id> --body "…"

Every subcommand prints JSON to stdout. Drafts are saved to the Drafts
mailbox with $draft set - this client NEVER sends mail.

The query/draft payload builders are pure functions so they stay unit-tested
without a network (tests/test_fastmail_client.py).
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import urllib.error
import urllib.request
from datetime import UTC, datetime, timedelta

SESSION_URL = "https://api.fastmail.com/jmap/session"
CAP_CORE = "urn:ietf:params:jmap:core"
CAP_MAIL = "urn:ietf:params:jmap:mail"
CAP_SUBMISSION = "urn:ietf:params:jmap:submission"
CAP_CALENDARS = "urn:ietf:params:jmap:calendars"

EMAIL_PROPERTIES = ["id", "threadId", "messageId", "inReplyTo", "references",
                    "subject", "from", "to", "cc", "replyTo", "receivedAt",
                    "preview", "keywords", "mailboxIds"]
EVENT_PROPERTIES = ["id", "calendarIds", "title", "description", "start",
                    "duration", "timeZone", "location", "participants", "status"]


class FastmailError(RuntimeError):
    pass


# -- pure builders (unit-tested, no I/O) -------------------------------------

def iso_days_ago(days: int, now: datetime | None = None) -> str:
    """UTC date-time string JMAP accepts, `days` before now."""
    now = now or datetime.now(UTC)
    return (now - timedelta(days=days)).strftime("%Y-%m-%dT%H:%M:%SZ")


def build_email_filter(query: str | None = None, from_: str | None = None,
                       subject: str | None = None, after: str | None = None,
                       before: str | None = None) -> dict:
    """Email/query FilterCondition. Conditions AND together server-side."""
    f = {}
    if query:
        f["text"] = query
    if from_:
        f["from"] = from_
    if subject:
        f["subject"] = subject
    if after:
        f["after"] = after
    if before:
        f["before"] = before
    return f


def build_event_filter(after: str | None = None, before: str | None = None,
                       query: str | None = None) -> dict:
    f = {}
    if after:
        f["after"] = after
    if before:
        f["before"] = before
    if query:
        f["text"] = query
    return f


def build_draft(drafts_mailbox_id: str, from_email: str, to: list[str],
                subject: str, body: str, in_reply_to: dict | None = None) -> dict:
    """Email/set `create` object for a plain-text draft.

    ``in_reply_to`` is the original Email object (JMAP properties) when the
    draft answers a message: threading headers (inReplyTo/references) are
    carried over per RFC 5322, the subject gains a single ``Re:``, and the
    recipients default to the original sender (honoring Reply-To).
    """
    draft_headers = {}
    if in_reply_to:
        orig_msg_ids = in_reply_to.get("messageId") or []
        if orig_msg_ids:
            draft_headers["inReplyTo"] = orig_msg_ids
            draft_headers["references"] = list(in_reply_to.get("references") or []) + orig_msg_ids
        if not subject:
            orig_subject = in_reply_to.get("subject") or ""
            subject = orig_subject if orig_subject.lower().startswith("re:") \
                else f"Re: {orig_subject}"
        if not to:
            counterparty = (in_reply_to.get("replyTo") or in_reply_to.get("from") or [])
            to = [p["email"] for p in counterparty if p.get("email")]
    return {
        "mailboxIds": {drafts_mailbox_id: True},
        "keywords": {"$draft": True, "$seen": True},
        "from": [{"email": from_email}],
        "to": [{"email": addr} for addr in to],
        "subject": subject,
        "bodyStructure": {"partId": "body", "type": "text/plain"},
        "bodyValues": {"body": {"value": body, "charset": "utf-8"}},
        **draft_headers,
    }


# -- client -------------------------------------------------------------------

class FastmailClient:
    def __init__(self, token: str | None = None, session_url: str | None = None):
        self.token = token or os.environ.get("FASTMAIL_API_TOKEN")
        if not self.token:
            raise FastmailError(
                "FASTMAIL_API_TOKEN is not set. Create one under Fastmail Settings "
                "-> Privacy & Security -> API tokens, or use the Fastmail MCP "
                "session instead (interactive path in /gather-recent-context)."
            )
        self.session_url = session_url or os.environ.get("FASTMAIL_SESSION_URL", SESSION_URL)
        self._session = None

    # transport ---------------------------------------------------------------

    def _http(self, url: str, body: dict | None = None) -> dict:
        req = urllib.request.Request(
            url,
            data=json.dumps(body).encode() if body is not None else None,
            method="POST" if body is not None else "GET",
            headers={"Authorization": f"Bearer {self.token}",
                     "Content-Type": "application/json"},
        )
        try:
            with urllib.request.urlopen(req, timeout=60) as resp:
                return json.loads(resp.read().decode())
        except urllib.error.HTTPError as e:
            raise FastmailError(f"{url} -> {e.code}: {e.read().decode(errors='replace')[:300]}") from e

    @property
    def session(self) -> dict:
        if self._session is None:
            self._session = self._http(self.session_url)
        return self._session

    def account_id(self, capability: str = CAP_MAIL) -> str:
        acct = self.session.get("primaryAccounts", {}).get(capability)
        if not acct:
            raise FastmailError(f"no primary account for {capability} "
                                "(token missing the scope?)")
        return acct

    def call(self, method_calls: list, using: list[str]) -> list:
        payload = self._http(self.session["apiUrl"],
                             {"using": [CAP_CORE, *using], "methodCalls": method_calls})
        responses = payload.get("methodResponses", [])
        for name, args, _tag in responses:
            if name == "error":
                raise FastmailError(f"JMAP error: {json.dumps(args)[:300]}")
        return responses

    @staticmethod
    def _response(responses: list, method: str) -> dict:
        for name, args, _tag in responses:
            if name == method:
                return args
        raise FastmailError(f"missing {method} response")

    # email --------------------------------------------------------------------

    def search_email(self, query: str | None = None, from_: str | None = None,
                     subject: str | None = None, days: int | None = None,
                     after: str | None = None, before: str | None = None,
                     limit: int = 30) -> list[dict]:
        acct = self.account_id(CAP_MAIL)
        filt = build_email_filter(query, from_, subject,
                                  after or (iso_days_ago(days) if days else None), before)
        responses = self.call([
            ["Email/query", {"accountId": acct, "filter": filt, "limit": limit,
                             "sort": [{"property": "receivedAt", "isAscending": False}],
                             "collapseThreads": True}, "q"],
            ["Email/get", {"accountId": acct, "properties": EMAIL_PROPERTIES,
                           "#ids": {"resultOf": "q", "name": "Email/query",
                                    "path": "/ids"}}, "g"],
        ], using=[CAP_MAIL])
        return self._response(responses, "Email/get")["list"]

    def get_thread(self, thread_id: str, with_bodies: bool = True) -> list[dict]:
        acct = self.account_id(CAP_MAIL)
        props = EMAIL_PROPERTIES + (["textBody", "bodyValues"] if with_bodies else [])
        get_args = {"accountId": acct, "properties": props,
                    "#ids": {"resultOf": "t", "name": "Thread/get", "path": "/list/0/emailIds"}}
        if with_bodies:
            get_args.update({"fetchTextBodyValues": True, "maxBodyValueBytes": 16384})
        responses = self.call([
            ["Thread/get", {"accountId": acct, "ids": [thread_id]}, "t"],
            ["Email/get", get_args, "g"],
        ], using=[CAP_MAIL])
        emails = self._response(responses, "Email/get")["list"]
        return sorted(emails, key=lambda e: e.get("receivedAt") or "")

    def get_email(self, email_id: str) -> dict:
        acct = self.account_id(CAP_MAIL)
        responses = self.call([
            ["Email/get", {"accountId": acct, "ids": [email_id],
                           "properties": EMAIL_PROPERTIES}, "g"],
        ], using=[CAP_MAIL])
        found = self._response(responses, "Email/get")["list"]
        if not found:
            raise FastmailError(f"email {email_id} not found")
        return found[0]

    # calendar -------------------------------------------------------------------

    def search_events(self, days: int = 7, past_days: int = 0,
                      query: str | None = None, limit: int = 50) -> list[dict]:
        acct = self.account_id(CAP_CALENDARS)
        now = datetime.now(UTC)
        filt = build_event_filter(
            after=(now - timedelta(days=past_days)).strftime("%Y-%m-%dT%H:%M:%SZ"),
            before=(now + timedelta(days=days)).strftime("%Y-%m-%dT%H:%M:%SZ"),
            query=query)
        responses = self.call([
            ["CalendarEvent/query", {"accountId": acct, "filter": filt, "limit": limit}, "q"],
            ["CalendarEvent/get", {"accountId": acct, "properties": EVENT_PROPERTIES,
                                   "#ids": {"resultOf": "q", "name": "CalendarEvent/query",
                                            "path": "/ids"}}, "g"],
        ], using=[CAP_CALENDARS])
        return self._response(responses, "CalendarEvent/get")["list"]

    # drafts -----------------------------------------------------------------------

    def _drafts_mailbox_id(self) -> str:
        acct = self.account_id(CAP_MAIL)
        responses = self.call([
            ["Mailbox/query", {"accountId": acct, "filter": {"role": "drafts"}}, "q"],
        ], using=[CAP_MAIL])
        ids = self._response(responses, "Mailbox/query")["ids"]
        if not ids:
            raise FastmailError("no Drafts mailbox found")
        return ids[0]

    def _from_email(self) -> str:
        # The session username is the primary sending identity on Fastmail.
        user = self.session.get("username")
        if not user:
            raise FastmailError("session carries no username to send from")
        return user

    def create_draft(self, to: list[str], subject: str, body: str,
                     reply_to_email_id: str | None = None) -> dict:
        """Save a plain-text draft; NEVER sends. Returns the created Email stub."""
        acct = self.account_id(CAP_MAIL)
        original = self.get_email(reply_to_email_id) if reply_to_email_id else None
        draft = build_draft(self._drafts_mailbox_id(), self._from_email(),
                            to, subject, body, in_reply_to=original)
        responses = self.call([
            ["Email/set", {"accountId": acct, "create": {"draft": draft}}, "s"],
        ], using=[CAP_MAIL])
        result = self._response(responses, "Email/set")
        if result.get("notCreated"):
            raise FastmailError(f"draft rejected: {json.dumps(result['notCreated'])[:300]}")
        return result["created"]["draft"]


# -- CLI -------------------------------------------------------------------------

def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    sub = ap.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("search", help="search email")
    p.add_argument("--query"); p.add_argument("--from", dest="from_")
    p.add_argument("--subject"); p.add_argument("--days", type=int)
    p.add_argument("--limit", type=int, default=30)

    p = sub.add_parser("thread", help="read a thread")
    p.add_argument("--id", required=True)
    p.add_argument("--no-bodies", action="store_true")

    p = sub.add_parser("events", help="search calendar events")
    p.add_argument("--days", type=int, default=7, help="days ahead (default 7)")
    p.add_argument("--past-days", type=int, default=0)
    p.add_argument("--query")

    p = sub.add_parser("draft", help="save a plain-text draft (never sends)")
    p.add_argument("--to", action="append", required=True)
    p.add_argument("--subject", required=True)
    p.add_argument("--body", required=True)

    p = sub.add_parser("reply", help="draft a reply to an email (never sends)")
    p.add_argument("--email-id", required=True)
    p.add_argument("--body", required=True)
    p.add_argument("--subject", default="")

    args = ap.parse_args(argv)
    client = FastmailClient()
    if args.cmd == "search":
        out = client.search_email(query=args.query, from_=args.from_,
                                  subject=args.subject, days=args.days, limit=args.limit)
    elif args.cmd == "thread":
        out = client.get_thread(args.id, with_bodies=not args.no_bodies)
    elif args.cmd == "events":
        out = client.search_events(days=args.days, past_days=args.past_days, query=args.query)
    elif args.cmd == "draft":
        out = client.create_draft(args.to, args.subject, args.body)
    else:  # reply
        out = client.create_draft([], args.subject, args.body,
                                  reply_to_email_id=args.email_id)
    print(json.dumps(out, indent=2))
    return 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except FastmailError as e:
        print(f"error: {e}", file=sys.stderr)
        sys.exit(1)