"""
MCP Email Server - MCP Streamable HTTP compatible with Open WebUI

Features:
- IMAP (login) + SMTP (login) for multiple accounts
- List/search all folders (including shared/public)
- New-message tracking and flagging
- Triage guidance (business description + personnel/chain-of-command)
- Extended tools: reply, forward, drafts, move, copy, labels, attachments, schedule_send, export_conversation
- Conflict check search
- Exposes MCP tools via Streamable HTTP (JSON-RPC 2.0)

Run:
- uvicorn server:app --host 0.0.0.0 --port 8000

Logging:
- All logs are written to stdout (and stderr) so Docker captures them.
- Set LOG_LEVEL=DEBUG for full trace logs.
"""

import json
import os
import sys
import time
import base64
import email
import email.policy
import imaplib
import smtplib
import logging
import threading
import uuid
from contextvars import ContextVar
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse

# Correlation ID support
request_id_var: ContextVar[str] = ContextVar("request_id", default="")


def get_request_id() -> str:
    return request_id_var.get() or ""


def log_context_prefix() -> str:
    rid = get_request_id()
    if rid:
        return f"[req:{rid}] "
    return ""


# -------------------- LOGGING (Docker-friendly) --------------------

log_level = os.getenv("LOG_LEVEL", "INFO").upper()

root_logger = logging.getLogger()
root_logger.handlers.clear()
root_logger.setLevel(logging.DEBUG)

logger = logging.getLogger("mcp-email-server")
logger.setLevel(getattr(logging, log_level, logging.INFO))

handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.DEBUG)
formatter = logging.Formatter(
    "%(asctime)s %(levelname)s %(name)s [%(threadName)s] %(message)s",
    datefmt="%Y-%m-%dT%H:%M:%S%z",
)
handler.setFormatter(formatter)
logger.addHandler(handler)


# -------------------- CONFIG (multi-account) --------------------

@dataclass
class AccountConfig:
    id: str
    name: str
    # IMAP
    imap_host: str
    imap_port: int
    imap_use_ssl: bool
    imap_username: str
    imap_password: str
    # SMTP
    smtp_host: str
    smtp_port: int
    smtp_use_tls: bool
    smtp_use_ssl: bool
    smtp_username: str
    smtp_password: str
    smtp_from: str


@dataclass
class GlobalConfig:
    # Accounts
    accounts: Dict[str, AccountConfig] = field(default_factory=dict)
    default_account_id: str = ""

    # Optional cc/bcc defaults (global)
    DEFAULT_CC: str = os.getenv("DEFAULT_CC", "")
    DEFAULT_BCC: str = os.getenv("DEFAULT_BCC", "")

    # Business/person description
    BUSINESS_DESCRIPTION: str = os.getenv("BUSINESS_DESCRIPTION", "")

    # Personnel / chain-of-command (JSON)
    PERSONNEL_JSON: str = os.getenv("PERSONNEL_JSON", "[]")

    # Scheduling
    SCHEDULED_SEND_INTERVAL: int = int(os.getenv("SCHEDULED_SEND_INTERVAL", "10"))

    # HTML signature (optional)
    # Can be:
    # - Raw HTML string
    # - File path: if the value is an existing file path, its content is used as the signature.
    EMAIL_HTML_SIGNATURE_RAW: str = os.getenv("EMAIL_HTML_SIGNATURE", "")

    EMAIL_HTML_SIGNATURE: str = ""

    @classmethod
    def resolve_html_signature(cls, raw: str) -> str:
        if not raw:
            return ""
        # If it looks like a path and file exists, treat as template path
        try:
            if os.path.isfile(raw):
                with open(raw, "r", encoding="utf-8") as f:
                    return f.read().strip()
        except Exception:
            pass
        # Otherwise, treat as inline HTML
        return raw.strip()

    @property
    def personnel(self) -> List[Dict[str, Any]]:
        try:
            return json.loads(self.PERSONNEL_JSON)
        except Exception:
            return []

    @property
    def default_cc_list(self) -> List[str]:
        return [e.strip() for e in self.DEFAULT_CC.split(",") if e.strip()]

    @property
    def default_bcc_list(self) -> List[str]:
        return [e.strip() for e in self.DEFAULT_BCC.split(",") if e.strip()]


def load_accounts_config() -> Dict[str, AccountConfig]:
    """
    Load accounts from:
      - ACCOUNTS_CONFIG_PATH (JSON file), or
      - ACCOUNTS_JSON (inline JSON), or
      - fallback to single account from existing env vars as id "default".
    """
    path = os.getenv("ACCOUNTS_CONFIG_PATH", "").strip()
    inline = os.getenv("ACCOUNTS_JSON", "").strip()

    raw = None

    if path:
        try:
            with open(path, "r", encoding="utf-8") as f:
                raw = json.loads(f.read())
        except Exception as e:
            logger.error(f"Error loading ACCOUNTS_CONFIG_PATH={path}: {e}")
    elif inline:
        try:
            raw = json.loads(inline)
        except Exception as e:
            logger.error(f"Error parsing ACCOUNTS_JSON: {e}")

    if isinstance(raw, list):
        accounts: Dict[str, AccountConfig] = {}
        for a in raw:
            acc = parse_account(a)
            accounts[acc.id] = acc
        return accounts

    if isinstance(raw, dict) and "accounts" in raw:
        accounts: Dict[str, AccountConfig] = {}
        for a in raw["accounts"]:
            acc = parse_account(a)
            accounts[acc.id] = acc
        return accounts

    # Fallback: single account from environment variables
    return {
        "default": build_fallback_account()
    }


def parse_account(a: Dict[str, Any]) -> AccountConfig:
    # Allow flexible fields; fall back to sensible defaults
    def boolify(v):
        return str(v).lower() in ("true", "1", "yes")

    return AccountConfig(
        id=str(a.get("id") or a.get("account_id") or "default"),
        name=str(a.get("name") or a.get("id") or a.get("account_id") or ""),
        imap_host=str(a.get("imap_host") or a.get("IMAP_HOST") or "imap.example.com"),
        imap_port=int(a.get("imap_port") or a.get("IMAP_PORT") or 993),
        imap_use_ssl=boolify(a.get("imap_use_ssl") or a.get("IMAP_USE_SSL") or True),
        imap_username=str(a.get("imap_username") or a.get("IMAP_USERNAME") or ""),
        imap_password=str(a.get("imap_password") or a.get("IMAP_PASSWORD") or ""),
        smtp_host=str(a.get("smtp_host") or a.get("SMTP_HOST") or "smtp.example.com"),
        smtp_port=int(a.get("smtp_port") or a.get("SMTP_PORT") or 587),
        smtp_use_tls=boolify(a.get("smtp_use_tls") or a.get("SMTP_USE_TLS") or True),
        smtp_use_ssl=boolify(a.get("smtp_use_ssl") or a.get("SMTP_USE_SSL") or False),
        smtp_username=str(a.get("smtp_username") or a.get("SMTP_USERNAME") or ""),
        smtp_password=str(a.get("smtp_password") or a.get("SMTP_PASSWORD") or ""),
        smtp_from=str(a.get("smtp_from") or a.get("SMTP_FROM") or ""),
    )


