from __future__ import annotations

import asyncio
import contextlib
import logging
import time
from contextlib import AsyncExitStack
from dataclasses import dataclass, field
from typing import Any

import httpx
from mcp import ClientSession
from mcp.client.sse import sse_client
from mcp.client.stdio import StdioServerParameters, stdio_client
from mcp.client.streamable_http import streamable_http_client
from mcp.types import Tool

from .config import McpServerConfig

logger = logging.getLogger(__name__)


@dataclass
class _Cmd:
    """A command enqueued to the client's runner task."""

    op: str
    fut: asyncio.Future | None
    payload: dict[str, Any] = field(default_factory=dict)


class McpClient:
    """Lightweight wrapper around the official Python MCP client SDK.

    Manages a single server connection (stdio, SSE, or streamable-http),
    exposes ``list_tools()`` and ``call_tool()``, and handles lifecycle
    (connect / disconnect).

    **Threading model.** The MCP client SDK (``stdio_client`` /
    ``sse_client`` / ``streamable_http_client`` and ``ClientSession``) is
    built on anyio task groups, whose cancel scopes require that
    ``__aenter__`` and ``__aexit__`` run in the *same* asyncio task. To
    guarantee that, every client owns a single long-running **runner task**
    that holds the ``AsyncExitStack`` and performs all transport enter/exit
    (connect, disconnect, reconnect) as well as every ``list_tools`` /
    ``call_tool``. The public methods just enqueue a command onto a queue and
    await its result Future, so callers can invoke them from any task
    (lifespan, health-check, request handler) without crossing cancel-scope
    task boundaries.

    Reconnect uses exponential backoff (base 1s, max 30s, ±20% jitter)
    so that a flapping server does not hammer the transport.
    """

    def __init__(self, name: str, config: McpServerConfig) -> None:
        self.name = name
        self.config = config

        # Mirrored snapshot of the runner's state, updated from inside the
        # runner task. asyncio is single-threaded, so these are safe to read
        # synchronously from any task via the properties below.
        self._connected: bool = False
        self._instructions: str | None = None

        # Runner task + command queue. Lazily created on first use.
        self._runner_task: asyncio.Task | None = None
        self._queue: asyncio.Queue[_Cmd] | None = None

    # ── Public, sync state ────────────────────────────────────────────────

    @property
    def connected(self) -> bool:
        return self._connected

    @property
    def instructions(self) -> str | None:
        """Server-provided instructions from MCP initialize handshake."""
        return self._instructions

    # ── Public, async API (enqueue + await) ───────────────────────────────

    async def connect(self) -> None:
        """Open transport, initialise session. Idempotent when already connected."""
        await self._send("connect")

    async def disconnect(self) -> None:
        """Close transport, stop the runner task, and reset state."""
        runner = self._runner_task
        queue = self._queue
        # Ask the runner to stop. It closes the transport *inside its own
        # task* — the anyio cancel scope is exited where it was entered.
        if queue is not None and runner is not None and not runner.done():
            with contextlib.suppress(Exception):
                queue.put_nowait(_Cmd("disconnect", None, {}))
        if runner is not None and not runner.done():
            # Shield so a lifespan cancellation does not interrupt the
            # runner's own transport teardown; the runner finishes in its
            # own task regardless.
            try:
                await asyncio.shield(runner)
            except BaseException:
                pass
        self._runner_task = None
        self._queue = None
        self._connected = False
        self._instructions = None

    async def list_tools(self) -> list[Tool]:
        """Return the tools exposed by the remote MCP server."""
        return await self._send("list_tools")  # type: ignore[return-value]

    async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None = None) -> tuple[str, bool]:
        """Execute a remote tool and return its output as a string.

        Returns (output_text, is_error).  ``is_error`` comes from the MCP
        ``CallToolResult.isError`` field so the caller can set ``success=False``.

        Text content is concatenated; images are reported as a placeholder.
        """
        return await self._send(  # type: ignore[return-value]
            "call_tool", tool_name=tool_name, arguments=arguments
        )

    async def mark_disconnected(self) -> None:
        """Force the client into a disconnected state without closing transport.

        Used by the health-check loop when a server drops silently: the next
        operation will reconnect (with backoff).
        """
        await self._send("mark_disconnected")

    # ── Runner plumbing ───────────────────────────────────────────────────

    def _ensure_runner(self) -> None:
        """Lazily start the runner task + queue if not already running."""
        if self._runner_task is not None and not self._runner_task.done():
            return
        loop = asyncio.get_running_loop()
        self._queue = asyncio.Queue()
        self._runner_task = loop.create_task(self._runner())

    async def _send(self, op: str, **payload: Any) -> Any:
        """Enqueue a command and await its result."""
        self._ensure_runner()
        loop = asyncio.get_running_loop()
        fut: asyncio.Future = loop.create_future()
        assert self._queue is not None
        await self._queue.put(_Cmd(op, fut, payload))
        return await fut

    async def _runner(self) -> None:
        """Owns the AsyncExitStack; processes commands until ``disconnect``.

        All transport ``enter_async_context`` / ``aclose`` calls happen here,
        in this single task, so anyio cancel scopes are always entered and
        exited by the same task.
        """
        queue = self._queue
        assert queue is not None

        exit_stack = AsyncExitStack()
        session: ClientSession | None = None
        connected: bool = False
        instructions: str | None = None
        backoff: float = 1.0
        last_attempt: float | None = None

        async def close_all() -> None:
            nonlocal session, connected, instructions
            try:
                await exit_stack.aclose()
            except SystemExit:
                raise
            except BaseException:
                # anyio raises RuntimeError / BaseExceptionGroup when an
                # async generator is closed; we are tearing down anyway.
                pass
            finally:
                session = None
                connected = False
                instructions = None
                self._connected = False
                self._instructions = None

        async def open_transport() -> None:
            nonlocal session, instructions
            cfg = self.config
            try:
                if cfg.is_stdio:
                    if not cfg.command:
                        raise ValueError("stdio transport requires 'command'")
                    params = StdioServerParameters(
                        command=cfg.command,
                        args=cfg.args or [],
                        env=cfg.env,
                        cwd=cfg.cwd,
                    )
                    transport = await exit_stack.enter_async_context(stdio_client(params))
                elif cfg.is_sse:
                    if not cfg.url:
                        raise ValueError("sse transport requires 'url'")
                    transport = await exit_stack.enter_async_context(
                        sse_client(cfg.url, headers=cfg.headers)
                    )
                elif cfg.is_streamable_http:
                    if not cfg.url:
                        raise ValueError("streamable_http transport requires 'url'")
                    http_client = await exit_stack.enter_async_context(
                        httpx.AsyncClient(headers=cfg.headers or {})
                    )
                    transport = await exit_stack.enter_async_context(
                        streamable_http_client(cfg.url, http_client=http_client)
                    )
                    # streamable_http_client returns (read, write, get_session_id).
                    transport = transport[:2]
                else:
                    raise ValueError(f"unknown transport: {cfg.transport}")

                read_stream, write_stream = transport
                session = await exit_stack.enter_async_context(
                    ClientSession(read_stream, write_stream)
                )
                init_result = await session.initialize()
                instructions = getattr(init_result, "instructions", None)
            except BaseException:
                # Roll back any half-entered contexts before propagating.
                await close_all()
                raise

        def check_backoff() -> bool:
            nonlocal last_attempt
            if last_attempt is None:
                return True
            elapsed = time.monotonic() - last_attempt
            jitter = backoff * 0.2 * (2 * (time.monotonic() % 1) - 1)
            return elapsed >= (backoff + jitter)

        async def ensure_connected() -> None:
            nonlocal connected, backoff, last_attempt, instructions
            if connected and session is not None:
                return
            if not check_backoff():
                remaining = backoff - (time.monotonic() - (last_attempt or 0))
                logger.warning(
                    "MCP server %r reconnect blocked by backoff (%.1fs remaining)",
                    self.name,
                    max(0, remaining),
                )
                raise RuntimeError(
                    f"MCP server {self.name!r} is disconnected and reconnect is throttled"
                )
            last_attempt = time.monotonic()
            logger.warning("MCP server %r disconnected, reconnecting...", self.name)
            await close_all()
            await open_transport()  # raises on failure -> caller retries later
            connected = True
            backoff = 1.0
            last_attempt = None
            self._connected = True
            self._instructions = instructions
            logger.info("MCP server %r connected (%s)", self.name, self.config.transport)

        try:
            while True:
                cmd = await queue.get()
                try:
                    if cmd.op == "connect":
                        await ensure_connected()
                        self._resolve(cmd.fut, None)
                    elif cmd.op == "disconnect":
                        break  # finally -> close_all
                    elif cmd.op == "mark_disconnected":
                        connected = False
                        self._connected = False
                        self._resolve(cmd.fut, None)
                    elif cmd.op == "list_tools":
                        await ensure_connected()
                        try:
                            result = await session.list_tools()  # type: ignore[union-attr]
                        except Exception:
                            await close_all()
                            await ensure_connected()
                            result = await session.list_tools()  # type: ignore[union-attr]
                        self._resolve(cmd.fut, list(result.tools))
                    elif cmd.op == "call_tool":
                        tool_name = cmd.payload["tool_name"]
                        arguments = cmd.payload.get("arguments")
                        await ensure_connected()
                        try:
                            result = await session.call_tool(  # type: ignore[union-attr]
                                tool_name, arguments or {}
                            )
                        except Exception:
                            await close_all()
                            await ensure_connected()
                            result = await session.call_tool(  # type: ignore[union-attr]
                                tool_name, arguments or {}
                            )
                        parts: list[str] = []
                        for item in result.content:
                            if item.type == "text":
                                parts.append(item.text)
                            elif item.type == "image":
                                parts.append(f"[image: {item.mimeType} ({len(item.data)} bytes)]")
                            else:
                                parts.append(f"[{item.type}]")
                        is_error = getattr(result, "isError", False)
                        self._resolve(cmd.fut, ("\n".join(parts), is_error))
                    else:
                        self._fail(cmd.fut, ValueError(f"unknown op: {cmd.op}"))
                except asyncio.CancelledError:
                    # Runner is being cancelled (e.g. app teardown). Release
                    # the in-flight caller, then let close_all run in finally.
                    self._cancel(cmd.fut)
                    raise
                except Exception as exc:
                    # On a connect/ensure failure, mirror the disconnected
                    # state before surfacing the error to the caller.
                    self._connected = connected
                    self._instructions = instructions
                    self._fail(cmd.fut, exc)
        finally:
            await close_all()

    @staticmethod
    def _resolve(fut: asyncio.Future | None, value: Any) -> None:
        if fut is not None and not fut.done():
            fut.set_result(value)

    @staticmethod
    def _fail(fut: asyncio.Future | None, exc: BaseException) -> None:
        if fut is not None and not fut.done():
            fut.set_exception(exc)

    @staticmethod
    def _cancel(fut: asyncio.Future | None) -> None:
        if fut is not None and not fut.done():
            fut.cancel()