"""Gmail (send/read/search/reply) in-process MCP for Jarvis — MULTI-ACCOUNT.

Gives Claude real Gmail hands across more than one Google account: Ahmed's
work mailbox (alrugaibfurniture.com) and his business Gmail, chosen per call
by a spoken label ("work" / "business"). "Deal with my work email" and
"reply to that on my business account" route to the right mailbox.

Same OAuth mechanism as ``gcal_control`` — the SAME desktop OAuth client
(``google_oauth.json`` at the project root), just with Gmail scopes and a
SEPARATE cached token PER ACCOUNT (``google_gmail_<label>.json``). The first
action on an account opens a browser once; sign in with THAT account (work vs
business) and the token caches + refreshes silently forever after. No API key,
no per-call billing.

Runs IN-PROCESS via the Agent SDK MCP server (key ``"gmail"``; Claude sees
``mcp__gmail__<name>``). Google client libs are imported lazily so importing
this module never hard-fails when they're missing.

Config (``.env`` / env):
  GMAIL_ACCOUNTS    comma-separated labels to offer, default "work,business"
  GMAIL_DEFAULT     label used when a tool call omits ``account``, default
                    the first in GMAIL_ACCOUNTS ("work")
  GOOGLE_OAUTH_JSON shared desktop OAuth client (default google_oauth.json)
Set ``MCP_GMAIL=0`` to disable. Enabling ALSO requires the Gmail API turned on
in the same Google Cloud project as Calendar and the ``gmail.*`` scopes added
to that project's OAuth consent screen (see README).
"""

from __future__ import annotations

import base64
import os
from email.message import EmailMessage
from email.utils import getaddresses
from pathlib import Path

_ROOT = Path(__file__).resolve().parent.parent
# modify = read + label/trash (not permanent-delete); send = send mail.
# Together: read, search, organise, send and reply. Both are sensitive/
# restricted scopes — fine for personal use as the project's own test user.
_SCOPES = [
    "https://www.googleapis.com/auth/gmail.modify",
    "https://www.googleapis.com/auth/gmail.send",
]

# label -> cached (service, email_address). Built on first use per account.
_SVC: dict[str, tuple] = {}


def _labels() -> list[str]:
    raw = os.environ.get("GMAIL_ACCOUNTS", "work,business")
    return [l.strip().lower() for l in raw.split(",") if l.strip()]


def _default_label() -> str:
    d = os.environ.get("GMAIL_DEFAULT", "").strip().lower()
    labels = _labels()
    return d if d in labels else (labels[0] if labels else "work")


def _oauth_path() -> Path:
    return _ROOT / os.environ.get("GOOGLE_OAUTH_JSON", "google_oauth.json")


def _token_path(label: str) -> Path:
    return _ROOT / f"google_gmail_{label}.json"


def enabled() -> bool:
    """Whether the Gmail MCP should be exposed to Claude. On when not disabled
    and the shared OAuth client exists (individual accounts consent on first
    use)."""
    if os.environ.get("MCP_GMAIL", "1") == "0":
        return False
    return _oauth_path().exists()


def _text(msg: str) -> dict:
    """Wrap a plain string as an MCP text content result."""
    return {"content": [{"type": "text", "text": msg}]}


def _resolve(label: str | None) -> str:
    """Normalise a spoken account label to a configured one (default if
    unknown/empty). 'my work'/'work email' -> 'work'."""
    if not label:
        return _default_label()
    l = label.strip().lower()
    for cfg in _labels():
        if cfg == l or cfg in l or l in cfg:
            return cfg
    return _default_label()


def _service(label: str):
    """Return (gmail_service, email_address) for one account, running the
    one-time browser consent if that account has no cached token. Blocking —
    call via asyncio.to_thread so the voice loop never stalls."""
    if label in _SVC:
        return _SVC[label]
    from google.oauth2.credentials import Credentials
    from google_auth_oauthlib.flow import InstalledAppFlow
    from google.auth.transport.requests import Request
    from googleapiclient.discovery import build

    token_path = _token_path(label)
    creds = None
    if token_path.exists():
        creds = Credentials.from_authorized_user_file(str(token_path), _SCOPES)
    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(Request())
        else:
            print(f"  [gmail] first-time sign-in for the '{label}' account — a "
                  f"browser will open; sign in with the {label.upper()} Google "
                  f"account and approve Jarvis.")
            flow = InstalledAppFlow.from_client_secrets_file(
                str(_oauth_path()), _SCOPES)
            creds = flow.run_local_server(port=0)
        token_path.write_text(creds.to_json())
    svc = build("gmail", "v1", credentials=creds, cache_discovery=False)
    addr = svc.users().getProfile(userId="me").execute().get(
        "emailAddress", label)
    _SVC[label] = (svc, addr)
    return _SVC[label]


