"""管理员托管的 Streamable HTTP MCP 外部系统连接器。"""

from __future__ import annotations

import base64
import binascii
import hashlib
import json
import re
import time
from dataclasses import dataclass, field
from functools import partial
from typing import Any, cast
from urllib.parse import urlparse

import anyio
import httpx
import httpx2
from mcp import ClientSession
from mcp.client.streamable_http import streamable_http_client
from mcp.types import PaginatedRequestParams

from .auth import ExternalAuthError, get_auth_strategy
from .results import MAX_STORED_RESULT_BYTES
from .runtime_cache import RUNTIME_CACHE

_AUTH_CACHE_TTL_SECONDS = 300.0
_CATALOG_CACHE_TTL_SECONDS = 300.0
_MAX_CATALOG_TOOLS = 1000
_SENSITIVE_KEY_RE = re.compile(
    r"(?:password|passwd|secret|token|api[_-]?key|authorization|cookie|credential)",
    re.IGNORECASE,
)


class McpConnectorError(RuntimeError):
    pass


def _bool_value(value: Any, default: bool) -> bool:
    if value is None:
        return default
    if isinstance(value, bool):
        return value
    return str(value).strip().lower() not in {"0", "false", "no", "off"}


def _validated_http_url(raw: str, label: str) -> str:
    value = (raw or "").strip().rstrip("/")
    parsed = urlparse(value)
    if parsed.scheme not in {"http", "https"} or not parsed.netloc:
        raise McpConnectorError(f"{label} 必须是有效的 http(s) URL")
    if parsed.username or parsed.password:
        raise McpConnectorError(f"{label} 不能内嵌凭据")
    return value


@dataclass(frozen=True)
class McpConfig:
    base_url: str
    mcp_url: str
    login_path: str
    timeout_seconds: float
    max_result_bytes: int
    max_total_result_bytes: int
    verify_tls: bool
    query_guidance: str
    recommended_operation_ids: tuple[str, ...]
    operation_mode: str = "upstream_managed"
    operation_policies: dict[str, str] = field(default_factory=dict)
    auth_type: str = "password_jwt"
    auth_config: dict[str, Any] = field(default_factory=dict)
    expected_server_name: str = ""
    max_response_bytes: int = MAX_STORED_RESULT_BYTES

    @classmethod
    def from_mapping(cls, data: dict[str, Any]) -> McpConfig:
        mcp_url = _validated_http_url(str(data.get("mcp_url") or ""), "mcp_url")
        parsed = urlparse(mcp_url)
        origin = f"{parsed.scheme}://{parsed.netloc}"
        base_url = _validated_http_url(str(data.get("base_url") or origin), "base_url")
        base = urlparse(base_url)
        if (base.scheme, base.netloc) != (parsed.scheme, parsed.netloc):
            raise McpConnectorError("mcp_url 必须与 base_url 同源")
        login_path = str(data.get("login_path") or "/api/auth/token/").strip()
        if not login_path.startswith("/") or "://" in login_path:
            raise McpConnectorError("login_path 必须是站内绝对路径")
        guidance = str(data.get("query_guidance") or "").strip()
        if len(guidance) > 4000:
            raise McpConnectorError("query_guidance 不能超过 4000 字符")
        raw_recommended = data.get("recommended_operation_ids", [])
        if isinstance(raw_recommended, str):
            raw_recommended = raw_recommended.split(",")
        if not isinstance(raw_recommended, (list, tuple, set)):
            raise McpConnectorError("recommended_operation_ids 必须是字符串数组")
        recommended = tuple(
            dict.fromkeys(
                str(item).strip() for item in raw_recommended if str(item).strip()
            )
        )
        if len(recommended) > 30 or any(len(item) > 200 for item in recommended):
            raise McpConnectorError(
                "recommended_operation_ids 最多 30 项且每项不超过 200 字符"
            )
        expected_name = str(data.get("expected_server_name") or "").strip()
        if len(expected_name) > 200:
            raise McpConnectorError("expected_server_name 不能超过 200 字符")
        max_result = max(4096, min(int(data.get("max_result_bytes", 65536)), 1048576))
        return cls(
            base_url=base_url,
            mcp_url=mcp_url,
            login_path=login_path,
            timeout_seconds=max(1.0, min(float(data.get("timeout_seconds", 15)), 60.0)),
            max_result_bytes=max_result,
            max_total_result_bytes=max(
                max_result,
                min(int(data.get("max_total_result_bytes", 262144)), 4194304),
            ),
            verify_tls=_bool_value(data.get("verify_tls"), True),
            query_guidance=guidance,
            recommended_operation_ids=recommended,
            auth_type=str(data.get("auth_type") or "password_jwt").strip(),
            auth_config={
                key: data[key]
                for key in (
                    "login_path",
                    "username_field",
                    "password_field",
                    "token_field",
                    "auth_header_name",
                    "auth_header_template",
                )
                if key in data
            },
            expected_server_name=expected_name,
            max_response_bytes=max(
                65536,
                min(
                    int(data.get("max_response_bytes", MAX_STORED_RESULT_BYTES)),
                    MAX_STORED_RESULT_BYTES,
                ),
            ),
        )


