from __future__ import annotations

import asyncio
import contextlib
import hashlib
import json
import logging
import os
import time
from contextlib import AsyncExitStack
from dataclasses import dataclass
from datetime import timedelta
from typing import Any, Callable

import httpx
import mcp.types as mcp_types
from mcp.client.session 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.shared.memory import create_connected_server_and_client_session

from comate_agent_sdk.llm.base import ToolDefinition
from comate_agent_sdk.llm.messages import ContentPartImageParam, ContentPartTextParam
from comate_agent_sdk.mcp.token_store import TokenStore
from comate_agent_sdk.mcp.types import McpServerConfig
from comate_agent_sdk.mcp.utils import sanitize_tool_name
from comate_agent_sdk.tools.decorator import Tool

logger = logging.getLogger("comate_agent_sdk.mcp.manager")


def _read_timeout_env(name: str, default: float) -> float:
    raw = os.getenv(name)
    if raw is None or not raw.strip():
        return default
    try:
        value = float(raw.strip())
    except ValueError:
        logger.warning(f"Invalid env var {name}={raw!r}; using default {default}.")
        return default
    if value <= 0:
        logger.warning(f"Invalid env var {name}={raw!r}; using default {default}.")
        return default
    return value


_MCP_TOOL_MARKER_ATTR = "_comate_agent_sdk_mcp_tool"
_MCP_TOOL_MARKER_VALUE = True
_START_TIMEOUT_S = _read_timeout_env("AGENT_SDK_MCP_START_TIMEOUT_S", 10.0)
_START_SAFETY_BUFFER_S = 2.0
_SHUTDOWN_TIMEOUT_S = _read_timeout_env("AGENT_SDK_MCP_SHUTDOWN_TIMEOUT_S", 5.0)
_SHUTDOWN_CANCEL_TIMEOUT_S = _read_timeout_env(
    "AGENT_SDK_MCP_SHUTDOWN_CANCEL_TIMEOUT_S",
    1.0,
)
_LOCAL_CONCURRENCY = int(os.getenv("AGENT_SDK_MCP_LOCAL_CONCURRENCY", "3"))
_REMOTE_CONCURRENCY = int(os.getenv("AGENT_SDK_MCP_REMOTE_CONCURRENCY", "20"))
_TOOLS_CACHE_TTL_S = _read_timeout_env("AGENT_SDK_MCP_TOOLS_CACHE_TTL_S", 300.0)
_CONN_CLOSE_TIMEOUT_S = _read_timeout_env("AGENT_SDK_MCP_CONN_CLOSE_TIMEOUT_S", 3.0)
_LOCAL_CONNECT_TIMEOUT_S = _read_timeout_env(
    "AGENT_SDK_MCP_LOCAL_CONNECT_TIMEOUT_S",
    30.0,
)
_REMOTE_CONNECT_TIMEOUT_S = _read_timeout_env(
    "AGENT_SDK_MCP_REMOTE_CONNECT_TIMEOUT_S",
    10.0,
)
# Two zero-sleep passes flush nested call_soon callbacks used by Proactor stdio teardown.
_WINDOWS_STDIO_DRAIN_PASSES = 2


@dataclass
class McpToolResult:
    """MCP 工具执行结果，显式携带 is_error 标记。"""
    content: str | list[ContentPartTextParam | ContentPartImageParam]
    is_error: bool = False

    def __str__(self) -> str:
        if isinstance(self.content, str):
            return self.content
        parts = []
        for part in self.content:
            if hasattr(part, "text"):
                parts.append(part.text)
            else:
                parts.append(str(part))
        return "\n".join(parts)


@dataclass(frozen=True)
class McpToolInfo:
    server_alias: str
    server_type: str
    remote_name: str
    mapped_name: str
    description: str
    input_schema: dict[str, Any]


@dataclass(frozen=True)
class McpServerRuntimeState:
    alias: str
    status: str
    reason: str | None
    tool_count: int
    instructions: str | None = None


@dataclass(frozen=True)
class _ServerLoadResult:
    server_alias: str
    session: ClientSession | None
    tool_infos: tuple[McpToolInfo, ...]
    tool_count: int
    failure_reason: str | None = None


@dataclass
class _CachedConnection:
    alias: str
    server_type: str
    config_hash: str
    session: ClientSession
    owner_task: asyncio.Task[None]
    shutdown_event: asyncio.Event
    get_session_id: Callable[[], str | None] | None = None


@dataclass
class _CachedToolList:
    alias: str
    tool_infos: tuple[McpToolInfo, ...]
    fetched_at: float


def is_mcp_tool(tool: Tool) -> bool:
    return getattr(tool, _MCP_TOOL_MARKER_ATTR, False) is True