# ---------------------------------------------------------------------------
# blocking helpers (run via asyncio.to_thread)
# ---------------------------------------------------------------------------

def _b64(msg: EmailMessage) -> str:
    return base64.urlsafe_b64encode(msg.as_bytes()).decode()


def _send(label: str, to: str, subject: str, body: str, cc: str) -> str:
    svc, addr = _service(label)
    msg = EmailMessage()
    msg["To"] = to
    if cc:
        msg["Cc"] = cc
    msg["Subject"] = subject
    msg.set_content(body)
    svc.users().messages().send(
        userId="me", body={"raw": _b64(msg)}).execute()
    return (f"Sent '{subject}' to {to}" + (f" (cc {cc})" if cc else "")
            + f" from your {label} account ({addr}).")


def _hdr(headers: list, name: str) -> str:
    for h in headers:
        if h.get("name", "").lower() == name.lower():
            return h.get("value", "")
    return ""


def _search(label: str, query: str, limit: int) -> str:
    svc, addr = _service(label)
    res = svc.users().messages().list(
        userId="me", q=query or "in:inbox", maxResults=limit).execute()
    ids = [m["id"] for m in res.get("messages", [])]
    if not ids:
        return f"No matching messages in the {label} account ({addr})."
    lines = [f"{label} account ({addr}):"]
    for mid in ids:
        m = svc.users().messages().get(
            userId="me", id=mid, format="metadata",
            metadataHeaders=["From", "Subject", "Date"]).execute()
        h = m.get("payload", {}).get("headers", [])
        snippet = m.get("snippet", "")[:100]
        lines.append(
            f"[id {mid}] {_hdr(h, 'Date')}\n  from: {_hdr(h, 'From')}\n"
            f"  subj: {_hdr(h, 'Subject')}\n  {snippet}")
    return "\n".join(lines)


def _walk_body(payload: dict) -> str:
    """Depth-first search for a text/plain part; fall back to any text/html."""
    stack = [payload]
    html = ""
    while stack:
        p = stack.pop(0)
        mime = p.get("mimeType", "")
        data = p.get("body", {}).get("data")
        if mime == "text/plain" and data:
            return base64.urlsafe_b64decode(data).decode("utf-8", "replace")
        if mime == "text/html" and data and not html:
            html = base64.urlsafe_b64decode(data).decode("utf-8", "replace")
        stack.extend(p.get("parts", []) or [])
    if html:
        import re
        return re.sub(r"<[^>]+>", "", html)  # crude strip; good enough for TTS
    return "(no readable body)"


def _read(label: str, mid: str) -> str:
    svc, addr = _service(label)
    m = svc.users().messages().get(userId="me", id=mid, format="full").execute()
    h = m.get("payload", {}).get("headers", [])
    body = _walk_body(m.get("payload", {}))
    if len(body) > 4000:
        body = body[:4000] + "\n…(truncated)"
    return (f"From: {_hdr(h, 'From')}\nTo: {_hdr(h, 'To')}\n"
            f"Subject: {_hdr(h, 'Subject')}\nDate: {_hdr(h, 'Date')}\n\n{body}")


def _reply(label: str, mid: str, body: str, reply_all: bool = False) -> str:
    """Reply in-thread to a message id, preserving threading headers.
    reply_all=True keeps every other recipient (original To + Cc, minus us)."""
    svc, addr = _service(label)
    orig = svc.users().messages().get(
        userId="me", id=mid, format="metadata",
        metadataHeaders=["From", "To", "Cc", "Subject", "Message-ID",
                         "References"]).execute()
    h = orig.get("payload", {}).get("headers", [])
    to = _hdr(h, "From")
    subj = _hdr(h, "Subject")
    if not subj.lower().startswith("re:"):
        subj = "Re: " + subj
    msg_id = _hdr(h, "Message-ID")
    refs = (_hdr(h, "References") + " " + msg_id).strip()
    msg = EmailMessage()
    msg["To"] = to
    msg["Subject"] = subj
    cc = ""
    if reply_all:
        # everyone else on the thread, minus our own address and the sender
        # (already in To) — de-duped, order preserved.
        seen = {addr.lower()} | {a.lower() for _, a in
                                 getaddresses([to]) if a}
        keep: list[str] = []
        for _n, a in getaddresses([_hdr(h, "To"), _hdr(h, "Cc")]):
            if a and a.lower() not in seen:
                seen.add(a.lower())
                keep.append(a)
        cc = ", ".join(keep)
        if cc:
            msg["Cc"] = cc
    if msg_id:
        msg["In-Reply-To"] = msg_id
        msg["References"] = refs
    msg.set_content(body)
    svc.users().messages().send(
        userId="me",
        body={"raw": _b64(msg), "threadId": orig.get("threadId")}).execute()
    return (f"Replied to {to}" + (f" (cc {cc})" if cc else "")
            + f" on your {label} account ({addr}).")


