from __future__ import annotations

import asyncio
import logging
from contextlib import asynccontextmanager
from typing import Any

import httpx
from mcp.server.fastmcp import Context, FastMCP

from .client import DropmailClient, DropmailError

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# GraphQL queries and mutations
# ---------------------------------------------------------------------------

_Q_LIST_DOMAINS = """
query ListDomains {
  domains {
    id
    name
    availableVia
  }
}
"""

_Q_CREATE_SESSION = """
mutation CreateSession($input: IntroduceSessionInput) {
  introduceSession(input: $input) {
    id
    expiresAt
    addresses {
      id
      address
      restoreKey
    }
  }
}
"""

_Q_GET_SESSION = """
query GetSession($id: ID!) {
  session(id: $id) {
    id
    expiresAt
    mailsCount
    addresses {
      id
      address
      restoreKey
    }
  }
}
"""

_Q_GET_INBOX = """
query GetInbox($id: ID!, $mailId: ID) {
  session(id: $id) {
    mailsAfterId(mailId: $mailId) {
      id
      receivedAt
      fromAddr
      headerFrom
      headerSubject
      toAddr
      rawSize
      text
    }
  }
}
"""

_Q_WAIT_FOR_EMAIL = """
subscription WaitForEmail($id: ID!, $mailsAfterId: ID) {
  sessionMailReceived(id: $id, mailsAfterId: $mailsAfterId) {
    id
    receivedAt
    fromAddr
    headerFrom
    headerSubject
    toAddr
    rawSize
    text
  }
}
"""

_Q_READ_EMAIL = """
query ReadEmail($id: ID!, $includeSanitizedHtml: Boolean = false) {
  node(id: $id) {
    ... on Mail {
      id
      receivedAt
      fromAddr
      headerFrom
      headerSubject
      toAddr
      rawSize
      text
      textSource
      decodeStatus
      hasHtml
      sanitizedHtml @include(if: $includeSanitizedHtml)
      downloadSanitizedHtmlUrl
      downloadUrl
      attachments {
        id
        name
        mime
        rawSize
        downloadUrl
      }
    }
  }
}
"""

_Q_ADD_ADDRESS = """
mutation AddAddress($input: IntroduceAddressInput!) {
  introduceAddress(input: $input) {
    id
    address
    restoreKey
  }
}
"""

_Q_RESTORE_ADDRESS = """
mutation RestoreAddress($input: RestoreAddressInput!) {
  restoreAddress(input: $input) {
    id
    address
    restoreKey
  }
}
"""

_Q_DELETE_ADDRESS = """
mutation DeleteAddress($input: DeleteAddressInput!) {
  deleteAddress(input: $input)
}
"""


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _text_preview(text: str | None) -> str:
    """Collapse whitespace and truncate to 1024 characters."""
    if not text:
        return ""
    return " ".join(text.split())[:1024]


def _shape_mail_preview(mail: dict[str, Any]) -> dict[str, Any]:
    """Shape a mail object for lightweight listing: rename text -> textPreview."""
    result = {k: v for k, v in mail.items() if k != "text"}
    result["textPreview"] = _text_preview(mail.get("text"))
    return result


# ---------------------------------------------------------------------------
# Server factory
# ---------------------------------------------------------------------------