def build_fallback_account() -> AccountConfig:
    return AccountConfig(
        id="default",
        name="default",
        imap_host=os.getenv("IMAP_HOST", "imap.example.com"),
        imap_port=int(os.getenv("IMAP_PORT", "993")),
        imap_use_ssl=os.getenv("IMAP_USE_SSL", "true").lower() in ("true", "1", "yes"),
        imap_username=os.getenv("IMAP_USERNAME", ""),
        imap_password=os.getenv("IMAP_PASSWORD", ""),
        smtp_host=os.getenv("SMTP_HOST", "smtp.example.com"),
        smtp_port=int(os.getenv("SMTP_PORT", "587")),
        smtp_use_tls=os.getenv("SMTP_USE_TLS", "true").lower() in ("true", "1", "yes"),
        smtp_use_ssl=os.getenv("SMTP_USE_SSL", "false").lower() in ("true", "1", "yes"),
        smtp_username=os.getenv("SMTP_USERNAME", ""),
        smtp_password=os.getenv("SMTP_PASSWORD", ""),
        smtp_from=os.getenv("SMTP_FROM", ""),
    )


def resolve_default_account_id(accounts: Dict[str, AccountConfig]) -> str:
    explicit = os.getenv("DEFAULT_ACCOUNT", "").strip()
    if explicit and explicit in accounts:
        return explicit
    if len(accounts) == 1:
        return next(iter(accounts))
    # Prefer "default" if present
    if "default" in accounts:
        return "default"
    # Fallback: first
    return next(iter(accounts))


# Build global config
accounts_map = load_accounts_config()
default_account_id = resolve_default_account_id(accounts_map)

config = GlobalConfig(
    accounts=accounts_map,
    default_account_id=default_account_id,
)

# Resolve EMAIL_HTML_SIGNATURE (inline or file path)
config.EMAIL_HTML_SIGNATURE = GlobalConfig.resolve_html_signature(config.EMAIL_HTML_SIGNATURE_RAW)

logger.info(f"Loaded {len(accounts_map)} account(s): {list(accounts_map.keys())}. Default: {default_account_id}")


def get_account(account_id: Optional[str] = None) -> AccountConfig:
    aid = (account_id or config.default_account_id).strip()
    if not aid:
        aid = config.default_account_id
    if aid not in config.accounts:
        raise ValueError(f"Account not found: {aid}. Available: {list(config.accounts.keys())}")
    return config.accounts[aid]


# -------------------- IMAP HELPERS (UID-based, per-account) --------------------

def create_imap(acc: AccountConfig):
    prefix = log_context_prefix()
    logger.info(f"{prefix}IMAP create connection: {acc.imap_host}:{acc.imap_port} ssl={acc.imap_use_ssl} account={acc.id}")
    if acc.imap_use_ssl:
        m = imaplib.IMAP4_SSL(acc.imap_host, acc.imap_port)
    else:
        m = imaplib.IMAP4(acc.imap_host, acc.imap_port)
        m.starttls()
    return m


def login_imap(imap, acc: AccountConfig):
    prefix = log_context_prefix()
    logger.info(f"{prefix}IMAP login user={acc.imap_username} account={acc.id}")
    try:
        imap.login(acc.imap_username, acc.imap_password)
        logger.debug(f"{prefix}IMAP login success user={acc.imap_username} account={acc.id}")
    except Exception as e:
        logger.error(f"{prefix}IMAP login failed user={acc.imap_username} account={acc.id}: {e}")
        raise


def list_all_folders(imap):
    _, data = imap.list()
    folders = []
    for line in data or []:
        if isinstance(line, bytes):
            line = line.decode("utf-8", errors="replace")
        parts = line.split('"')
        if len(parts) >= 4:
            name = parts[3]
            folders.append({"path": name})
    return folders


def ensure_selected(imap, folder: str):
    prefix = log_context_prefix()
    logger.debug(f"{prefix}IMAP select folder={folder}")
    imap.select(folder, readonly=False)


def parse_message(imap, uid: int):
    prefix = log_context_prefix()
    logger.debug(f"{prefix}IMAP fetch message uid={uid}")
    status, msg_data = imap.uid("FETCH", str(uid), "(RFC822)")
    if status != "OK" or not msg_data or msg_data[0] is None:
        logger.error(f"{prefix}IMAP fetch failed for uid={uid}")
        raise RuntimeError(f"Failed to fetch message UID {uid}")
    msg = email.message_from_bytes(msg_data[0][1], policy=email.policy.default)

    subject = msg.get("Subject", "") or ""
    from_header = msg.get("From", "") or ""
    to_header = msg.get("To", "") or ""
    cc_header = msg.get("Cc", "") or ""
    date_header = msg.get("Date", "") or ""
    message_id = msg.get("Message-ID", "") or ""
    in_reply_to = msg.get("In-Reply-To", "") or ""
    references = msg.get("References", "") or ""

    body_plain = ""
    body_html = ""

    if msg.is_multipart():
        for part in msg.walk():
            ptype = part.get_content_type()
            pdisp = (part.get("Content-Disposition") or "").lower()
            if "attachment" in pdisp:
                continue
            if ptype == "text/plain":
                body_plain += part.get_payload(decode=True).decode(
                    part.get_content_charset() or "utf-8", errors="replace"
                )
            elif ptype == "text/html":
                body_html += part.get_payload(decode=True).decode(
                    part.get_content_charset() or "utf-8", errors="replace"
                )
    else:
        ptype = msg.get_content_type()
        if ptype in ("text/plain", "text/html"):
            payload = msg.get_payload(decode=True).decode(
                msg.get_content_charset() or "utf-8", errors="replace"
            )
            if ptype == "text/plain":
                body_plain = payload
            else:
                body_html = payload

    _, flags_data = imap.uid("FETCH", str(uid), "(FLAGS)")
    flags = []
    if flags_data:
        line = flags_data[0].decode("utf-8", errors="replace")
        start = line.find("(")
        end = line.find(")")
        if start != -1 and end != -1:
            flags = [f.strip() for f in line[start+1:end].split() if f.strip()]

    return {
        "uid": uid,
        "subject": subject,
        "from": from_header,
        "to": to_header,
        "cc": cc_header,
        "date": date_header,
        "body_plain": body_plain,
        "body_html": body_html,
        "flags": flags,
        "message_id": message_id,
        "in_reply_to": in_reply_to,
        "references": references,
    }