class _LimitedAsyncStream(httpx2.AsyncByteStream):
    def __init__(self, stream: httpx2.AsyncByteStream, limit: int):
        self._stream = stream
        self._limit = limit

    async def __aiter__(self):
        total = 0
        async for chunk in self._stream:
            total += len(chunk)
            if total > self._limit:
                raise McpConnectorError("MCP 响应超过安全下载上限")
            yield chunk

    async def aclose(self) -> None:
        await self._stream.aclose()


class _LimitedTransport(httpx2.AsyncBaseTransport):
    def __init__(self, *, verify: bool, limit: int):
        self._transport = httpx2.AsyncHTTPTransport(verify=verify, retries=0)
        self._limit = limit

    async def handle_async_request(self, request: httpx2.Request) -> httpx2.Response:
        response = await self._transport.handle_async_request(request)
        raw_length = response.headers.get("content-length")
        try:
            content_length = int(raw_length) if raw_length else None
        except ValueError:
            content_length = None
        if content_length is not None and content_length > self._limit:
            await response.aclose()
            raise McpConnectorError("MCP 响应超过安全下载上限")
        response.stream = _LimitedAsyncStream(
            cast(httpx2.AsyncByteStream, response.stream),
            self._limit,
        )
        return response

    async def aclose(self) -> None:
        await self._transport.aclose()


def _auth_cache_ttl(headers: dict[str, str]) -> float:
    authorization = next(
        (value for name, value in headers.items() if name.lower() == "authorization"),
        "",
    )
    token = authorization.split(" ", 1)[-1].strip()
    parts = token.split(".")
    if len(parts) == 3:
        try:
            padding = "=" * (-len(parts[1]) % 4)
            payload = json.loads(
                base64.urlsafe_b64decode(parts[1] + padding).decode("utf-8")
            )
            return max(
                0.0,
                min(
                    _AUTH_CACHE_TTL_SECONDS,
                    float(payload.get("exp")) - time.time() - 30,
                ),
            )
        except (binascii.Error, TypeError, ValueError, UnicodeDecodeError):
            pass
    return _AUTH_CACHE_TTL_SECONDS


def _exception_has_status(exc: BaseException, status_code: int) -> bool:
    response = getattr(exc, "response", None)
    if getattr(response, "status_code", None) == status_code:
        return True
    return any(
        _exception_has_status(child, status_code)
        for child in getattr(exc, "exceptions", ())
        if isinstance(child, BaseException)
    )


def _find_connector_error(exc: BaseException) -> McpConnectorError | None:
    if isinstance(exc, McpConnectorError):
        return exc
    for child in getattr(exc, "exceptions", ()):
        if isinstance(child, BaseException):
            found = _find_connector_error(child)
            if found is not None:
                return found
    return None


def _redacted(value: Any) -> Any:
    if isinstance(value, dict):
        return {
            str(key): "[REDACTED]"
            if _SENSITIVE_KEY_RE.search(str(key))
            else _redacted(item)
            for key, item in value.items()
        }
    if isinstance(value, list):
        return [_redacted(item) for item in value]
    return value


