"""MCP streamable HTTP client for the MCPBundles Hub.

Manages session lifecycle, auth header injection, and protocol operations.
All discovery (list_tools, list_resources, list_prompts) is always live.
Session IDs are persisted only to avoid re-running the MCP initialize handshake.
"""

import json
import logging
from contextlib import asynccontextmanager
from typing import Any

import httpx
from mcp import ClientSession
from mcp.client.streamable_http import streamable_http_client
from mcp.types import CallToolResult, TextContent, ImageContent

from mcpbundles.config import AuthError

logger = logging.getLogger(__name__)


class ConnectionError(Exception):
    pass


class ToolCallError(Exception):
    def __init__(self, message: str, content: list[Any] | None = None):
        super().__init__(message)
        self.content = content


def _build_headers(auth: str, session_id: str | None = None) -> dict[str, str]:
    headers: dict[str, str] = {}
    if auth.startswith("mb_"):
        headers["X-API-Key"] = auth
    else:
        headers["Authorization"] = f"Bearer {auth}"
    if session_id:
        headers["Mcp-Session-Id"] = session_id
    return headers


def _unwrap_exception(exc: BaseException) -> BaseException:
    """Unwrap nested ExceptionGroups from anyio TaskGroups to get the root cause."""
    while isinstance(exc, BaseExceptionGroup) and exc.exceptions:
        exc = exc.exceptions[0]
    return exc


@asynccontextmanager
async def hub_session(
    auth: str,
    base_url: str,
    endpoint: str = "/hub/",
    session_id: str | None = None,
):
    """Open an MCP session to the Hub or a bundle endpoint."""
    url = f"{base_url}{endpoint}"
    headers = _build_headers(auth, session_id)
    session_ready = False

    try:
        async with httpx.AsyncClient(headers=headers, timeout=httpx.Timeout(30.0)) as http_client:
            async with streamable_http_client(url, http_client=http_client) as (read, write, _):
                async with ClientSession(read, write) as session:
                    await session.initialize()
                    logger.info("MCP session initialized: %s", url)
                    session_ready = True
                    yield session
    except BaseException as exc:
        if session_ready:
            # Session was used successfully — this is a cleanup exception
            # from MCP SDK task group teardown. Safe to suppress.
            logger.debug("Session cleanup exception (suppressed): %s", exc)
            return
        root = _unwrap_exception(exc)
        if isinstance(root, httpx.HTTPStatusError):
            if root.response.status_code in (401, 403):
                raise AuthError(f"Authentication failed ({root.response.status_code}). Run 'mcpbundles login'.") from root
            raise ConnectionError(f"Server error: {root.response.status_code}") from root
        if isinstance(root, httpx.ConnectError):
            raise ConnectionError(f"Cannot connect to {url}: {root}") from root
        from mcp.shared.exceptions import McpError as SdkMcpError
        if isinstance(root, SdkMcpError):
            raise ConnectionError(f"MCP error: {root.error.message}") from root
        if isinstance(root, (AuthError, ConnectionError, ToolCallError)):
            raise
        raise ConnectionError(f"Connection failed: {root}") from root


async def call_tool(session: ClientSession, name: str, arguments: dict[str, Any] | None = None) -> CallToolResult:
    result = await session.call_tool(name, arguments or {})
    if result.isError:
        text_parts = [c.text for c in result.content if isinstance(c, TextContent)]
        raise ToolCallError("\n".join(text_parts) if text_parts else "Tool call failed", result.content)
    return result


def extract_text(result: CallToolResult) -> str:
    parts: list[str] = []
    for block in result.content:
        if isinstance(block, TextContent):
            parts.append(block.text)
    return "\n".join(parts)


def extract_json(result: CallToolResult) -> Any:
    text = extract_text(result)
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return text