def search_messages(imap, folder: str, query: str, since: Optional[str], max_results: int):
    prefix = log_context_prefix()
    logger.info(f"{prefix}IMAP search folder={folder} query={query} since={since} max_results={max_results}")
    ensure_selected(imap, folder)

    criteria = []
    if query:
        criteria.append(f'SUBJECT "{query}"')
    else:
        criteria.append("ALL")

    if since:
        criteria.append(f"SINCE {since}")

    search_str = " ".join(criteria)
    logger.debug(f"{prefix}IMAP search criteria: {search_str}")
    status, data = imap.uid("SEARCH", None, search_str)
    if status != "OK" or not data or not data[0]:
        logger.debug(f"{prefix}IMAP search returned no results for folder={folder}")
        return []

    ids = (data[0].split() if data[0] else [])
    ids = ids[-max_results:]

    results = []
    for uid_bytes in ids:
        uid = int(uid_bytes)
        try:
            m = parse_message(imap, uid)
            results.append(m)
        except Exception as e:
            logger.error(f"Error parsing message UID {uid}: {e}")
    return results


def mark_as_read(imap, folder: str, uid: int):
    ensure_selected(imap, folder)
    imap.uid("STORE", str(uid), "+FLAGS", "\\Seen")


def flag_message(imap, folder: str, uid: int):
    ensure_selected(imap, folder)
    imap.uid("STORE", str(uid), "+FLAGS", "\\Flagged")


def move_message(imap, folder: str, uid: int, dest_folder: str):
    ensure_selected(imap, folder)
    imap.uid("COPY", str(uid), dest_folder)
    imap.uid("STORE", str(uid), "+FLAGS", "\\Deleted")
    imap.expunge()


def copy_message(imap, folder: str, uid: int, dest_folder: str):
    ensure_selected(imap, folder)
    imap.uid("COPY", str(uid), dest_folder)


def add_flags(imap, folder: str, uid: int, flags: List[str]):
    ensure_selected(imap, folder)
    imap.uid("STORE", str(uid), "+FLAGS", tuple(flags))


def remove_flags(imap, folder: str, uid: int, flags: List[str]):
    ensure_selected(imap, folder)
    imap.uid("STORE", str(uid), "-FLAGS", tuple(flags))


def get_labels(imap, folder: str, uid: int):
    _, flags_data = imap.uid("FETCH", str(uid), "(FLAGS)")
    if not flags_data:
        return []
    line = flags_data[0].decode("utf-8", errors="replace")
    start = line.find("(")
    end = line.find(")")
    if start == -1 or end == -1:
        return []
    flags = [f.strip() for f in line[start+1:end].split() if f.strip()]
    return [f for f in flags if f.startswith("$Label_")]


def list_available_labels(imap):
    labels = set()
    try:
        folders = list_all_folders(imap)
        for f in folders[:20]:
            path = f["path"]
            try:
                ensure_selected(imap, path)
                status, data = imap.uid("SEARCH", None, "ALL")
                if status != "OK" or not data or not data[0]:
                    continue
                ids = data[0].split()
                for uid_bytes in ids[:20]:
                    uid = int(uid_bytes)
                    _, flags_data = imap.uid("FETCH", str(uid), "(FLAGS)")
                    if not flags_data:
                        continue
                    line = flags_data[0].decode("utf-8", errors="replace")
                    start = line.find("(")
                    end = line.find(")")
                    if start == -1 or end == -1:
                        continue
                    flags = [f.strip() for f in line[start+1:end].split() if f.strip()]
                    for fl in flags:
                        if fl.startswith("$Label_"):
                            labels.add(fl.replace("$Label_", "", 1))
            except Exception:
                continue
    except Exception:
        pass
    return sorted(labels)


# -------------------- SMTP HELPERS --------------------

def send_email(
    acc: AccountConfig,
    to: List[str],
    subject: str,
    body: str,
    cc: Optional[List[str]] = None,
    bcc: Optional[List[str]] = None,
    html: bool = False,
    in_reply_to: Optional[str] = None,
    references: Optional[str] = None,
):
    prefix = log_context_prefix()
    if not acc.smtp_from:
        raise ValueError("SMTP_FROM not configured for account")

    if not cc:
        cc = []
    if not bcc:
        bcc = []

    cc = list(set(cc + config.default_cc_list))
    bcc = list(set(bcc + config.default_bcc_list))

    all_recipients = to + cc + bcc
    logger.info(f"{prefix}SMTP send_email from={acc.smtp_from} to={to} cc={cc} bcc={bcc} account={acc.id}")

    from email.mime.text import MIMEText
    from email.mime.multipart import MIMEMultipart

    if html:
        m = MIMEMultipart("alternative")
        m.attach(MIMEText(body, "html"))
    else:
        m = MIMEText(body, "plain")

    m["From"] = acc.smtp_from
    m["To"] = ", ".join(to)
    if cc:
        m["Cc"] = ", ".join(cc)
    m["Subject"] = subject
    domain = acc.smtp_from.split("@")[1] if "@" in acc.smtp_from else "example.com"
    m["Message-ID"] = f"<{uuid.uuid4().hex}@{domain}>"

    if in_reply_to:
        m["In-Reply-To"] = in_reply_to
    if references:
        m["References"] = references

    try:
        if acc.smtp_use_ssl:
            s = smtplib.SMTP_SSL(acc.smtp_host, acc.smtp_port)
        else:
            s = smtplib.SMTP(acc.smtp_host, acc.smtp_port)
            if acc.smtp_use_tls:
                s.starttls()

        logger.info(f"{prefix}SMTP connect host={acc.smtp_host} port={acc.smtp_port} account={acc.id}")
        s.login(acc.smtp_username, acc.smtp_password)
        logger.info(f"{prefix}SMTP login success user={acc.smtp_username} account={acc.id}")
        s.sendmail(acc.smtp_from, all_recipients, m.as_string())
        logger.info(f"{prefix}SMTP send success to={all_recipients} account={acc.id}")
        s.quit()
    except Exception as e:
        logger.error(f"{prefix}SMTP send failed: {e}")
        raise


# -------------------- NEW-MESSAGE TRACKING --------------------

last_seen_uids: Dict[str, set] = {}


def get_new_messages(imap, folder: str):
    ensure_selected(imap, folder)
    status, data = imap.uid("SEARCH", None, "ALL")
    if status != "OK" or not data or not data[0]:
        return []
    ids = data[0].split()
    uids = {int(u) for u in ids}

    prev = last_seen_uids.setdefault(folder, set())
    new_uids = sorted(uids - prev)
    last_seen_uids[folder] = uids

    results = []
    for uid in new_uids:
        m = parse_message(imap, uid)
        results.append(m)
    return results


def sync_seen(imap, folder: str):
    ensure_selected(imap, folder)
    status, data = imap.uid("SEARCH", None, "ALL")
    if status != "OK" or not data or not data[0]:
        last_seen_uids[folder] = set()
        return
    ids = data[0].split()
    last_seen_uids[folder] = {int(u) for u in ids}


# -------------------- TRIAGE HELPERS --------------------