# ---------------------------------------------------------------------------
# HUD email card — full threads, untruncated bodies, real attachments
# (voice/mail_view.py turns these into the `email_thread` event). Everything
# above stays as it was: those feed Jarvis' VOICE, these feed his SCREEN.
# ---------------------------------------------------------------------------

_INLINE_MAX = 4 * 1024 * 1024   # don't inline a "cid:" image bigger than this


def _b64d(data: str) -> bytes:
    """Decode Gmail's base64url part data (it comes back unpadded)."""
    s = (data or "").replace("-", "+").replace("_", "/")
    return base64.b64decode(s + "=" * (-len(s) % 4))


def _walk_parts(payload: dict):
    """Yield every MIME part of a message, depth-first (payload included)."""
    stack = [payload]
    while stack:
        p = stack.pop(0)
        yield p
        stack.extend(p.get("parts", []) or [])


def _cid(part: dict) -> str:
    """The part's Content-ID with the <angle brackets> stripped."""
    v = _hdr(part.get("headers", []) or [], "Content-ID").strip()
    return v[1:-1] if v.startswith("<") and v.endswith(">") else v


def _parts_of(svc, m: dict) -> dict:
    """Split one full-format message into its text body, HTML body, inline
    (cid:) images and real attachments. Bodies are NOT truncated — the card
    shows the whole thing, unlike ``_read`` which is a TTS summary."""
    text, html = "", ""
    inline: dict[str, dict] = {}
    atts: list[dict] = []
    for p in _walk_parts(m.get("payload", {}) or {}):
        mime = p.get("mimeType", "") or ""
        if mime.startswith("multipart/"):
            continue
        body = p.get("body", {}) or {}
        data = body.get("data")
        fname = p.get("filename") or ""
        cid = _cid(p)
        if cid and mime.startswith("image/"):
            raw = data
            if not raw and body.get("attachmentId") and \
                    int(body.get("size") or 0) <= _INLINE_MAX:
                raw = svc.users().messages().attachments().get(
                    userId="me", messageId=m.get("id", ""),
                    id=body["attachmentId"]).execute().get("data")
            if raw:  # data: URIs want STANDARD base64, not Gmail's url-safe
                inline[cid] = {"mime": mime,
                               "b64": base64.b64encode(_b64d(raw)).decode()}
            continue
        if fname:
            atts.append({"att_id": body.get("attachmentId", "") or "",
                         "filename": fname, "mime": mime,
                         "size": int(body.get("size") or 0), "path": ""})
            continue
        if mime == "text/plain" and data:
            text += _b64d(data).decode("utf-8", "replace")
        elif mime == "text/html" and data:
            html += _b64d(data).decode("utf-8", "replace")
    return {"body_text": text, "body_html": html, "inline": inline,
            "attachments": atts}


def search_ids(label: str, query: str, limit: int = 5) -> list[str]:
    """Message ids matching a Gmail query, newest first (empty = inbox)."""
    svc, _addr = _service(label)
    res = svc.users().messages().list(
        userId="me", q=query or "in:inbox", maxResults=limit).execute()
    return [m["id"] for m in res.get("messages", []) or []]


def thread(label: str, ident: str) -> dict:
    """A whole THREAD as structured data for the HUD card. `ident` is a thread
    id OR any message id inside it (Gmail ids are interchangeable-looking, so
    a message id is resolved to its threadId on the fly). Oldest first."""
    svc, addr = _service(label)
    try:
        t = svc.users().threads().get(
            userId="me", id=ident, format="full").execute()
    except Exception:  # noqa: BLE001 — it was a message id, not a thread id
        m = svc.users().messages().get(
            userId="me", id=ident, format="minimal").execute()
        t = svc.users().threads().get(
            userId="me", id=m.get("threadId", ident), format="full").execute()
    msgs: list[dict] = []
    for m in t.get("messages", []) or []:
        h = m.get("payload", {}).get("headers", []) or []
        msgs.append({
            "msg_id": m.get("id", ""),
            "from": _hdr(h, "From"), "to": _hdr(h, "To"), "cc": _hdr(h, "Cc"),
            "subject": _hdr(h, "Subject"), "date": _hdr(h, "Date"),
            "ts": int(m.get("internalDate") or 0) / 1000.0,
            # SENT is Gmail's own word for "Ahmed wrote this" — authoritative
            "outgoing": "SENT" in (m.get("labelIds") or []),
            **_parts_of(svc, m),
        })
    msgs.sort(key=lambda x: x["ts"])
    subject = next((m["subject"] for m in msgs if m["subject"]), "(no subject)")
    return {"thread_id": t.get("id", ident), "account_addr": addr,
            "subject": subject, "messages": msgs}