class McpManager:
    """MCP tools 管理器：负责连接 server、拉取 tools、以及将其封装为 SDK Tool。"""

    def __init__(
        self,
        servers: dict[str, McpServerConfig],
        *,
        connect_timeout_s: float | None = None,
        call_timeout_s: float = 60.0,
        token_store: TokenStore | None = None,
    ) -> None:
        self._servers = dict(servers)
        self._connect_timeout_s_override = (
            float(connect_timeout_s) if connect_timeout_s is not None else None
        )
        self._call_timeout_s = float(call_timeout_s)
        self._token_store = token_store or TokenStore()

        self._tool_info_by_mapped: dict[str, McpToolInfo] = {}
        self._tools: list[Tool] = []
        self._lifecycle_task: asyncio.Task[None] | None = None
        self._shutdown_event: asyncio.Event = asyncio.Event()
        self._init_done: asyncio.Future[None] | None = None
        self._lifecycle_lock: asyncio.Lock = asyncio.Lock()

        self._conn_cache: dict[str, _CachedConnection] = {}
        self._tools_list_cache: dict[str, _CachedToolList] = {}
        self._state_lock: asyncio.Lock = asyncio.Lock()

        self._local_semaphore = asyncio.Semaphore(_LOCAL_CONCURRENCY)
        self._remote_semaphore = asyncio.Semaphore(_REMOTE_CONCURRENCY)

        self._on_tools_ready: Callable[[list[Tool]], None] | None = None
        self._all_workers_done: asyncio.Event | None = None
        self._tools_hash: str = ""
        self._server_states: dict[str, McpServerRuntimeState] = {
            alias: McpServerRuntimeState(
                alias=alias,
                status="idle",
                reason=None,
                tool_count=0,
            )
            for alias in self._servers
        }

    @property
    def tools(self) -> list[Tool]:
        return list(self._tools)

    @property
    def tool_infos(self) -> list[McpToolInfo]:
        return list(self._tool_info_by_mapped.values())

    @property
    def failed_servers(self) -> list[tuple[str, str]]:
        """Return list of (alias, friendly_reason) for servers that failed to connect."""
        failed: list[tuple[str, str]] = []
        for alias in sorted(self._server_states.keys()):
            state = self._server_states[alias]
            if state.status == "failed" and state.reason:
                failed.append((alias, state.reason))
        return failed

    @property
    def server_states(self) -> dict[str, McpServerRuntimeState]:
        return dict(self._server_states)

    @property
    def tools_hash(self) -> str:
        return self._tools_hash

    @staticmethod
    async def _drain_windows_stdio_cleanup() -> None:
        """Flush deferred Proactor pipe-close callbacks before loop shutdown."""
        if os.name != "nt":
            return
        for _ in range(_WINDOWS_STDIO_DRAIN_PASSES):
            try:
                await asyncio.sleep(0)
            except RuntimeError as exc:
                logger.debug(
                    "Skipping remaining Windows stdio drain passes: %s",
                    exc,
                )
                return

    def _compute_tools_hash(self) -> str:
        """Compute stable hash of current tool set for cache invalidation detection."""
        items = sorted(
            self._tool_info_by_mapped.values(),
            key=lambda info: info.mapped_name,
        )
        summary = [
            (info.mapped_name, info.description, info.input_schema)
            for info in items
        ]
        raw = json.dumps(summary, sort_keys=True, default=str)
        return hashlib.sha256(raw.encode()).hexdigest()[:16]

    async def _set_server_state(
        self,
        alias: str,
        *,
        status: str,
        reason: str | None = None,
        tool_count: int = 0,
    ) -> None:
        async with self._state_lock:
            self._server_states[alias] = McpServerRuntimeState(
                alias=alias,
                status=status,
                reason=reason,
                tool_count=int(tool_count),
            )

    async def _mark_connect_timeout_failed(
        self,
        alias: str,
        cfg: McpServerConfig,
        *,
        timeout_s: float,
    ) -> None:
        try:
            await asyncio.sleep(timeout_s)
        except asyncio.CancelledError:
            return

        current = self._server_states.get(alias)
        if current is None or current.status != "connecting":
            return

        await self._set_server_state(
            alias,
            status="failed",
            reason=self._friendly_error_message(alias, cfg, asyncio.TimeoutError()),
        )

    def _connect_timeout_for_cfg(self, cfg: McpServerConfig) -> float:
        if self._connect_timeout_s_override is not None:
            return self._connect_timeout_s_override
        server_type = self._server_type(cfg)
        if server_type in {"stdio", "sdk"}:
            return _LOCAL_CONNECT_TIMEOUT_S
        return _REMOTE_CONNECT_TIMEOUT_S

    @staticmethod
    def _friendly_error_message(alias: str, cfg: McpServerConfig, exc: Exception) -> str:
        """Classify common MCP connection exceptions into actionable English hints."""
        # Unwrap ExceptionGroup to find the root cause exception
        unwrapped = exc
        while isinstance(unwrapped, BaseExceptionGroup):
            if unwrapped.exceptions:
                unwrapped = unwrapped.exceptions[0]
            else:
                break
        if unwrapped is not exc and isinstance(unwrapped, Exception):
            return McpManager._friendly_error_message(alias, cfg, unwrapped)

        msg = str(exc)

        # None header value — typically an unset env var
        if "Header value" in msg and "NoneType" in msg:
            # Try to identify which header key is None
            headers = cfg.get("headers") or {}  # type: ignore[attr-defined]
            null_keys = [k for k, v in headers.items() if v is None] if isinstance(headers, dict) else []
            if null_keys:
                return f"header contains null value for {null_keys}, check related env vars"
            return "header contains null value (likely unset env var), check config"

        # Connection refused / connect error
        exc_name = type(exc).__name__
        if exc_name in ("ConnectionRefusedError", "ConnectError") or "Connection refused" in msg:
            return "cannot connect to server, verify URL and server status"

        # Timeout
        if isinstance(exc, (TimeoutError, asyncio.TimeoutError)) or "timeout" in msg.lower():
            return "connection timed out, check network or server availability"

        # stdio command not found
        if isinstance(exc, FileNotFoundError):
            command = cfg.get("command", "")  # type: ignore[attr-defined]
            return f"command '{command}' not found, verify it is installed and in PATH"

        # HTTP 404 during initialize → upstream SDK reports "Session terminated"
        # which is misleading; the real cause is almost always a wrong URL path.
        if exc_name == "McpError" and "Session terminated" in msg:
            url = cfg.get("url", "")  # type: ignore[attr-defined]
            return f"server returned 404, verify the URL path is correct (current: {url})"

        from comate_agent_sdk.mcp.oauth import (
            AuthorizationError,
            CallbackStateMismatchError,
            CallbackTimeoutError,
            McpOAuthError,
            MetadataDiscoveryError,
            NoBrowserError,
            RefreshFailedError,
            TokenExchangeError,
            TokenStorageError,
            _redact_value,
        )

        if isinstance(exc, McpOAuthError):
            redacted_msg = _redact_value(msg)
            if isinstance(exc, NoBrowserError):
                return "OAuth requires a browser for first-time authorization"
            if isinstance(exc, CallbackTimeoutError):
                return "OAuth authorization timed out before callback completed"
            if isinstance(exc, CallbackStateMismatchError):
                return "OAuth callback rejected due to state mismatch"
            if isinstance(exc, MetadataDiscoveryError):
                return "OAuth metadata discovery failed; verify issuer or explicit endpoints"
            if isinstance(exc, AuthorizationError):
                return "OAuth authorization denied by server"
            if isinstance(exc, TokenExchangeError):
                return f"OAuth token exchange failed: {redacted_msg}"
            if isinstance(exc, RefreshFailedError):
                return f"OAuth token refresh failed; run retry_server('{alias}') to re-authorize"
            if isinstance(exc, TokenStorageError):
                return f"OAuth token storage failed: {redacted_msg}"

        # ValueError from our own validation (already friendly)
        if isinstance(exc, ValueError):
            return msg

        # Fallback: type + message
        return f"{exc_name}: {msg}"

    @staticmethod
    def _is_session_expired(exc: Exception) -> bool:
        """Detect session expiry signals that warrant reconnection."""
        if isinstance(exc, (ConnectionResetError, BrokenPipeError)):
            return True
        msg = str(exc)
        exc_name = type(exc).__name__
        if exc_name == "McpError" and "Session terminated" in msg:
            return True
        if "-32001" in msg:
            return True
        return False

    async def start(
        self,
        on_tools_ready: Callable[[list[Tool]], None] | None = None,
    ) -> None:
        async with self._lifecycle_lock:
            if self._lifecycle_task is not None and not self._lifecycle_task.done():
                init_done = self._init_done
                if init_done is None:
                    init_done = asyncio.get_running_loop().create_future()
                    self._init_done = init_done
            else:
                # 若存在已结束的 task，清理引用后按冷启动处理。
                # 不清空 _tool_info_by_mapped 和 _tools，由 _run_lifecycle 按需清理。
                self._on_tools_ready = on_tools_ready
                self._lifecycle_task = None
                self._shutdown_event = asyncio.Event()
                self._all_workers_done = asyncio.Event()
                init_done = asyncio.get_running_loop().create_future()
                self._init_done = init_done
                self._lifecycle_task = asyncio.create_task(
                    self._run_lifecycle(),
                    name=f"mcp_lifecycle_{id(self)}",
                )

        _safety_timeout = _START_TIMEOUT_S + _START_SAFETY_BUFFER_S
        try:
            await asyncio.wait_for(
                asyncio.shield(init_done), timeout=_safety_timeout
            )
        except asyncio.TimeoutError:
            logger.error(
                f"MCP lifecycle task itself timed out (safety net: {_safety_timeout:.1f}s)"
            )
            await self.aclose()
            raise
        except Exception:
            await self.aclose()
            raise

    async def start_and_wait(self) -> None:
        """Start and block until all servers are connected or failed."""
        await self.start()
        if self._all_workers_done is not None:
            await self._all_workers_done.wait()

    async def aclose(self) -> None:
        async with self._lifecycle_lock:
            lifecycle_task = self._lifecycle_task
            if lifecycle_task is None:
                return
            self._shutdown_event.set()

        try:
            await asyncio.wait_for(
                asyncio.shield(lifecycle_task),
                timeout=_SHUTDOWN_TIMEOUT_S,
            )
        except asyncio.TimeoutError:
            logger.debug(f"MCP lifecycle task cleanup timed out ({_SHUTDOWN_TIMEOUT_S:.1f}s), force cancelling")
            lifecycle_task.cancel()
            done, _pending = await asyncio.wait(
                {lifecycle_task},
                timeout=_SHUTDOWN_CANCEL_TIMEOUT_S,
            )
            if not done:
                logger.debug(
                    "MCP lifecycle task did not exit after cancellation within "
                    f"{_SHUTDOWN_CANCEL_TIMEOUT_S:.1f}s; forcing bounded return",
                )
            else:
                try:
                    await lifecycle_task
                except asyncio.CancelledError:
                    pass
                except Exception as e:
                    logger.debug(f"MCP lifecycle task error after cancellation: {e}", exc_info=True)
        except Exception as e:
            logger.debug(f"MCP lifecycle task error: {e}", exc_info=True)
        finally:
            async with self._lifecycle_lock:
                if self._lifecycle_task is lifecycle_task:
                    self._lifecycle_task = None
                self._init_done = None
                for alias in list(self._conn_cache):
                    await self._close_connection(alias)
                self._tool_info_by_mapped.clear()
                self._tools = []

    async def call_tool(
        self,
        mapped_name: str,
        arguments: dict[str, Any],
        cancel_event: asyncio.Event | None = None,
    ) -> McpToolResult:
        info = self._tool_info_by_mapped.get(mapped_name)
        if info is None:
            return McpToolResult(content=f"Unknown MCP tool '{mapped_name}'", is_error=True)

        cached = self._conn_cache.get(info.server_alias)
        if cached is None:
            return McpToolResult(content=f"MCP server '{info.server_alias}' is not connected", is_error=True)

        if cancel_event is not None and cancel_event.is_set():
            return McpToolResult(
                content=json.dumps({"status": "cancelled", "reason": "user_interrupt"}, ensure_ascii=False),
                is_error=True,
            )

        try:
            if cancel_event is not None:
                call_task = asyncio.create_task(
                    cached.session.call_tool(
                        info.remote_name,
                        arguments=arguments,
                        read_timeout_seconds=timedelta(seconds=self._call_timeout_s),
                    )
                )
                cancel_task = asyncio.create_task(cancel_event.wait())
                done, pending = await asyncio.wait(
                    [call_task, cancel_task],
                    return_when=asyncio.FIRST_COMPLETED,
                )
                for p in pending:
                    p.cancel()
                    with contextlib.suppress(asyncio.CancelledError):
                        await p
                if cancel_task in done:
                    return McpToolResult(
                        content=json.dumps({"status": "cancelled", "reason": "user_interrupt"}, ensure_ascii=False),
                        is_error=True,
                    )
                result = call_task.result()
            else:
                result = await cached.session.call_tool(
                    info.remote_name,
                    arguments=arguments,
                    read_timeout_seconds=timedelta(seconds=self._call_timeout_s),
                )
        except Exception as e:
            if isinstance(e, httpx.HTTPStatusError) and e.response.status_code == 401:
                return McpToolResult(
                    content=(
                        f"MCP tool call failed with 401; session may need re-authorization "
                        f"via retry_server('{info.server_alias}')"
                    ),
                    is_error=True,
                )
            if not self._is_session_expired(e):
                return McpToolResult(content=f"MCP tool call failed: {type(e).__name__}: {e}", is_error=True)

            # Attempt single reconnect
            cfg = self._servers.get(info.server_alias)
            if cfg is None:
                return McpToolResult(content=f"MCP server '{info.server_alias}' config not found for reconnect", is_error=True)

            logger.info(f"MCP server '{info.server_alias}' session expired, attempting reconnect...")
            await self._close_connection(info.server_alias)
            try:
                # spec N9：session-expired 重连忽略 instructions，
                # 不覆写 _server_states[alias].instructions 避免"老用法"被意外刷空。
                _, _ = await self._open_connection(info.server_alias, cfg)
                new_cached = self._conn_cache.get(info.server_alias)
                if new_cached is None:
                    return McpToolResult(content="MCP reconnect succeeded but session not found", is_error=True)
                result = await new_cached.session.call_tool(
                    info.remote_name,
                    arguments=arguments,
                    read_timeout_seconds=timedelta(seconds=self._call_timeout_s),
                )
            except Exception as reconnect_err:
                return McpToolResult(content=f"MCP reconnect failed: {type(reconnect_err).__name__}: {reconnect_err}", is_error=True)

        return self._convert_call_tool_result(result)

    async def retry_server(self, alias: str) -> bool:
        """Manually retry a failed server. Returns True if successful."""
        if alias not in self._servers:
            raise ValueError(f"Unknown MCP server alias: '{alias}'")

        cfg = self._servers[alias]
        await self._close_connection(alias)
        self._tools_list_cache.pop(alias, None)
        await self._set_server_state(alias, status="connecting")

        sem = self._get_semaphore(cfg)
        connect_timeout_s = self._connect_timeout_for_cfg(cfg)
        async with sem:
            try:
                timeout_marker = asyncio.create_task(
                    self._mark_connect_timeout_failed(
                        alias,
                        cfg,
                        timeout_s=connect_timeout_s,
                    )
                )
                try:
                    session, instructions = await asyncio.wait_for(
                        self._open_connection(alias, cfg),
                        timeout=connect_timeout_s,
                    )
                finally:
                    if not timeout_marker.done():
                        timeout_marker.cancel()
                        with contextlib.suppress(asyncio.CancelledError):
                            await timeout_marker
                tool_list = await asyncio.wait_for(
                    self._list_all_tools(session),
                    timeout=connect_timeout_s,
                )
            except asyncio.CancelledError:
                await self._set_server_state(alias, status="failed", reason="cancelled")
                raise
            except Exception as exc:
                friendly = self._friendly_error_message(alias, cfg, exc)
                await self._set_server_state(
                    alias,
                    status="failed",
                    reason=friendly,
                )
                logger.warning(f"MCP server '{alias}' retry failed: {friendly}")
                return False

        tool_infos = tuple(
            McpToolInfo(
                server_alias=alias, server_type=self._server_type(cfg),
                remote_name=t.name,
                mapped_name=self._map_tool_name(alias, t.name),
                description=(t.description or "").strip(),
                input_schema=self._normalize_input_schema(t.inputSchema),
            )
            for t in tool_list
        )

        async with self._state_lock:
            for info in tool_infos:
                self._tool_info_by_mapped[info.mapped_name] = info
            self._tools_list_cache[alias] = _CachedToolList(
                alias=alias, tool_infos=tool_infos, fetched_at=time.monotonic(),
            )
            new_tools = [self._build_tool(info.mapped_name, info) for info in tool_infos]
            self._tools = [
                self._build_tool(m, i) for m, i in self._tool_info_by_mapped.items()
            ]
            self._tools_hash = self._compute_tools_hash()
            self._server_states[alias] = McpServerRuntimeState(
                alias=alias,
                status="connected",
                reason=None,
                tool_count=len(tool_infos),
                instructions=instructions,
            )

        if self._on_tools_ready is not None:
            self._on_tools_ready(new_tools)

        logger.info(f"MCP server '{alias}' retry succeeded ({len(tool_infos)} tools)")
        return True

    def build_server_instructions_text(self) -> str:
        """按 server 分组渲染 MCP server instructions（对齐 Claude Code 的
        `# MCP Server Instructions` section）。

        仅包含 status=="connected" 且 instructions 非空的 server；无任何 server
        满足条件时返回空串（上游 runtime_mcp 应据此移除 MCP_TOOL session_state item）。

        外壳保留 `<mcp_tools>` XML 以维持 ContextIR session_state marker 兼容。

        并发语义：沿用 build_overview_text 无锁读模式。_server_states 的每项
        构造都是 McpServerRuntimeState frozen dataclass 的整体替换（而非就地
        修改），字典遍历读到的状态永远是一致快照。

        spec §4.3.2
        """
        blocks: list[str] = []
        for alias in sorted(self._server_states.keys()):
            state = self._server_states[alias]
            if state.status != "connected":
                continue
            if not state.instructions:
                continue
            blocks.append(f"## {alias}\n{state.instructions}")

        if not blocks:
            return ""

        body = (
            "# MCP Server Instructions\n\n"
            "The following MCP servers have provided instructions for how to "
            "use their tools and resources:\n\n"
            + "\n\n".join(blocks)
        )
        return f"<mcp_tools>\n{body}\n</mcp_tools>"

    def build_server_instructions_metadata(self) -> dict[str, Any]:
        """机器可读的 server instructions 结构，保存在 ContextItem.metadata，不进 prompt。

        spec §4.3.3
        """
        servers = [
            {
                "alias": alias,
                "status": state.status,
                "tool_count": state.tool_count,
                "has_instructions": bool(state.instructions),
                "instructions": state.instructions,
            }
            for alias, state in sorted(self._server_states.items())
        ]
        return {"servers": servers}

    async def refresh_tools_if_stale(self) -> bool:
        """Check TTL, refresh expired tool lists using cached connections. Returns True if refreshed."""
        now = time.monotonic()
        refreshed = False
        for alias, cached_tools in list(self._tools_list_cache.items()):
            if now - cached_tools.fetched_at < _TOOLS_CACHE_TTL_S:
                continue
            cached_conn = self._conn_cache.get(alias)
            if cached_conn is None:
                continue
            cfg = self._servers.get(alias, {})
            connect_timeout_s = self._connect_timeout_for_cfg(cfg)
            try:
                tool_list = await asyncio.wait_for(
                    self._list_all_tools(cached_conn.session),
                    timeout=connect_timeout_s,
                )
                new_infos = tuple(
                    McpToolInfo(
                        server_alias=alias, server_type=self._server_type(cfg),
                        remote_name=t.name,
                        mapped_name=self._map_tool_name(alias, t.name),
                        description=(t.description or "").strip(),
                        input_schema=self._normalize_input_schema(t.inputSchema),
                    )
                    for t in tool_list
                )
                self._tools_list_cache[alias] = _CachedToolList(
                    alias=alias, tool_infos=new_infos, fetched_at=time.monotonic(),
                )
                async with self._state_lock:
                    for info in new_infos:
                        self._tool_info_by_mapped[info.mapped_name] = info
                refreshed = True
                logger.info(f"MCP server '{alias}' tools refreshed ({len(new_infos)} tools)")
            except Exception as e:
                logger.warning(f"MCP server '{alias}' tools refresh failed: {e}")
        if refreshed:
            async with self._state_lock:
                self._tools = [
                    self._build_tool(m, i) for m, i in self._tool_info_by_mapped.items()
                ]
                self._tools_hash = self._compute_tools_hash()
        return refreshed

    # ===== internal helpers =====

    @staticmethod
    def _compute_config_hash(cfg: McpServerConfig) -> str:
        server_type = cfg.get("type", "stdio")
        if server_type == "sdk":
            instance = cfg.get("instance")
            raw = f"sdk:{id(instance)}"
        else:
            serializable = {k: v for k, v in cfg.items() if k != "instance"}
            if str(serializable.get("type", "stdio")).strip().lower() == "streamable-http":
                serializable["type"] = "http"
            raw = json.dumps(serializable, sort_keys=True, default=str)
        return hashlib.sha256(raw.encode()).hexdigest()[:16]

    async def _open_connection(
        self, alias: str, cfg: McpServerConfig
    ) -> tuple[ClientSession, str | None]:
        server_type = self._server_type(cfg)
        config_hash = self._compute_config_hash(cfg)
        shutdown_event = asyncio.Event()
        session_future: asyncio.Future[
            tuple[ClientSession, Callable[[], str | None] | None, str | None]
        ] = (
            asyncio.get_running_loop().create_future()
        )

        async def _connection_owner() -> None:
            """Owner task: enter and exit MCP transport context in the same task."""
            stack = AsyncExitStack()
            try:
                session, get_session_id, instructions = await self._create_session(
                    stack, alias, cfg, server_type
                )
                if not session_future.done():
                    session_future.set_result((session, get_session_id, instructions))
                await shutdown_event.wait()
            except BaseException as exc:
                if not session_future.done():
                    session_future.set_exception(exc)
                elif not isinstance(exc, asyncio.CancelledError):
                    logger.warning(f"MCP connection '{alias}' died unexpectedly: {exc}")
            finally:
                try:
                    await stack.aclose()
                except Exception as close_exc:
                    logger.debug(f"Error closing stack for '{alias}': {close_exc}")

        owner_task = asyncio.create_task(
            _connection_owner(),
            name=f"mcp_conn_owner_{alias}_{id(self)}",
        )

        try:
            session, get_session_id, instructions = await session_future
        except BaseException:
            shutdown_event.set()
            with contextlib.suppress(Exception):
                await asyncio.wait_for(owner_task, timeout=_SHUTDOWN_CANCEL_TIMEOUT_S)
            if not owner_task.done():
                owner_task.cancel()
            raise

        self._conn_cache[alias] = _CachedConnection(
            alias=alias,
            server_type=server_type,
            config_hash=config_hash,
            session=session,
            owner_task=owner_task,
            shutdown_event=shutdown_event,
            get_session_id=get_session_id,
        )
        return session, instructions

    async def _create_session(
        self,
        stack: AsyncExitStack,
        alias: str,
        cfg: McpServerConfig,
        server_type: str,
    ) -> tuple[ClientSession, Callable[[], str | None] | None, str | None]:
        """在 exit_stack 中建立 transport 并返回 (session, get_session_id, instructions)。

        spec §4.2.1-§4.2.2：
        - sdk 分支 ctx manager 吞掉 InitializeResult，instructions=None（N8）
        - stdio/sse/http 分支接住 init_result.instructions
        """
        connect_timeout_s = self._connect_timeout_for_cfg(cfg)
        if server_type == "sdk":
            instance = cfg.get("instance")  # type: ignore[attr-defined]
            if instance is None:
                raise ValueError("sdk server 缺少 instance")
            session = await stack.enter_async_context(
                create_connected_server_and_client_session(instance)
            )
            return session, None, None

        if server_type == "stdio":
            command = cfg.get("command")  # type: ignore[attr-defined]
            if not isinstance(command, str) or not command.strip():
                raise ValueError("stdio server 缺少 command")
            args = cfg.get("args") or []  # type: ignore[attr-defined]
            env = cfg.get("env")  # type: ignore[attr-defined]
            params = StdioServerParameters(command=command, args=list(args), env=env)
            errlog = stack.enter_context(open(os.devnull, "w", encoding="utf-8"))
            read_stream, write_stream = await stack.enter_async_context(
                stdio_client(params, errlog=errlog)
            )
            session = await stack.enter_async_context(
                ClientSession(
                    read_stream,
                    write_stream,
                    read_timeout_seconds=timedelta(seconds=connect_timeout_s),
                )
            )
            init_result = await session.initialize()
            return session, None, init_result.instructions

        if server_type == "sse":
            url = cfg.get("url")  # type: ignore[attr-defined]
            if not isinstance(url, str) or not url.strip():
                raise ValueError("sse server 缺少 url")
            headers = cfg.get("headers")  # type: ignore[attr-defined]
            read_stream, write_stream = await stack.enter_async_context(
                sse_client(
                    url=url,
                    headers=headers,
                    timeout=connect_timeout_s,
                )
            )
            session = await stack.enter_async_context(
                ClientSession(
                    read_stream,
                    write_stream,
                    read_timeout_seconds=timedelta(seconds=connect_timeout_s),
                )
            )
            init_result = await session.initialize()
            return session, None, init_result.instructions

        if server_type == "http":
            url = cfg.get("url")  # type: ignore[attr-defined]
            if not isinstance(url, str) or not url.strip():
                raise ValueError("http server 缺少 url")

            from comate_agent_sdk.mcp.utils import expand_env_in_value

            raw_headers = cfg.get("headers") or {}  # type: ignore[attr-defined]
            expanded_headers: dict[str, str] = {}
            missing_all: list[str] = []
            for key, value in raw_headers.items():
                if isinstance(value, str):
                    expanded, missing = expand_env_in_value(value)
                    expanded_headers[key] = expanded
                    missing_all.extend(missing)
                else:
                    expanded_headers[key] = value
            if missing_all:
                logger.warning(
                    f"MCP server '{alias}' headers reference unset env vars: {missing_all}"
                )

            auth_provider: httpx.Auth | None = None
            oauth_cfg = cfg.get("oauth")  # type: ignore[attr-defined]
            if oauth_cfg:
                from comate_agent_sdk.mcp.oauth import build_auth as oauth_build_auth

                auth_provider = await oauth_build_auth(
                    alias,
                    oauth_cfg,
                    token_store=self._token_store,
                )

            http_client = await stack.enter_async_context(
                httpx.AsyncClient(
                    headers=expanded_headers or None,
                    timeout=connect_timeout_s,
                    auth=auth_provider,
                )
            )
            read_stream, write_stream, get_session_id = await stack.enter_async_context(
                streamable_http_client(url, http_client=http_client)
            )
            session = await stack.enter_async_context(
                ClientSession(
                    read_stream,
                    write_stream,
                    read_timeout_seconds=timedelta(seconds=connect_timeout_s),
                )
            )
            init_result = await session.initialize()
            return session, get_session_id, init_result.instructions

        raise ValueError(f"Unsupported MCP server type: {server_type} (alias={alias})")

    async def _close_connection(self, alias: str) -> None:
        cached = self._conn_cache.pop(alias, None)
        if cached is None:
            return

        # Signal the owner task to shut down
        cached.shutdown_event.set()

        # Wait for owner task to gracefully exit (it will aclose the stack internally)
        try:
            await asyncio.wait_for(cached.owner_task, timeout=_CONN_CLOSE_TIMEOUT_S)
        except asyncio.TimeoutError:
            logger.warning(f"MCP connection '{alias}' owner task timed out, force cancelling")
            cached.owner_task.cancel()
            with contextlib.suppress(asyncio.CancelledError, Exception):
                await asyncio.wait_for(
                    cached.owner_task, timeout=_SHUTDOWN_CANCEL_TIMEOUT_S
                )
        except asyncio.CancelledError:
            raise
        except Exception as e:
            logger.debug(f"Error closing connection for '{alias}': {e}")
        finally:
            if cached.server_type == "stdio":
                await self._drain_windows_stdio_cleanup()

    def _server_type(self, cfg: McpServerConfig) -> str:
        t = cfg.get("type")  # type: ignore[attr-defined]
        if t is None:
            return "stdio"
        normalized = str(t).strip().lower()
        if normalized == "streamable-http":
            return "http"
        return normalized

    def _get_semaphore(self, cfg: McpServerConfig) -> asyncio.Semaphore:
        server_type = self._server_type(cfg)
        if server_type in ("stdio", "sdk"):
            return self._local_semaphore
        return self._remote_semaphore

    def _map_tool_name(self, server_alias: str, remote_tool_name: str) -> str:
        alias = sanitize_tool_name(server_alias)
        tool = sanitize_tool_name(remote_tool_name)
        return f"mcp__{alias}__{tool}"

    def _normalize_input_schema(self, schema: Any) -> dict[str, Any]:
        if isinstance(schema, dict):
            # FastMCP / MCP server 往往会返回标准 JSON Schema（type=object + properties）。
            # 这里尽量原样保留，避免与 server 的实际校验/期望不一致。
            if schema.get("type") == "object" or "properties" in schema:
                return schema

        # 最小修复：无 schema 或非 object -> 空 object
        return {"type": "object", "properties": {}, "required": []}

    def _build_tool(self, mapped_name: str, info: McpToolInfo) -> Tool:
        async def _handler(**kwargs: Any) -> Any:
            # Extract cancel_event if injected by tool execution layer
            cancel_event = kwargs.pop("_cancel_event", None)
            return await self.call_tool(mapped_name, kwargs, cancel_event=cancel_event)

        tool = Tool(func=_handler, description=info.description or "", name=mapped_name, ephemeral=False)
        # 标记来源，方便刷新/移除
        setattr(tool, _MCP_TOOL_MARKER_ATTR, _MCP_TOOL_MARKER_VALUE)

        tool._definition = ToolDefinition(  # type: ignore[attr-defined]
            name=mapped_name,
            description=info.description or "",
            parameters=info.input_schema,
            strict=False,
        )
        return tool

    async def _list_all_tools(self, session: ClientSession) -> list[mcp_types.Tool]:
        tools: list[mcp_types.Tool] = []
        cursor: str | None = None

        while True:
            result = await session.list_tools(cursor=cursor)
            tools.extend(list(result.tools or []))
            cursor = result.nextCursor
            if not cursor:
                break

        return tools

    def _convert_call_tool_result(
        self, result: mcp_types.CallToolResult
    ) -> McpToolResult:
        is_error = bool(result.isError)

        if is_error:
            text = self._render_mcp_content_as_text(result.content)
            return McpToolResult(
                content=text or "MCP tool returned error",
                is_error=True,
            )

        from comate_agent_sdk.mcp.content_convert import convert_content_blocks
        content = convert_content_blocks(result.content or [])
        return McpToolResult(content=content, is_error=False)

    def _render_mcp_content_as_text(self, content: list[mcp_types.Content] | None) -> str:
        if not content:
            return ""
        texts: list[str] = []
        for item in content:
            if isinstance(item, mcp_types.TextContent):
                texts.append(item.text)
                continue
            try:
                texts.append(item.model_dump_json())
            except Exception:
                texts.append(str(item))
        return "\n".join(texts).strip()

    async def _run_lifecycle(self) -> None:
        """Lifecycle management task: ensure exit stack is created and cleaned up in the same task."""
        init_done = self._init_done
        worker_tasks: list[asyncio.Task[None]] = []
        ready_signals: dict[str, asyncio.Future[bool]] = {}
        success_servers: list[str] = []

        try:
            # Clean up servers that were removed from config
            for alias in list(self._conn_cache):
                if alias not in self._servers:
                    await self._close_connection(alias)
                    async with self._state_lock:
                        self._tool_info_by_mapped = {
                            k: v for k, v in self._tool_info_by_mapped.items()
                            if v.server_alias != alias
                        }
                        self._server_states.pop(alias, None)
                    self._tools_list_cache.pop(alias, None)

            # Check cache and start workers for uncached servers
            for alias, cfg in self._servers.items():
                async with self._state_lock:
                    if alias not in self._server_states:
                        self._server_states[alias] = McpServerRuntimeState(
                            alias=alias,
                            status="idle",
                            reason=None,
                            tool_count=0,
                        )
                cached = self._conn_cache.get(alias)
                if cached is not None:
                    expected_hash = self._compute_config_hash(cfg)
                    if cached.config_hash == expected_hash:
                        cached_tools = self._tools_list_cache.get(alias)
                        if cached_tools is not None:
                            # Cache hit — tools already in _tool_info_by_mapped from previous start
                            # Re-insert to be safe (idempotent)
                            async with self._state_lock:
                                for info in cached_tools.tool_infos:
                                    self._tool_info_by_mapped[info.mapped_name] = info
                                # spec §4.2.4：cache-hit 重建 state 必须继承旧 instructions，
                                # 否则 McpServerRuntimeState 默认 None 会覆写，导致
                                # prompt 里该 server 的 # MCP Server Instructions section 凭空消失。
                                prev = self._server_states.get(alias)
                                self._server_states[alias] = McpServerRuntimeState(
                                    alias=alias,
                                    status="connected",
                                    reason=None,
                                    tool_count=len(cached_tools.tool_infos),
                                    instructions=prev.instructions if prev else None,
                                )
                            logger.debug(f"MCP server '{alias}' using cached connection")
                            success_servers.append(alias)
                            continue
                    else:
                        # Config changed — close old connection
                        await self._close_connection(alias)
                        self._tools_list_cache.pop(alias, None)
                        async with self._state_lock:
                            self._tool_info_by_mapped = {
                                k: v for k, v in self._tool_info_by_mapped.items()
                                if v.server_alias != alias
                            }
                            self._server_states[alias] = McpServerRuntimeState(
                                alias=alias,
                                status="idle",
                                reason=None,
                                tool_count=0,
                            )

                ready = asyncio.get_running_loop().create_future()
                ready_signals[alias] = ready
                worker_tasks.append(asyncio.create_task(
                    self._run_server_worker(alias, cfg, ready),
                    name=f"mcp_server_{alias}_{id(self)}",
                ))

            # Collection window: wait for workers within timeout
            if ready_signals:
                pending = set(ready_signals.values())
                deadline = time.monotonic() + _START_TIMEOUT_S
                while pending:
                    remaining = deadline - time.monotonic()
                    if remaining <= 0:
                        break
                    done_set, pending = await asyncio.wait(
                        pending, timeout=remaining,
                        return_when=asyncio.FIRST_COMPLETED,
                    )

                # Workers that haven't completed within the window will continue
                # running in the background — they will process their own results
                # when done (updating tools/state in-place).

            # Rebuild tools list from current state
            async with self._state_lock:
                self._tools = [
                    self._build_tool(m, i)
                    for m, i in self._tool_info_by_mapped.items()
                ]
                self._tools_hash = self._compute_tools_hash()

            # Log warnings
            if self._servers and not self._tools and not success_servers:
                failed = self.failed_servers
                logger.warning(
                    f"MCP configured {len(self._servers)} server(s) but no tools were loaded. "
                    f"succeeded: {success_servers}, failed: {failed}"
                )

            # Signal start() to return
            if init_done is not None and not init_done.done():
                init_done.set_result(None)

            # Wait for ALL remaining workers (background completion)
            if worker_tasks:
                remaining_tasks = [t for t in worker_tasks if not t.done()]
                if remaining_tasks:
                    await asyncio.wait(remaining_tasks)

        except asyncio.CancelledError:
            raise
        except Exception as e:
            logger.debug(f"MCP lifecycle task error: {e}", exc_info=True)
            if init_done is not None and not init_done.done():
                init_done.set_exception(e)
            raise
        finally:
            for task in worker_tasks:
                if not task.done():
                    task.cancel()
            if worker_tasks:
                pending_tasks = [t for t in worker_tasks if not t.done()]
                if pending_tasks:
                    await asyncio.wait(pending_tasks, timeout=_SHUTDOWN_CANCEL_TIMEOUT_S)
                for task in worker_tasks:
                    if task.done() and not task.cancelled():
                        try:
                            task.result()
                        except Exception:
                            pass

            if init_done is not None and not init_done.done():
                init_done.set_result(None)

            if self._all_workers_done is not None:
                self._all_workers_done.set()

    async def _run_server_worker(
        self,
        alias: str,
        cfg: McpServerConfig,
        ready: asyncio.Future[bool],
    ) -> None:
        sem = self._get_semaphore(cfg)
        connect_timeout_s = self._connect_timeout_for_cfg(cfg)
        async with sem:
            await self._set_server_state(alias, status="connecting")
            try:
                timeout_marker = asyncio.create_task(
                    self._mark_connect_timeout_failed(
                        alias,
                        cfg,
                        timeout_s=connect_timeout_s,
                    )
                )
                try:
                    session, instructions = await asyncio.wait_for(
                        self._open_connection(alias, cfg),
                        timeout=connect_timeout_s,
                    )
                finally:
                    if not timeout_marker.done():
                        timeout_marker.cancel()
                        with contextlib.suppress(asyncio.CancelledError):
                            await timeout_marker
                try:
                    tool_list = await asyncio.wait_for(
                        self._list_all_tools(session),
                        timeout=connect_timeout_s,
                    )
                except asyncio.TimeoutError:
                    await self._close_connection(alias)
                    friendly = "list_tools timed out"
                    await self._set_server_state(
                        alias,
                        status="failed",
                        reason=friendly,
                    )
                    if not ready.done():
                        ready.set_result(False)
                    return

                # Build tool infos and process results internally
                tool_infos = tuple(
                    McpToolInfo(
                        server_alias=alias,
                        server_type=self._server_type(cfg),
                        remote_name=tool.name,
                        mapped_name=self._map_tool_name(alias, tool.name),
                        description=(tool.description or "").strip(),
                        input_schema=self._normalize_input_schema(tool.inputSchema),
                    )
                    for tool in tool_list
                )
                new_tools: list[Tool] = []

                async with self._state_lock:
                    for info in tool_infos:
                        self._tool_info_by_mapped[info.mapped_name] = info
                    self._tools_list_cache[alias] = _CachedToolList(
                        alias=alias, tool_infos=tool_infos, fetched_at=time.monotonic(),
                    )
                    new_tools = [self._build_tool(info.mapped_name, info) for info in tool_infos]
                    self._tools = [
                        self._build_tool(m, i) for m, i in self._tool_info_by_mapped.items()
                    ]
                    self._tools_hash = self._compute_tools_hash()
                    self._server_states[alias] = McpServerRuntimeState(
                        alias=alias,
                        status="connected",
                        reason=None,
                        tool_count=len(tool_infos),
                        instructions=instructions,
                    )

                if self._on_tools_ready is not None and new_tools:
                    try:
                        self._on_tools_ready(new_tools)
                    except Exception as e:
                        logger.warning(f"on_tools_ready callback error: {e}")

                logger.info(f"MCP server loaded: {alias} ({len(tool_infos)} tools)")
                if not ready.done():
                    ready.set_result(True)
                return

            except asyncio.CancelledError:
                await self._set_server_state(alias, status="failed", reason="cancelled")
                if not ready.done():
                    ready.cancel()
                raise
            except Exception as exc:
                await self._close_connection(alias)
                friendly = self._friendly_error_message(alias, cfg, exc)
                await self._set_server_state(
                    alias,
                    status="failed",
                    reason=friendly,
                )
                if not ready.done():
                    ready.set_result(False)
                logger.debug(f"MCP server '{alias}' raw error", exc_info=True)