def build_triage_hint(msg: Dict[str, Any]) -> str:
    if not config.BUSINESS_DESCRIPTION and not config.personnel:
        return ""

    hint_parts = []

    if config.BUSINESS_DESCRIPTION:
        hint_parts.append(f"Context: {config.BUSINESS_DESCRIPTION}")

    from email.utils import parseaddr
    sender_email = parseaddr(msg.get("from", ""))[1].lower()
    personnel_by_email = {
        p["email"].lower(): p for p in config.personnel if p.get("email")
    }

    sender_info = personnel_by_email.get(sender_email)
    if sender_info:
        hint_parts.append(
            f"Sender is internal: {sender_info['name']} ({sender_info['role']})."
        )
        if sender_info.get("escalates_to"):
            hint_parts.append(
                f"If escalation needed, escalate to: {sender_info['escalates_to']}."
            )

    subject = (msg.get("subject") or "").lower()
    if any(w in subject for w in ["urgent", "asap", "escalat", "priority"]):
        hint_parts.append("Message appears urgent; consider escalation.")

    return " ".join(hint_parts)


# -------------------- CONFLICT CHECK SEARCH --------------------

def conflict_check_search(imap, terms: List[str], folders: Optional[List[str]] = None, max_results: int = 100):
    prefix = log_context_prefix()
    logger.info(f"{prefix}conflict_check_search terms={terms} folders={folders} max_results={max_results}")
    if not folders:
        folders = ["INBOX"]

    all_results = []
    for folder in folders:
        try:
            ensure_selected(imap, folder)
        except Exception:
            logger.warning(f"{prefix}conflict_check_search cannot select folder={folder}")
            continue

        for term in terms:
            try:
                status, data = imap.uid("SEARCH", None, f'SUBJECT "{term}"')
                if status != "OK" or not data or not data[0]:
                    continue
                ids = data[0].split()
                for uid_bytes in ids[-max_results:]:
                    uid = int(uid_bytes)
                    msg = parse_message(imap, uid)
                    all_results.append({
                        "folder": folder,
                        "uid": uid,
                        "subject": msg["subject"],
                        "from": msg["from"],
                        "to": msg["to"],
                        "cc": msg["cc"],
                        "date": msg["date"],
                        "matched_term": term,
                        "snippet": (msg["body_plain"] or "")[:400],
                    })
            except Exception as e:
                logger.error(f"{prefix}Conflict search error for term '{term}' in {folder}: {e}")

    seen = set()
    unique = []
    for r in all_results:
        key = (r["folder"], r["uid"])
        if key not in seen:
            seen.add(key)
            unique.append(r)
    return unique


# -------------------- ATTACHMENTS --------------------

def list_attachments(imap, folder: str, uid: int) -> List[Dict[str, Any]]:
    ensure_selected(imap, folder)
    status, msg_data = imap.uid("FETCH", str(uid), "(RFC822)")
    if status != "OK" or not msg_data or msg_data[0] is None:
        return []
    msg = email.message_from_bytes(msg_data[0][1], policy=email.policy.default)

    attachments = []
    for part in msg.walk():
        pdisp = (part.get("Content-Disposition") or "").lower()
        if "attachment" not in pdisp:
            continue
        filename = part.get_filename() or "attachment"
        size = len(part.get_payload(decode=True) or b"")
        attachments.append({
            "filename": filename,
            "content_type": part.get_content_type(),
            "size": size,
        })
    return attachments


def download_attachment(imap, folder: str, uid: int, filename: str) -> Dict[str, Any]:
    ensure_selected(imap, folder)
    status, msg_data = imap.uid("FETCH", str(uid), "(RFC822)")
    if status != "OK" or not msg_data or msg_data[0] is None:
        raise ValueError("Failed to fetch message for attachment")
    msg = email.message_from_bytes(msg_data[0][1], policy=email.policy.default)

    for part in msg.walk():
        pdisp = (part.get("Content-Disposition") or "").lower()
        if "attachment" not in pdisp:
            continue
        part_filename = part.get_filename() or "attachment"
        if part_filename != filename:
            continue
        payload = part.get_payload(decode=True) or b""
        b64 = base64.b64encode(payload).decode("ascii")
        return {
            "filename": part_filename,
            "content_type": part.get_content_type(),
            "size": len(payload),
            "content_base64": b64,
        }

    raise ValueError(f"Attachment not found: {filename}")


def search_attachments(imap, folder: str, file_pattern: str, max_results: int):
    ensure_selected(imap, folder)
    status, data = imap.uid("SEARCH", None, "ALL")
    if status != "OK" or not data or not data[0]:
        return []
    ids = data[0].split()
    ids = ids[-max_results * 5:]

    pattern = file_pattern.lower()
    results = []
    for uid_bytes in ids:
        if len(results) >= max_results:
            break
        uid = int(uid_bytes)
        try:
            status, msg_data = imap.uid("FETCH", str(uid), "(RFC822)")
            if status != "OK" or not msg_data or msg_data[0] is None:
                continue
            msg = email.message_from_bytes(msg_data[0][1], policy=email.policy.default)
            for part in msg.walk():
                pdisp = (part.get("Content-Disposition") or "").lower()
                if "attachment" not in pdisp:
                    continue
                filename = (part.get_filename() or "").lower()
                if pattern and pattern not in filename:
                    continue
                results.append({
                    "uid": uid,
                    "subject": msg.get("Subject", ""),
                    "from": msg.get("From", ""),
                    "date": msg.get("Date", ""),
                    "filename": part.get_filename() or "attachment",
                })
                break
        except Exception as e:
            logger.error(f"Error scanning message UID {uid}: {e}")
    return results


# -------------------- DRAFTS (IMAP Drafts folder) --------------------

def save_draft_to_imap(imap, acc: AccountConfig, folder: str, to: List[str], subject: str, body: str, html: bool = False):
    from email.mime.text import MIMEText
    from email.mime.multipart import MIMEMultipart

    m = MIMEMultipart()
    m["From"] = acc.smtp_from
    m["To"] = ", ".join(to)
    m["Subject"] = subject
    domain = acc.smtp_from.split("@")[1] if "@" in acc.smtp_from else "example.com"
    m["Message-ID"] = f"<{uuid.uuid4().hex}@{domain}>"
    m.attach(MIMEText(body, "html" if html else "plain"))

    ensure_selected(imap, folder)
    imap.append(folder, None, None, m.as_string().encode("utf-8"))
    return {"status": "draft_saved", "folder": folder}


# -------------------- SCHEDULED SEND (in-memory) --------------------

scheduled_sends: Dict[str, Dict[str, Any]] = {}
schedule_lock = threading.Lock()


def add_scheduled_send(account_id, to, subject, body, cc, bcc, html, send_at, in_reply_to, references):
    task_id = uuid.uuid4().hex
    with schedule_lock:
        scheduled_sends[task_id] = {
            "account_id": account_id,
            "to": to,
            "subject": subject,
            "body": body,
            "cc": cc,
            "bcc": bcc,
            "html": html,
            "send_at": float(send_at),
            "in_reply_to": in_reply_to,
            "references": references,
        }
    return task_id