def attachment(label: str, msg_id: str, att_id: str, filename: str,
               thread_id: str = "") -> str:
    """Download one attachment to ~/.jarvis/attachments/<thread>/ and return
    its local path. The filename is sanitised (no path traversal)."""
    from voice import mail_view  # lazy: mail_view imports us back

    svc, _addr = _service(label)
    a = svc.users().messages().attachments().get(
        userId="me", messageId=msg_id, id=att_id).execute()
    path = mail_view.dest_path(thread_id or msg_id, filename)
    path.write_bytes(_b64d(a.get("data", "")))
    return str(path)


def build_server():
    """Build and return the in-process ``gmail`` MCP server."""
    import asyncio

    from claude_agent_sdk import tool, create_sdk_mcp_server

    accts = ", ".join(_labels()) or "work"
    acct_help = (f" `account` picks the mailbox ({accts}); omit it for "
                 f"'{_default_label()}'.")

    @tool("gmail_send",
          "Send an email from one of Ahmed's Gmail accounts." + acct_help
          + " `to`/`cc` are comma-separated addresses; `cc` optional.",
          {"to": str, "subject": str, "body": str, "cc": str, "account": str})
    async def gmail_send(args: dict) -> dict:
        to = str(args.get("to", "")).strip()
        if not to:
            return _text("gmail_send failed: no recipient.")
        try:
            out = await asyncio.to_thread(
                _send, _resolve(args.get("account")), to,
                str(args.get("subject", "")), str(args.get("body", "")),
                str(args.get("cc", "")).strip())
            return _text(out)
        except Exception as e:  # noqa: BLE001
            return _text(f"gmail_send failed: {e}")

    @tool("gmail_search",
          "Search a Gmail account using Gmail search syntax (e.g. "
          "'from:elie is:unread', 'subject:invoice newer_than:7d', empty = "
          "inbox). Returns a list with a message id for each; pass an id to "
          "gmail_read." + acct_help,
          {"query": str, "limit": int, "account": str})
    async def gmail_search(args: dict) -> dict:
        try:
            limit = max(1, min(int(args.get("limit") or 10), 25))
            out = await asyncio.to_thread(
                _search, _resolve(args.get("account")),
                str(args.get("query", "")), limit)
            return _text(out)
        except Exception as e:  # noqa: BLE001
            return _text(f"gmail_search failed: {e}")

    @tool("gmail_read",
          "Read the full body of one message by its id (from gmail_search)."
          + acct_help, {"id": str, "account": str})
    async def gmail_read(args: dict) -> dict:
        mid = str(args.get("id", "")).strip()
        if not mid:
            return _text("gmail_read failed: no message id.")
        try:
            out = await asyncio.to_thread(
                _read, _resolve(args.get("account")), mid)
            return _text(out)
        except Exception as e:  # noqa: BLE001
            return _text(f"gmail_read failed: {e}")

    @tool("gmail_reply",
          "Reply in-thread to a message by its id (from gmail_search), keeping "
          "the subject and conversation threading. Goes to the original "
          "sender." + acct_help, {"id": str, "body": str, "account": str})
    async def gmail_reply(args: dict) -> dict:
        mid = str(args.get("id", "")).strip()
        if not mid:
            return _text("gmail_reply failed: no message id.")
        try:
            out = await asyncio.to_thread(
                _reply, _resolve(args.get("account")), mid,
                str(args.get("body", "")))
            return _text(out)
        except Exception as e:  # noqa: BLE001
            return _text(f"gmail_reply failed: {e}")

    @tool("gmail_accounts",
          "List the Gmail accounts Jarvis can use and which one is the "
          "default.", {})
    async def gmail_accounts(args: dict) -> dict:
        lines = []
        for l in _labels():
            authed = _token_path(l).exists()
            mark = "default" if l == _default_label() else ""
            state = "signed in" if authed else "not signed in yet"
            lines.append(f"- {l} ({state}){' — ' + mark if mark else ''}")
        return _text("\n".join(lines) or "No accounts configured.")

    return create_sdk_mcp_server(
        name="gmail",
        version="1.0.0",
        tools=[gmail_send, gmail_search, gmail_read, gmail_reply,
               gmail_accounts],
    )