class McpClient:
    def __init__(
        self,
        credentials: dict[str, str],
        cfg: McpConfig,
        *,
        cache_namespace: str = "",
    ):
        self.credentials = credentials
        self.cfg = cfg
        identity = json.dumps(
            {
                "namespace": cache_namespace,
                "credentials": credentials,
                "config": cfg.__dict__,
            },
            sort_keys=True,
            ensure_ascii=False,
            default=str,
        )
        self._runtime_identity = hashlib.sha256(identity.encode("utf-8")).hexdigest()

    def authenticate(self, *, force: bool = False) -> dict[str, str]:
        if force:
            RUNTIME_CACHE.invalidate_auth(self._runtime_identity)
        else:
            cached = RUNTIME_CACHE.get_auth(self._runtime_identity)
            if cached is not None:
                return cached

        def load() -> dict[str, str]:
            if not force:
                cached = RUNTIME_CACHE.get_auth(self._runtime_identity)
                if cached is not None:
                    return cached
            try:
                with httpx.Client(
                    timeout=self.cfg.timeout_seconds,
                    verify=self.cfg.verify_tls,
                    follow_redirects=False,
                ) as client:
                    headers = get_auth_strategy(self.cfg.auth_type).headers(
                        client=client,
                        base_url=self.cfg.base_url,
                        credentials=self.credentials,
                        config=self.cfg.auth_config,
                    )
            except ExternalAuthError as exc:
                raise McpConnectorError(str(exc)) from exc
            RUNTIME_CACHE.set_auth(
                self._runtime_identity,
                headers,
                ttl_seconds=_auth_cache_ttl(headers),
            )
            return dict(headers)

        return RUNTIME_CACHE.singleflight("auth", self._runtime_identity, load)

    async def _session_operation(
        self, operation: str, payload: Any, *, force_auth: bool
    ):
        headers = self.authenticate(force=force_auth)
        transport = _LimitedTransport(
            verify=self.cfg.verify_tls,
            limit=self.cfg.max_response_bytes,
        )
        async with (
            httpx2.AsyncClient(
                headers=headers,
                timeout=self.cfg.timeout_seconds,
                follow_redirects=False,
                transport=transport,
            ) as http_client,
            streamable_http_client(
                self.cfg.mcp_url,
                http_client=http_client,
            ) as streams,
            ClientSession(*streams) as session,
        ):
            initialized = await session.initialize()
            server_info = initialized.server_info
            if (
                self.cfg.expected_server_name
                and server_info.name != self.cfg.expected_server_name
            ):
                raise McpConnectorError(
                    "MCP Server 身份不匹配："
                    f"期望 {self.cfg.expected_server_name}，实际 {server_info.name}"
                )
            if operation == "list":
                tools: list[dict[str, Any]] = []
                cursor = None
                while True:
                    params = (
                        PaginatedRequestParams(cursor=cursor)
                        if cursor is not None
                        else None
                    )
                    page = await session.list_tools(params=params)
                    tools.extend(
                        tool.model_dump(by_alias=True, exclude_none=True)
                        for tool in page.tools
                    )
                    if len(tools) > _MAX_CATALOG_TOOLS:
                        raise McpConnectorError("MCP 工具目录超过 1000 项安全上限")
                    cursor = page.next_cursor
                    if not cursor:
                        break
                return {
                    "server": server_info.model_dump(by_alias=True, exclude_none=True),
                    "tools": tools,
                }
            if operation == "call":
                name, arguments = payload
                available: set[str] = set()
                cursor = None
                while True:
                    params = (
                        PaginatedRequestParams(cursor=cursor)
                        if cursor is not None
                        else None
                    )
                    page = await session.list_tools(params=params)
                    available.update(tool.name for tool in page.tools)
                    if len(available) > _MAX_CATALOG_TOOLS:
                        raise McpConnectorError("MCP 工具目录超过 1000 项安全上限")
                    cursor = page.next_cursor
                    if not cursor:
                        break
                if name not in available:
                    raise McpConnectorError("MCP 工具已不存在，请重新搜索工具目录")
                result = await session.call_tool(
                    name,
                    arguments=arguments,
                    read_timeout_seconds=self.cfg.timeout_seconds,
                )
                return result.model_dump(by_alias=True, exclude_none=True)
            raise AssertionError(f"unknown MCP operation: {operation}")

    def _run(self, operation: str, payload: Any = None) -> Any:
        for attempt in range(2):
            try:
                return anyio.run(
                    partial(
                        self._session_operation,
                        operation,
                        payload,
                        force_auth=attempt == 1,
                    )
                )
            except BaseException as exc:
                if attempt == 0 and _exception_has_status(exc, 401):
                    RUNTIME_CACHE.invalidate_auth(self._runtime_identity)
                    continue
                connector_error = _find_connector_error(exc)
                if connector_error is not None:
                    raise connector_error
                raise McpConnectorError(f"MCP 调用失败: {type(exc).__name__}") from exc
        raise McpConnectorError("MCP 认证失败")

    def _catalog(self, *, force: bool = False) -> dict[str, Any]:
        if force:
            RUNTIME_CACHE.invalidate_mcp_catalog(self._runtime_identity)
        else:
            cached = RUNTIME_CACHE.get_mcp_catalog(self._runtime_identity)
            if cached is not None:
                return cached

        def load() -> dict[str, Any]:
            if not force:
                cached = RUNTIME_CACHE.get_mcp_catalog(self._runtime_identity)
                if cached is not None:
                    return cached
            catalog = self._run("list")
            RUNTIME_CACHE.set_mcp_catalog(
                self._runtime_identity,
                catalog,
                ttl_seconds=_CATALOG_CACHE_TTL_SECONDS,
            )
            return catalog

        return RUNTIME_CACHE.singleflight("mcp-catalog", self._runtime_identity, load)

    @staticmethod
    def _operation_id(name: str) -> str:
        return f"mcp/{name}"

    @staticmethod
    def _remote_name(operation_id: str) -> str:
        value = str(operation_id or "").strip()
        return value.removeprefix("mcp/")

    def test_connection(self) -> dict[str, Any]:
        catalog = self._catalog(force=True)
        tool_count = len(catalog["tools"])
        return {
            "server": catalog["server"],
            # 对外测试连接响应与 OpenAPI 统一使用 operation_count；保留
            # tool_count，兼容已依赖 MCP 连接器原始响应的调用方。
            "operation_count": tool_count,
            "tool_count": tool_count,
        }

    def search(self, query: str, limit: int = 12) -> list[dict[str, Any]]:
        safe_limit = max(1, min(int(limit), 30))
        terms = [term.lower() for term in str(query or "").split() if term]
        recommended = {
            self._remote_name(item): index
            for index, item in enumerate(self.cfg.recommended_operation_ids)
        }
        scored: list[tuple[int, int, str, dict[str, Any]]] = []
        for tool in self._catalog()["tools"]:
            name = str(tool.get("name") or "")
            title = str(tool.get("title") or "")
            description = str(tool.get("description") or "")
            haystack = f"{name} {title} {description}".lower()
            score = sum(
                4 if term in name.lower() else 1 for term in terms if term in haystack
            )
            is_recommended = name in recommended
            if terms and score == 0 and not is_recommended:
                continue
            item = {
                "operation_id": self._operation_id(name),
                "name": name,
                "title": title,
                "summary": description,
                "input_schema": tool.get("inputSchema") or {"type": "object"},
                "output_schema": tool.get("outputSchema"),
                "annotations": tool.get("annotations"),
                "recommended": is_recommended,
            }
            scored.append(
                (
                    0 if is_recommended else 1,
                    recommended.get(name, -score),
                    name,
                    item,
                )
            )
        scored.sort(key=lambda row: (row[0], row[1], row[2]))
        return [row[3] for row in scored[:safe_limit]]

    def call(
        self,
        operation_id: str,
        arguments: dict[str, Any] | None = None,
        body: Any = None,
    ) -> dict[str, Any]:
        if body is not None:
            raise McpConnectorError("MCP 工具参数请全部放入 arguments，不使用 body")
        name = self._remote_name(operation_id)
        catalog = self._catalog()
        if name not in {str(tool.get("name") or "") for tool in catalog["tools"]}:
            raise McpConnectorError("MCP 工具不存在，请先搜索工具目录")
        raw = self._run("call", (name, dict(arguments or {})))
        if raw.get("isError"):
            details = json.dumps(
                _redacted(raw.get("content") or []),
                ensure_ascii=False,
                default=str,
            )
            raise McpConnectorError(f"MCP 工具返回错误: {details[:1000]}")
        data = raw.get("structuredContent")
        if data is None:
            data = {"content": raw.get("content") or []}
        normalized = json.loads(json.dumps(data, ensure_ascii=False, default=str))
        response_bytes = len(
            json.dumps(normalized, ensure_ascii=False, separators=(",", ":")).encode(
                "utf-8"
            )
        )
        return {
            "operation_id": self._operation_id(name),
            "status_code": None,
            "response_bytes": response_bytes,
            "truncated": False,
            "data": normalized,
        }