def scheduled_send_loop():
    logger.info("Scheduled send loop started")
    while True:
        now = time.time()
        with schedule_lock:
            ready = [tid for tid, t in scheduled_sends.items() if t["send_at"] <= now]
        if ready:
            logger.info(f"Scheduled send loop: {len(ready)} messages ready to send")
        for tid in ready:
            with schedule_lock:
                t = scheduled_sends.pop(tid, None)
            if not t:
                continue
            logger.info(f"Scheduled send: task_id={tid} to={t['to']} subject={t['subject'][:100]} account={t['account_id']}")
            try:
                acc = get_account(t["account_id"])
                send_email(
                    acc=acc,
                    to=t["to"],
                    subject=t["subject"],
                    body=t["body"],
                    cc=t["cc"],
                    bcc=t["bcc"],
                    html=t["html"],
                    in_reply_to=t.get("in_reply_to"),
                    references=t.get("references"),
                )
                logger.info(f"Scheduled send: task_id={tid} sent successfully")
            except Exception as e:
                logger.error(f"Scheduled send error for {tid}: {e}")
        time.sleep(config.SCHEDULED_SEND_INTERVAL)


scheduler_thread = threading.Thread(target=scheduled_send_loop, daemon=True)
scheduler_thread.start()


# -------------------- EXPORT CONVERSATION --------------------

def export_conversation(imap, folder: str, message_uid: int, max_messages: int):
    ensure_selected(imap, folder)
    status, msg_data = imap.uid("FETCH", str(message_uid), "(RFC822)")
    if status != "OK" or not msg_data or msg_data[0] is None:
        return "Failed to fetch root message."
    root = email.message_from_bytes(msg_data[0][1], policy=email.policy.default)
    subject = (root.get("Subject", "") or "").strip()
    msg_id = (root.get("Message-ID", "") or "").strip()
    refs = set((root.get("References", "") or "").split())
    in_reply = (root.get("In-Reply-To", "") or "").strip()
    if in_reply:
        refs.add(in_reply)

    status, data = imap.uid("SEARCH", None, f'SUBJECT "{subject}"')
    if status != "OK" or not data or not data[0]:
        ids = []
    else:
        ids = data[0].split()

    thread_uids = []
    for uid_bytes in ids:
        uid = int(uid_bytes)
        if len(thread_uids) >= max_messages:
            break
        status, mdata = imap.uid("FETCH", str(uid), "(RFC822)")
        if status != "OK" or not mdata or mdata[0] is None:
            continue
        m = email.message_from_bytes(mdata[0][1], policy=email.policy.default)
        mid = (m.get("Message-ID", "") or "").strip()
        mrefs = set((m.get("References", "") or "").split())
        mirt = (m.get("In-Reply-To", "") or "").strip()
        if uid == message_uid or mid == msg_id:
            thread_uids.append(uid)
        elif refs and mrefs & refs:
            thread_uids.append(uid)
        elif mirt and (mirt == msg_id or mirt in refs):
            thread_uids.append(uid)
        elif mid and (mid in refs or mid == in_reply):
            thread_uids.append(uid)

    lines = []
    for uid in thread_uids:
        try:
            msg = parse_message(imap, uid)
        except Exception:
            continue
        lines.append(f"From: {msg['from']}")
        lines.append(f"To: {msg['to']}")
        lines.append(f"Date: {msg['date']}")
        lines.append(f"Subject: {msg['subject']}")
        lines.append("")
        lines.append(msg["body_plain"].strip())
        lines.append("\n---\n")

    return "\n".join(lines)


# -------------------- MCP TOOL DEFINITIONS --------------------

TOOLS = [
    {
        "name": "list_accounts",
        "description": "List all configured email accounts.",
        "inputSchema": {
            "type": "object",
            "properties": {},
            "required": []
        }
    },
    {
        "name": "list_folders",
        "description": "List all IMAP folders for an account.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"}
            },
            "required": []
        }
    },
    {
        "name": "search_messages",
        "description": "Search messages in a folder by subject and optional date.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string", "description": "IMAP folder (default: INBOX)"},
                "query": {"type": "string", "description": "Search query (subject)"},
                "since": {"type": "string", "description": "Optional SINCE date, e.g. '01-Jan-2024'"},
                "max_results": {"type": "integer", "description": "Max results (default: 50)"}
            },
            "required": []
        }
    },
    {
        "name": "get_new_messages",
        "description": "Get new messages in a folder since last check.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string", "description": "IMAP folder (default: INBOX)"}
            },
            "required": []
        }
    },
    {
        "name": "sync_seen",
        "description": "Sync seen UIDs for a folder to reset new-message tracking.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string", "description": "IMAP folder (default: INBOX)"}
            },
            "required": []
        }
    },
    {
        "name": "mark_as_read",
        "description": "Mark a message as read.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string"},
                "uid": {"type": "integer"}
            },
            "required": ["folder", "uid"]
        }
    },
    {
        "name": "flag_message",
        "description": "Flag a message.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string"},
                "uid": {"type": "integer"}
            },
            "required": ["folder", "uid"]
        }
    },
    {
        "name": "move_message",
        "description": "Move a message to another folder.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string"},
                "uid": {"type": "integer"},
                "dest_folder": {"type": "string"}
            },
            "required": ["folder", "uid", "dest_folder"]
        }
    },
    {
        "name": "copy_message",
        "description": "Copy a message to another folder.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string"},
                "uid": {"type": "integer"},
                "dest_folder": {"type": "string"}
            },
            "required": ["folder", "uid", "dest_folder"]
        }
    },
    {
        "name": "apply_label",
        "description": "Apply a label (custom flag) to a message.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string"},
                "uid": {"type": "integer"},
                "label": {"type": "string", "description": "Label name (without prefix)"}
            },
            "required": ["folder", "uid", "label"]
        }
    },
    {
        "name": "remove_label",
        "description": "Remove a label from a message.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string"},
                "uid": {"type": "integer"},
                "label": {"type": "string", "description": "Label name (without prefix)"}
            },
            "required": ["folder", "uid", "label"]
        }
    },
    {
        "name": "list_labels",
        "description": "List labels for a message or available labels in a folder.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string", "description": "IMAP folder (default: INBOX)"},
                "uid": {"type": "integer", "description": "If set, list labels for that message; else list used labels in folder."}
            },
            "required": []
        }
    },
    {
        "name": "get_unread_summary",
        "description": "Get unread message summary for a folder.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string", "description": "IMAP folder (default: INBOX)"}
            },
            "required": []
        }
    },
    {
        "name": "send_email",
        "description": "Send an email using a specific account (or default).",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID to send from (optional, uses default if omitted)"},
                "to": {"type": "array", "items": {"type": "string"}},
                "subject": {"type": "string"},
                "body": {"type": "string"},
                "cc": {"type": "array", "items": {"type": "string"}},
                "bcc": {"type": "array", "items": {"type": "string"}},
                "html": {"type": "boolean", "description": "Treat body as HTML"},
                "in_reply_to": {"type": "string"},
                "references": {"type": "string"}
            },
            "required": ["to", "subject", "body"]
        }
    },
    {
        "name": "reply_to_message",
        "description": "Reply to a specific message using the same account.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string"},
                "uid": {"type": "integer"},
                "body": {"type": "string"},
                "html": {"type": "boolean"},
                "to": {"type": "array", "items": {"type": "string"}, "description": "Override reply recipients"},
                "cc": {"type": "array", "items": {"type": "string"}},
                "bcc": {"type": "array", "items": {"type": "string"}}
            },
            "required": ["folder", "uid", "body"]
        }
    },
    {
        "name": "forward_message",
        "description": "Forward a message.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string"},
                "uid": {"type": "integer"},
                "to": {"type": "array", "items": {"type": "string"}},
                "cc": {"type": "array", "items": {"type": "string"}},
                "bcc": {"type": "array", "items": {"type": "string"}},
                "note": {"type": "string", "description": "Optional note to prepend"},
                "html": {"type": "boolean"}
            },
            "required": ["folder", "uid", "to"]
        }
    },
    {
        "name": "save_draft",
        "description": "Save a draft to an IMAP folder.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "to": {"type": "array", "items": {"type": "string"}},
                "subject": {"type": "string"},
                "body": {"type": "string"},
                "html": {"type": "boolean"},
                "folder": {"type": "string", "description": "IMAP folder (default: Drafts)"}
            },
            "required": ["to", "subject", "body"]
        }
    },
    {
        "name": "list_attachments",
        "description": "List attachments for a message.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string"},
                "uid": {"type": "integer"}
            },
            "required": ["folder", "uid"]
        }
    },
    {
        "name": "download_attachment",
        "description": "Download an attachment (returns base64 content).",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string"},
                "uid": {"type": "integer"},
                "filename": {"type": "string"}
            },
            "required": ["folder", "uid", "filename"]
        }
    },
    {
        "name": "search_attachments",
        "description": "Search for attachments by filename pattern.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string", "description": "IMAP folder (default: INBOX)"},
                "file_pattern": {"type": "string", "description": "Filename substring to match (case-insensitive)"},
                "max_results": {"type": "integer", "description": "Max results (default: 50)"}
            },
            "required": []
        }
    },
    {
        "name": "schedule_send",
        "description": "Schedule an email to be sent at a future time using a specific account (or default).",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "to": {"type": "array", "items": {"type": "string"}},
                "subject": {"type": "string"},
                "body": {"type": "string"},
                "cc": {"type": "array", "items": {"type": "string"}},
                "bcc": {"type": "array", "items": {"type": "string"}},
                "html": {"type": "boolean"},
                "send_at": {"type": "number", "description": "Unix timestamp when to send"},
                "in_reply_to": {"type": "string"},
                "references": {"type": "string"}
            },
            "required": ["to", "subject", "body", "send_at"]
        }
    },
    {
        "name": "export_conversation",
        "description": "Export a conversation thread as plain text.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "folder": {"type": "string"},
                "message_uid": {"type": "integer"},
                "max_messages": {"type": "integer", "description": "Max messages (default: 50)"}
            },
            "required": ["folder", "message_uid"]
        }
    },
    {
        "name": "conflict_check_search",
        "description": "Search for potential conflicts by terms (names, firms, counsel).",
        "inputSchema": {
            "type": "object",
            "properties": {
                "account": {"type": "string", "description": "Account ID (optional, uses default if omitted)"},
                "terms": {"type": "array", "items": {"type": "string"}},
                "folders": {"type": "array", "items": {"type": "string"}},
                "max_results": {"type": "integer", "description": "Max results (default: 100)"}
            },
            "required": ["terms"]
        }
    },
    {
        "name": "get_triage_config",
        "description": "Get triage configuration (business description and personnel).",
        "inputSchema": {
            "type": "object",
            "properties": {},
            "required": []
        }
    },
]


