"""Email connector — send via SMTP, search via IMAP, or just draft a message.

The SMTP/IMAP calls are isolated in small methods so tests can stub them without
a live mail server.
"""

from __future__ import annotations

import imaplib
import smtplib
from email.message import EmailMessage

from legate.agent.tooling import ToolContext, ToolResult, ToolSpec
from legate.connectors.base import BaseConnector, ConnectionResult


class EmailConnector(BaseConnector):
    type = "email"
    title = "Email (SMTP/IMAP)"
    config_schema = {
        "type": "object",
        "properties": {
            "smtp_host": {"type": "string"},
            "smtp_port": {"type": "integer"},
            "imap_host": {"type": "string"},
            "imap_port": {"type": "integer"},
            "username": {"type": "string"},
            "from_addr": {"type": "string"},
            "use_tls": {"type": "boolean"},
        },
        "required": ["smtp_host", "username"],
    }
    secret_schema = {
        "type": "object",
        "properties": {"password": {"type": "string"}},
        "required": ["password"],
    }

    def list_tools(self) -> list[ToolSpec]:
        compose_params = {
            "type": "object",
            "properties": {
                "to": {"type": "string"},
                "subject": {"type": "string"},
                "body": {"type": "string"},
            },
            "required": ["to", "subject", "body"],
        }
        return [
            ToolSpec(
                name="email.draft",
                description="Compose an email without sending it.",
                parameters=compose_params,
                is_sensitive=False,
                scopes=["email:write"],
            ),
            ToolSpec(
                name="email.send",
                description="Send an email via SMTP.",
                parameters=compose_params,
                is_sensitive=True,
                scopes=["email:send"],
            ),
            ToolSpec(
                name="email.search",
                description="Search a mailbox via IMAP and return message headers.",
                parameters={
                    "type": "object",
                    "properties": {
                        "mailbox": {"type": "string"},
                        "limit": {"type": "integer"},
                    },
                },
                is_sensitive=False,
                scopes=["email:read"],
            ),
        ]

    def _compose(self, config: dict, args: dict) -> EmailMessage:
        message = EmailMessage()
        message["From"] = config.get("from_addr") or config.get("username", "")
        message["To"] = args.get("to", "")
        message["Subject"] = args.get("subject", "")
        message.set_content(args.get("body", ""))
        return message

    # -- network seams (stubbed in tests) ---------------------------------

    def _smtp_send(self, config: dict, secrets: dict, message: EmailMessage) -> None:
        host = config["smtp_host"]
        port = int(config.get("smtp_port", 587))
        with smtplib.SMTP(host, port) as server:
            if config.get("use_tls", True):
                server.starttls()
            server.login(config["username"], secrets["password"])
            server.send_message(message)

    def _imap_search(self, config: dict, secrets: dict, mailbox: str, limit: int) -> list[dict]:
        host = config.get("imap_host", config["smtp_host"])
        port = int(config.get("imap_port", 993))
        server = imaplib.IMAP4_SSL(host, port)
        try:
            server.login(config["username"], secrets["password"])
            server.select(mailbox)
            _, data = server.search(None, "ALL")
            ids = data[0].split()[:limit]
            results = []
            for msg_id in ids:
                _, msg_data = server.fetch(msg_id, "(BODY[HEADER.FIELDS (FROM SUBJECT DATE)])")
                first = msg_data[0] if msg_data else None
                headers = ""
                if isinstance(first, tuple) and len(first) >= 2 and isinstance(first[1], bytes):
                    headers = first[1].decode(errors="replace")
                results.append({"id": msg_id.decode(), "headers": headers})
            return results
        finally:
            server.logout()

    def _smtp_login(self, config: dict, secrets: dict) -> None:
        host = config["smtp_host"]
        port = int(config.get("smtp_port", 587))
        with smtplib.SMTP(host, port) as server:
            if config.get("use_tls", True):
                server.starttls()
            server.login(config["username"], secrets["password"])

    # -- connector API ----------------------------------------------------

    async def test_connection(self, config: dict, secrets: dict) -> ConnectionResult:
        try:
            self._smtp_login(config, secrets)
        except Exception as exc:  # smtplib raises a variety of errors
            return ConnectionResult(ok=False, message=f"SMTP login failed: {exc}")
        return ConnectionResult(ok=True, message="SMTP login succeeded")

    async def execute(
        self, tool: str, args: dict, config: dict, secrets: dict, ctx: ToolContext
    ) -> ToolResult:
        if tool == "email.draft":
            message = self._compose(config, args)
            return ToolResult(
                ok=True,
                data={
                    "to": args.get("to"),
                    "subject": args.get("subject"),
                    "mime": message.as_string(),
                },
            )
        if tool == "email.send":
            try:
                self._smtp_send(config, secrets, self._compose(config, args))
            except Exception as exc:
                return ToolResult(ok=False, error=f"send failed: {exc}")
            return ToolResult(ok=True, data={"sent": True, "to": args.get("to")})
        if tool == "email.search":
            try:
                results = self._imap_search(
                    config, secrets, args.get("mailbox", "INBOX"), int(args.get("limit", 10))
                )
            except Exception as exc:
                return ToolResult(ok=False, error=f"search failed: {exc}")
            return ToolResult(ok=True, data={"messages": results})
        return ToolResult(ok=False, error=f"unknown email tool: {tool}")