def create_server(client: DropmailClient) -> FastMCP:
    @asynccontextmanager
    async def lifespan(server: FastMCP):  # type: ignore[type-arg]
        try:
            yield
        finally:
            await client.aclose()

    mcp = FastMCP(
        "dropmail",
        instructions=(
            "DropMail provides disposable temporary email addresses. "
            "Typical use cases: "
            '"Give me a temp email and wait for the confirmation from example.com"; '
            '"Sign me up for the newsletter at example.com using a throwaway address"; '
            '"I need to sign up for a product trial on example.com — give me a .com temp email, '
            'wait for their email, and pull out the activation code"; '
            '"Check if the password reset email has arrived and give me the link". '
            "Workflow: (optionally list_domains) → create_session → (optionally add_address / "
            "restore_address if restore_key available) → "
            "wait_for_email → read_email."
        ),
        lifespan=lifespan,
    )

    # ------------------------------------------------------------------

    @mcp.tool()
    async def list_domains() -> list[dict[str, Any]]:
        """List all email domains available for creating temporary addresses.

        Returns each domain's id (needed for create_session / add_address),
        name (e.g. '@dropmail.me'), and availableVia interfaces.
        """
        logger.info("tool=list_domains")
        data, errors = await client.execute(_Q_LIST_DOMAINS)
        return client.unwrap(data, errors, "domains")

    # ------------------------------------------------------------------

    @mcp.tool()
    async def create_session(domain_id: str | None = None) -> dict[str, Any]:
        """Create a new DropMail session with one temporary email address.

        Returns session_id (needed for all subsequent calls), the email address,
        restore_key (save this to restore the address later), and expiry time.
        The session lifetime is extended automatically on each access.

        Args:
            domain_id: Optional domain ID from list_domains. Random if omitted.
        """
        logger.info("tool=create_session domain_id=%s", domain_id)
        variables: dict[str, Any] = {
            "input": {"withAddress": True, **({"domainId": domain_id} if domain_id else {})},
        }
        data, errors = await client.execute(_Q_CREATE_SESSION, variables)
        session = client.unwrap(data, errors, "introduceSession")
        address = (session.get("addresses") or [{}])[0]
        return {
            "session_id": session["id"],
            "expires_at": session["expiresAt"],
            "address": address.get("address"),
            "address_id": address.get("id"),
            "restore_key": address.get("restoreKey"),
        }

    # ------------------------------------------------------------------

    @mcp.tool()
    async def get_session(session_id: str) -> dict[str, Any]:
        """Get session metadata: active addresses, total mail count, and expiry.

        Use get_inbox to retrieve the actual emails.

        Args:
            session_id: Session ID returned by create_session.
        """
        logger.info("tool=get_session session_id=%s", session_id)
        data, errors = await client.execute(_Q_GET_SESSION, {"id": session_id})
        return client.unwrap(data, errors, "session")

    # ------------------------------------------------------------------

    @mcp.tool()
    async def get_inbox(
        session_id: str,
        after_id: str | None = None,
    ) -> list[dict[str, Any]]:
        """List emails received by the session, newest first (up to 100).

        Each entry includes id, receivedAt, fromAddr, headerFrom, headerSubject,
        toAddr, rawSize, and textPreview (first 1024 chars, whitespace collapsed).
        Use read_email for the full content of a specific message.

        Args:
            session_id: Session ID returned by create_session.
            after_id:   Only return emails received after this mail ID.
                        Pass the last seen mail's id to paginate incrementally.
        """
        logger.info("tool=get_inbox session_id=%s after_id=%s", session_id, after_id)
        data, errors = await client.execute(
            _Q_GET_INBOX, {"id": session_id, "mailId": after_id}
        )
        session = client.unwrap(data, errors, "session")
        mails: list[dict[str, Any]] = session.get("mailsAfterId") or []
        return [_shape_mail_preview(m) for m in mails]

    # ------------------------------------------------------------------

    @mcp.tool()
    async def wait_for_email(
        session_id: str,
        after_id: str | None = None,
        timeout: float = 180.0,
        ctx: Context = None,
    ) -> dict[str, Any]:
        """Block until a new email arrives in the session (long-poll, up to timeout seconds).

        On arrival returns the mail with the same fields as get_inbox entries,
        plus {"arrived": true}. Use the returned mail id as after_id on the next
        call to avoid re-receiving the same message.
        Returns {"arrived": false, "reason": "timeout"} if no email arrives in time —
        this is a normal outcome, not an error. Simply call again with the same after_id.

        Progress notifications are sent every 15 seconds so MCP clients know the
        tool is still working.

        Args:
            session_id: Session ID returned by create_session.
            after_id:   Only notify for emails received after this mail ID.
            timeout:    Maximum seconds to wait (default 180).
        """
        logger.info("tool=wait_for_email session_id=%s after_id=%s timeout=%s", session_id, after_id, timeout)

        async def _heartbeat() -> None:
            elapsed = 0.0
            while True:
                await asyncio.sleep(15)
                elapsed += 15
                remaining = max(0.0, timeout - elapsed)
                try:
                    if ctx is not None:
                        await ctx.report_progress(
                            elapsed, timeout,
                            f"Waiting for email… ({int(remaining)}s remaining)"
                        )
                except Exception:
                    pass

        heartbeat = asyncio.create_task(_heartbeat())
        variables: dict[str, Any] = {
            "id": session_id,
            "mailsAfterId": after_id,
        }
        try:
            data, errors = await client.execute(
                _Q_WAIT_FOR_EMAIL, variables, timeout=timeout + 10.0
            )
        except (httpx.TimeoutException, asyncio.CancelledError, asyncio.TimeoutError):
            logger.info("tool=wait_for_email result=timeout")
            return {"arrived": False, "reason": "timeout"}
        finally:
            heartbeat.cancel()
        for e in errors or []:
            code = (e.get("extensions") or {}).get("code")
            if code == "SESSION_NOT_FOUND":
                logger.info("tool=wait_for_email result=session_not_found")
                client.unwrap(None, errors, "sessionMailReceived")
            if code == "timeout":
                logger.info("tool=wait_for_email result=timeout (server)")
                return {"arrived": False, "reason": "timeout"}
        mail = (data or {}).get("sessionMailReceived")
        if mail is not None:
            logger.info("tool=wait_for_email result=arrived mail_id=%s", mail.get("id"))
            return {"arrived": True, **_shape_mail_preview(mail)}
        logger.info("tool=wait_for_email result=timeout (null)")
        return {"arrived": False, "reason": "timeout"}

    # ------------------------------------------------------------------

    @mcp.tool()
    async def read_email(
        mail_id: str,
        include_html: bool = False,
    ) -> dict[str, Any]:
        """Fetch the full content of a specific email.

        Returns full plain-text body, decode status, attachment list with
        download URLs, and URLs for the raw MIME and sanitized HTML versions.

        Links are preserved as Markdown [anchor](url) in the text field when
        the email was HTML-only (textSource=HTML). If textSource=TEXT and the
        expected information (link, code) is missing from text, retry with
        include_html=True — the sender's HTML body may contain more complete
        content.

        Args:
            mail_id:      Mail ID from get_inbox or wait_for_email.
            include_html: Include sanitizedHtml inline (can be large). Default false.
        """
        logger.info("tool=read_email mail_id=%s include_html=%s", mail_id, include_html)
        data, errors = await client.execute(
            _Q_READ_EMAIL,
            {"id": mail_id, "includeSanitizedHtml": include_html},
        )
        return client.unwrap(data, errors, "node")

    # ------------------------------------------------------------------

    @mcp.tool()
    async def add_address(
        session_id: str,
        domain_id: str | None = None,
    ) -> dict[str, Any]:
        """Add a second temporary email address to an existing session.

        All addresses in a session share the same inbox.

        Args:
            session_id: Session ID returned by create_session.
            domain_id:  Optional domain ID from list_domains. Random if omitted.
        """
        logger.info("tool=add_address session_id=%s domain_id=%s", session_id, domain_id)
        variables: dict[str, Any] = {
            "input": {
                "sessionId": session_id,
                **({"domainId": domain_id} if domain_id else {}),
            },
        }
        data, errors = await client.execute(_Q_ADD_ADDRESS, variables)
        return client.unwrap(data, errors, "introduceAddress")

    # ------------------------------------------------------------------

    @mcp.tool()
    async def restore_address(
        session_id: str,
        address: str,
        restore_key: str,
    ) -> dict[str, Any]:
        """Restore a previously used address into a session using its restore key.

        The address must have been obtained from create_session or add_address.
        Note: emails from the old session are permanently deleted; only the
        address itself is restored.

        Args:
            session_id:  Session ID to restore the address into.
            address:     Email address string (login@domain).
            restore_key: The restoreKey returned when the address was created.
        """
        logger.info("tool=restore_address session_id=%s address=%s", session_id, address)
        variables: dict[str, Any] = {
            "input": {
                "sessionId": session_id,
                "mailAddress": address,
                "restoreKey": restore_key,
            },
        }
        data, errors = await client.execute(_Q_RESTORE_ADDRESS, variables)
        return client.unwrap(data, errors, "restoreAddress")

    # ------------------------------------------------------------------

    @mcp.tool()
    async def delete_address(address_id: str) -> bool:
        """Remove an address from its session. The address can be restored later.

        Emails already received are not deleted — only the address is removed.

        Args:
            address_id: Address ID from create_session, add_address, or get_session.
        """
        logger.info("tool=delete_address address_id=%s", address_id)
        variables: dict[str, Any] = {"input": {"addressId": address_id}}
        data, errors = await client.execute(_Q_DELETE_ADDRESS, variables)
        return bool(client.unwrap(data, errors, "deleteAddress"))

    return mcp