# -------------------- MCP JSON-RPC HANDLER --------------------

initialized = False


def mcp_error(rid, code, message):
    return {
        "jsonrpc": "2.0",
        "id": rid,
        "error": {
            "code": code,
            "message": message
        }
    }


def handle_rpc(msg):
    jsonrpc = msg.get("jsonrpc")
    method = msg.get("method")
    params = msg.get("params", {})
    rid = msg.get("id")
    prefix = log_context_prefix()

    if jsonrpc != "2.0":
        logger.info(f"{prefix}MCP invalid jsonrpc: {jsonrpc}")
        return mcp_error(rid, -32600, "Invalid Request")

    logger.debug(f"{prefix}MCP RPC method={method} params_keys={list(params.keys())}")

    try:
        if method == "initialize":
            global initialized
            initialized = True
            logger.info(f"{prefix}MCP: initialize")
            return {
                "jsonrpc": "2.0",
                "id": rid,
                "result": {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {
                        "tools": {}
                    },
                    "serverInfo": {
                        "name": "mcp-email-server",
                        "version": "1.0.0"
                    }
                }
            }

        if method == "tools/list":
            logger.info(f"{prefix}MCP: tools/list")
            return {
                "jsonrpc": "2.0",
                "id": rid,
                "result": {
                    "tools": TOOLS
                }
            }

        if method == "tools/call":
            tool_name = params.get("name")
            tool_args = params.get("arguments", {})
            logger.info(f"{prefix}MCP: tools/call name={tool_name} args_keys={list(tool_args.keys())}")
            if logger.isEnabledFor(logging.DEBUG):
                safe_args = {k: (str(v)[:400] if isinstance(v, str) else v) for k, v in tool_args.items()}
                logger.debug(f"{prefix}MCP: tools/call args={json.dumps(safe_args, default=str)}")
            result = call_tool(tool_name, tool_args)
            return {
                "jsonrpc": "2.0",
                "id": rid,
                "result": {
                    "content": [
                        {"type": "text", "text": json.dumps(result)}
                    ]
                }
            }

        logger.warning(f"{prefix}MCP: unknown method={method}")
        return mcp_error(rid, -32601, "Method not found")

    except Exception as e:
        logger.exception(f"{prefix}MCP tool error for method={method}")
        return mcp_error(rid, -32603, f"Internal error: {str(e)}")


def call_tool(name, args):
    prefix = log_context_prefix()
    logger.info(f"{prefix}call_tool name={name}")

    # IMAP tools (run in thread pool)
    def imap_op(fn):
        logger.info(f"{prefix}IMAP operation: {fn.__name__}")
        try:
            result = executor.submit(fn, **args).result()
            logger.debug(f"{prefix}IMAP operation {fn.__name__} completed successfully")
            return result
        except Exception as e:
            logger.error(f"{prefix}IMAP operation {fn.__name__} failed: {e}")
            raise

    if name == "list_accounts":
        logger.info(f"{prefix}call_tool list_accounts")
        return list_accounts_impl(args)

    if name == "list_folders":
        return imap_op(list_folders_impl)

    if name == "search_messages":
        return imap_op(search_messages_impl)

    if name == "get_new_messages":
        return imap_op(get_new_messages_impl)

    if name == "sync_seen":
        return imap_op(sync_seen_impl)

    if name == "mark_as_read":
        return imap_op(mark_as_read_impl)

    if name == "flag_message":
        return imap_op(flag_message_impl)

    if name == "move_message":
        return imap_op(move_message_impl)

    if name == "copy_message":
        return imap_op(copy_message_impl)

    if name == "apply_label":
        return imap_op(apply_label_impl)

    if name == "remove_label":
        return imap_op(remove_label_impl)

    if name == "list_labels":
        return imap_op(list_labels_impl)

    if name == "get_unread_summary":
        return imap_op(get_unread_summary_impl)

    if name == "reply_to_message":
        return imap_op(reply_to_message_impl)

    if name == "forward_message":
        return imap_op(forward_message_impl)

    if name == "save_draft":
        return imap_op(save_draft_impl)

    if name == "list_attachments":
        return imap_op(list_attachments_impl)

    if name == "download_attachment":
        return imap_op(download_attachment_impl)

    if name == "search_attachments":
        return imap_op(search_attachments_impl)

    if name == "export_conversation":
        return imap_op(export_conversation_impl)

    if name == "conflict_check_search":
        return imap_op(conflict_check_search_impl)

    # Non-IMAP tools
    if name == "send_email":
        logger.info(f"{prefix}call_tool send_email")
        return send_email_impl(args)

    if name == "schedule_send":
        logger.info(f"{prefix}call_tool schedule_send")
        return schedule_send_impl(args)

    if name == "get_triage_config":
        logger.info(f"{prefix}call_tool get_triage_config")
        return get_triage_config_impl(args)

    raise ValueError(f"Unknown tool: {name}")


# -------------------- TOOL IMPLEMENTATIONS --------------------

import concurrent.futures

executor = concurrent.futures.ThreadPoolExecutor(max_workers=4)


def list_accounts_impl(args):
    accounts = []
    for aid, acc in config.accounts.items():
        accounts.append({
            "id": acc.id,
            "name": acc.name or acc.id,
            "imap_username": acc.imap_username,
            "smtp_from": acc.smtp_from,
            "is_default": aid == config.default_account_id,
        })
    return {
        "accounts": accounts,
        "default_account": config.default_account_id,
    }


def list_folders_impl(account=None):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        folders = list_all_folders(imap)
        return {"folders": folders}
    finally:
        imap.logout()


def search_messages_impl(account=None, folder="INBOX", query="", since=None, max_results=50):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        results = search_messages(imap, folder=folder, query=query, since=since, max_results=max_results)
        return {"messages": results}
    finally:
        imap.logout()


def get_new_messages_impl(account=None, folder="INBOX"):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        new = get_new_messages(imap, folder)
        out = []
        for m in new:
            m["triage_hint"] = build_triage_hint(m)
            out.append(m)
        return {"messages": out}
    finally:
        imap.logout()


def sync_seen_impl(account=None, folder="INBOX"):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        sync_seen(imap, folder)
        return {"status": "ok", "folder": folder}
    finally:
        imap.logout()


def mark_as_read_impl(account=None, folder=None, uid=None):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        mark_as_read(imap, folder, uid)
        return {"status": "ok"}
    finally:
        imap.logout()


def flag_message_impl(account=None, folder=None, uid=None):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        flag_message(imap, folder, uid)
        return {"status": "ok"}
    finally:
        imap.logout()


def move_message_impl(account=None, folder=None, uid=None, dest_folder=None):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        move_message(imap, folder, uid, dest_folder)
        return {"status": "moved", "dest_folder": dest_folder}
    finally:
        imap.logout()


def copy_message_impl(account=None, folder=None, uid=None, dest_folder=None):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        copy_message(imap, folder, uid, dest_folder)
        return {"status": "copied", "dest_folder": dest_folder}
    finally:
        imap.logout()


def apply_label_impl(account=None, folder=None, uid=None, label=None):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        add_flags(imap, folder, uid, [f"$Label_{label}"])
        return {"status": "applied", "label": label}
    finally:
        imap.logout()


def remove_label_impl(account=None, folder=None, uid=None, label=None):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        remove_flags(imap, folder, uid, [f"$Label_{label}"])
        return {"status": "removed", "label": label}
    finally:
        imap.logout()


def list_labels_impl(account=None, folder="INBOX", uid=None):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        if uid:
            labels = get_labels(imap, folder, uid)
            return {"labels": [l.replace("$Label_", "", 1) for l in labels]}
        else:
            labels = list_available_labels(imap)
            return {"labels": labels}
    finally:
        imap.logout()


def get_unread_summary_impl(account=None, folder="INBOX"):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        ensure_selected(imap, folder)
        status, data = imap.uid("SEARCH", None, "UNSEEN")
        if status != "OK" or not data or not data[0]:
            ids = []
        else:
            ids = data[0].split()
        uids = [int(u) for u in ids]
        count = len(uids)
        previews = []
        for uid in uids[:20]:
            msg = parse_message(imap, uid)
            previews.append({
                "uid": uid,
                "subject": msg["subject"],
                "from": msg["from"],
                "date": msg["date"],
                "snippet": (msg["body_plain"] or "")[:300],
            })
        return {
            "folder": folder,
            "unread_count": count,
            "previews": previews,
        }
    finally:
        imap.logout()


def send_email_impl(args):
    account_id = args.get("account") or config.default_account_id
    acc = get_account(account_id)

    # Default to HTML unless explicitly set to false
    html = args.get("html")
    if html is None:
        html = True

    body = args["body"]
    if html and config.EMAIL_HTML_SIGNATURE:
        body = body.rstrip("\n") + "\n" + config.EMAIL_HTML_SIGNATURE

    try:
        send_email(
            acc=acc,
            to=args["to"],
            subject=args["subject"],
            body=body,
            cc=args.get("cc"),
            bcc=args.get("bcc"),
            html=bool(html),
            in_reply_to=args.get("in_reply_to"),
            references=args.get("references"),
        )
        return {"status": "sent", "account": acc.id}
    except Exception as e:
        raise RuntimeError(str(e))


def schedule_send_impl(args):
    account_id = args.get("account") or config.default_account_id

    # Default to HTML unless explicitly set to false
    html = args.get("html")
    if html is None:
        html = True

    body = args["body"]
    if html and config.EMAIL_HTML_SIGNATURE:
        body = body.rstrip("\n") + "\n" + config.EMAIL_HTML_SIGNATURE

    task_id = add_scheduled_send(
        account_id=account_id,
        to=args["to"],
        subject=args["subject"],
        body=body,
        cc=args.get("cc"),
        bcc=args.get("bcc"),
        html=bool(html),
        send_at=float(args["send_at"]),
        in_reply_to=args.get("in_reply_to"),
        references=args.get("references"),
    )
    return {"status": "scheduled", "task_id": task_id, "account": account_id}


def reply_to_message_impl(account=None, folder=None, uid=None, body=None, html=None, to=None, cc=None, bcc=None):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        msg = parse_message(imap, folder, uid)

        from email.utils import getaddresses
        if to:
            reply_to_list = to
        else:
            addrs = []
            for _, addr in getaddresses([msg["to"], msg["from"]]):
                if addr:
                    addrs.append(addr)
            reply_to_list = addrs

        # Default to HTML unless explicitly set to false
        if html is None:
            html = True

        subject = msg["subject"]
        if not subject.lower().startswith("re:"):
            subject = f"Re: {subject}"

        in_reply_to = msg.get("message_id") or None
        refs = (msg.get("references") or "").strip()
        if in_reply_to and refs:
            refs = f"{refs} {in_reply_to}"
        elif in_reply_to:
            refs = in_reply_to

        # Append HTML signature if sending as HTML
        if html and config.EMAIL_HTML_SIGNATURE:
            body = (body or "").rstrip("\n") + "\n" + config.EMAIL_HTML_SIGNATURE

        send_email(
            acc=acc,
            to=reply_to_list,
            subject=subject,
            body=body,
            cc=cc,
            bcc=bcc,
            html=html,
            in_reply_to=in_reply_to,
            references=refs,
        )
        return {"status": "replied", "account": acc.id}
    finally:
        imap.logout()


def forward_message_impl(account=None, folder=None, uid=None, to=None, cc=None, bcc=None, note="", html=None, body=None):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        msg = parse_message(imap, folder, uid)

        # Default to HTML unless explicitly set to false
        if html is None:
            html = True

        subject = msg["subject"]
        if not subject.lower().startswith("fwd:"):
            subject = f"Fwd: {subject}"

        body = body or ""
        if note:
            body = f"{note}\n\n" + body

        original = msg["body_plain"] or (msg["body_html"] or "")
        if original:
            quoted = "\n".join("> " + line for line in original.splitlines())
            body = (body + "\n\n" + quoted).strip()

        # Append HTML signature if sending as HTML
        if html and config.EMAIL_HTML_SIGNATURE:
            body = body.rstrip("\n") + "\n" + config.EMAIL_HTML_SIGNATURE

        send_email(
            acc=acc,
            to=to,
            subject=subject,
            body=body,
            cc=cc,
            bcc=bcc,
            html=html,
        )
        return {"status": "forwarded", "account": acc.id}
    finally:
        imap.logout()


def save_draft_impl(account=None, to=None, subject=None, body=None, html=False, folder="Drafts"):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        return save_draft_to_imap(imap, acc, folder, to, subject, body, html)
    finally:
        imap.logout()


def list_attachments_impl(account=None, folder=None, uid=None):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        atts = list_attachments(imap, folder, uid)
        return {"attachments": atts}
    finally:
        imap.logout()


def download_attachment_impl(account=None, folder=None, uid=None, filename=None):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        info = download_attachment(imap, folder, uid, filename)
        return info
    except ValueError as e:
        raise RuntimeError(str(e))
    finally:
        imap.logout()


def search_attachments_impl(account=None, folder="INBOX", file_pattern="", max_results=50):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        results = search_attachments(imap, folder, file_pattern, max_results)
        return {"results": results}
    finally:
        imap.logout()


def export_conversation_impl(account=None, folder=None, message_uid=None, max_messages=50):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        text = export_conversation(imap, folder, message_uid, max_messages)
        return {"export": text}
    finally:
        imap.logout()


def conflict_check_search_impl(account=None, terms=None, folders=None, max_results=100):
    acc = get_account(account)
    imap = create_imap(acc)
    try:
        login_imap(imap, acc)
        results = conflict_check_search(imap, terms, folders, max_results)
        return {"results": results}
    finally:
        imap.logout()


def get_triage_config_impl(args):
    return {
        "business_description": config.BUSINESS_DESCRIPTION,
        "personnel": config.personnel,
    }


# -------------------- FASTAPI / MCP ENDPOINT --------------------

app = FastAPI(
    title="MCP Email Server",
    description="Email assistant MCP server exposing IMAP/SMTP and conflict-check operations via MCP Streamable HTTP.",
    version="1.0.0",
)

API_KEY = os.getenv("API_KEY", "").strip()


@app.exception_handler(Exception)
async def global_exception_handler(request, exc):
    logger.exception("Unhandled exception")
    return JSONResponse(
        status_code=500,
        content={"detail": "Internal server error"}
    )


def require_bearer(request: Request):
    if not API_KEY:
        return
    auth = (request.headers.get("authorization") or "").strip()
    if not auth.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
    token = auth[len("Bearer "):].strip()
    if token != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API key")


@app.get("/health")
async def health_check():
    return {
        "status": "ok",
        "service": "mcp-email-server",
        "version": "1.0.0"
    }


@app.get("/mcp/openapi.json")
async def mcp_openapi_spec(request: Request):
    require_bearer(request)
    return {
        "openapi": "3.1.0",
        "info": {
            "title": "MCP Email Server",
            "version": "1.0.0",
            "description": "Email assistant MCP server. Tools are exposed via MCP Streamable HTTP at POST /mcp."
        },
        "paths": {
            "/mcp": {
                "post": {
                    "summary": "MCP JSON-RPC endpoint",
                    "description": "Primary MCP endpoint (Streamable HTTP). Use JSON-RPC 2.0 with methods: initialize, tools/list, tools/call.",
                    "requestBody": {
                        "required": True,
                        "content": {
                            "application/json": {
                                "schema": {
                                    "type": "object",
                                    "description": "JSON-RPC 2.0 request or batch of requests"
                                }
                            }
                        }
                    },
                    "responses": {
                        "200": {
                            "description": "JSON-RPC 2.0 response",
                            "content": {
                                "application/json": {
                                    "schema": {
                                        "type": "object"
                                    }
                                }
                            }
                        },
                        "400": {
                            "description": "Invalid JSON or JSON-RPC request"
                        },
                        "500": {
                            "description": "Internal server error"
                        }
                    }
                }
            },
            "/health": {
                "get": {
                    "summary": "Health check",
                    "responses": {
                        "200": {
                            "description": "Service is running",
                            "content": {
                                "application/json": {
                                    "schema": {
                                        "type": "object",
                                        "properties": {
                                            "status": {"type": "string"},
                                            "service": {"type": "string"},
                                            "version": {"type": "string"}
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }


@app.post("/mcp")
async def mcp_endpoint(request: Request):
    require_bearer(request)
    if request.method != "POST":
        raise HTTPException(status_code=405, detail="Method Not Allowed")

    req_id = str(uuid.uuid4())
    request_id_var.set(req_id)
    prefix = f"[req:{req_id}]"

    logger.info(f"{prefix} POST /mcp")
    if logger.isEnabledFor(logging.DEBUG):
        logger.debug(f"{prefix} POST /mcp headers={dict(request.headers)}")

    try:
        body = await request.json()
    except Exception:
        logger.warning(f"{prefix} POST /mcp invalid JSON")
        return JSONResponse(
            status_code=400,
            content=mcp_error(None, -32700, "Parse error")
        )

    if isinstance(body, list):
        results = [handle_rpc(msg) for msg in body]
        logger.debug(f"{prefix} POST /mcp batch response count={len(results)}")
        return JSONResponse(content=results)
    else:
        result = handle_rpc(body)
        logger.debug(f"{prefix} POST /mcp response keys={list(result.keys())}")
        return JSONResponse(content=result)


# -------------------- ENTRYPOINT --------------------

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("server:app", host="0.0.0.0", port=8000, reload=False)