"""Core MCP JSON-RPC handler — transport-agnostic.

The handler maps the MCP wire protocol onto an :class:`A2AAgent`:

- ``initialize`` advertises tool support and echoes the server identity
  from the agent's :class:`AgentCard`.
- ``tools/list`` enumerates the agent's skills.
- ``tools/call`` dispatches into :meth:`A2AAgent.invoke_json`.

A transport (stdio, HTTP) calls :meth:`MCPServer.handle` with each
decoded JSON-RPC message and forwards the returned dict back to the
client. Notifications return ``None``.
"""
from __future__ import annotations

import json
import logging
from typing import Any, Awaitable, Callable

from ..agent import A2AAgent, SkillInputError, SkillNotFound, SkillSpec
from ..auth import NoAuth
from ..context import LocalRunContext, MissingScopes, RunContext

MCP_PROTOCOL_VERSION = "2025-06-18"

# JSON-RPC error codes — standard plus MCP additions.
_PARSE_ERROR = -32700
_INVALID_REQUEST = -32600
_METHOD_NOT_FOUND = -32601
_INVALID_PARAMS = -32602
_INTERNAL_ERROR = -32603

log = logging.getLogger(__name__)


ContextBuilder = Callable[[str], Awaitable[RunContext[Any]] | RunContext[Any]]
"""Factory called per tools/call with the skill name; returns the RunContext."""


def _apply_caller_meta(ctx: RunContext[Any], meta: dict[str, Any]) -> None:
    """Mutate ``ctx`` to carry caller-supplied credentials forwarded by an
    MCP gateway (e.g. ``a2amcp``).

    Recognized keys: ``cp_jwt`` + ``cp_url`` (forwarded to skills via the
    ``ctx.cp_jwt`` / ``ctx.cp_url`` properties); ``bucket`` (mints a
    workspace whose ``bucket`` attribute resolves, AND overrides
    ``ctx.write_artifact`` to persist to MinIO so outputs survive past
    the request — the default in-memory artifacts dict drops them on
    pod restart, which the early version of this shim accidentally
    inherited).
    """
    import os

    cp_jwt = meta.get("cp_jwt")
    cp_url = meta.get("cp_url")
    bucket = meta.get("bucket")
    if isinstance(cp_jwt, str) and cp_jwt:
        ctx._cp_jwt = cp_jwt  # noqa: SLF001
    if isinstance(cp_url, str) and cp_url:
        ctx._cp_url = cp_url  # noqa: SLF001
    if not isinstance(bucket, str) or not bucket:
        return

    if getattr(ctx, "_workspace", None) is None:
        from ..workspace import (
            LocalWorkspaceClient,
            MinIOWorkspaceClient,
            WorkspaceAccess,
            WorkspaceMode,
        )

        access = WorkspaceAccess.dynamic(
            max_files=256,
            allowed_modes=(
                WorkspaceMode.READ_ONLY,
                WorkspaceMode.READ_WRITE_OVERLAY,
            ),
            require_reason=False,
            require_human_approval=False,
        )
        endpoint = os.environ.get("A2A_MINIO_ENDPOINT")
        access_key = os.environ.get("A2A_MINIO_ACCESS_KEY", "minioadmin")
        secret_key = os.environ.get("A2A_MINIO_SECRET_KEY", "minioadmin")
        if endpoint:
            ctx._workspace = MinIOWorkspaceClient(  # noqa: SLF001
                bucket=bucket,
                endpoint_url=endpoint,
                access_key_id=access_key,
                secret_access_key=secret_key,
                access=access,
                issuer="mcp-gateway",
            )
        else:
            ctx._workspace = LocalWorkspaceClient(  # noqa: SLF001
                files={}, access=access, bucket=bucket, issuer="mcp-gateway",
            )

    # Persist ctx.write_artifact outputs to MinIO under
    # ``outputs/{task_id}/{name}`` so users see them in their bucket
    # instead of memory://. Pod needs ``A2A_MINIO_*`` env — when not
    # set, fall back to the in-memory default (no regression).
    endpoint = os.environ.get("A2A_MINIO_ENDPOINT")
    access_key = os.environ.get("A2A_MINIO_ACCESS_KEY", "minioadmin")
    secret_key = os.environ.get("A2A_MINIO_SECRET_KEY", "minioadmin")
    if endpoint:
        _bind_s3_write_artifact(ctx, bucket, endpoint, access_key, secret_key)


def _bind_s3_write_artifact(
    ctx: Any,
    bucket: str,
    endpoint: str,
    access_key: str,
    secret_key: str,
) -> None:
    """Replace ``ctx.write_artifact`` with one that PUTs to MinIO."""
    try:
        import boto3
        from botocore.config import Config as _BotoConfig
    except Exception:  # noqa: BLE001
        return

    from ..context import ArtifactRef

    s3 = boto3.client(
        "s3",
        endpoint_url=endpoint,
        aws_access_key_id=access_key,
        aws_secret_access_key=secret_key,
        region_name="us-east-1",
        config=_BotoConfig(s3={"addressing_style": "path"}),
    )

    task_id = getattr(ctx, "task_id", "task")

    async def write_artifact(
        name: str, data: bytes, mime_type: str
    ) -> ArtifactRef:
        key = f"outputs/{task_id}/{name}"
        s3.put_object(
            Bucket=bucket, Key=key, Body=data, ContentType=mime_type,
        )
        # Also keep a memory ref so tests / introspection that read
        # ctx.artifacts still work.
        try:
            ctx.artifacts[name] = data  # type: ignore[attr-defined]
        except Exception:  # noqa: BLE001
            pass
        return ArtifactRef(
            name=name,
            uri=f"s3://{bucket}/{key}",
            mime_type=mime_type,
            size_bytes=len(data),
        )

    ctx.write_artifact = write_artifact  # type: ignore[assignment]


def _default_context_builder(agent: A2AAgent) -> ContextBuilder:
    """Build a permissive :class:`LocalRunContext` keyed by skill name."""

    auth_cls = type(agent).auth_model

    def _build(skill_name: str) -> LocalRunContext[Any]:
        try:
            auth = auth_cls()
        except Exception:  # noqa: BLE001 — best-effort default
            auth = NoAuth()
        return LocalRunContext(auth=auth, task_id=f"mcp-{skill_name}")

    return _build


def skills_to_tools(agent: A2AAgent) -> list[dict[str, Any]]:
    """Render the agent's skills as MCP tool descriptors."""
    tools: list[dict[str, Any]] = []
    for spec in agent.skills.values():
        tool: dict[str, Any] = {
            "name": spec.name,
            "description": spec.description or f"{agent.name}.{spec.name}",
            "inputSchema": spec.input_schema,
        }
        # outputSchema is optional in MCP 2025-06-18 but lets clients
        # validate structuredContent. Always emit when we have one.
        if spec.output_schema:
            tool["outputSchema"] = _ensure_object_schema(spec.output_schema)
        tools.append(tool)
    return tools


def _ensure_object_schema(schema: dict[str, Any]) -> dict[str, Any]:
    """MCP requires outputSchema to be ``type: object``. Wrap scalars."""
    if schema.get("type") == "object":
        return schema
    return {
        "type": "object",
        "properties": {"result": schema},
        "required": ["result"],
    }


def tool_call_result(spec: SkillSpec, value: Any) -> dict[str, Any]:
    """Wrap a skill return value as an MCP CallToolResult."""
    output_schema = spec.output_schema if spec else None
    if output_schema and output_schema.get("type") == "object":
        structured = value
    else:
        structured = {"result": value}
    text = json.dumps(value, ensure_ascii=False, default=str)
    return {
        "content": [{"type": "text", "text": text}],
        "structuredContent": structured,
        "isError": False,
    }


def tool_call_error(message: str) -> dict[str, Any]:
    """Wrap a tool-side failure as an MCP CallToolResult with isError=true."""
    return {
        "content": [{"type": "text", "text": message}],
        "isError": True,
    }


class MCPServer:
    """Transport-agnostic MCP request handler for an :class:`A2AAgent`."""

    def __init__(
        self,
        agent: A2AAgent,
        *,
        context_builder: ContextBuilder | None = None,
    ) -> None:
        self.agent = agent
        self._build_ctx = context_builder or _default_context_builder(agent)
        self._initialized = False

    async def handle(self, message: dict[str, Any]) -> dict[str, Any] | None:
        """Dispatch one JSON-RPC message. Returns response or ``None`` for notifications."""
        if not isinstance(message, dict):
            return _error(None, _INVALID_REQUEST, "request must be a JSON object")

        method = message.get("method")
        msg_id = message.get("id")
        is_notification = "id" not in message

        if not isinstance(method, str):
            return _error(msg_id, _INVALID_REQUEST, "missing method")

        params = message.get("params") or {}
        if not isinstance(params, dict):
            return _error(msg_id, _INVALID_PARAMS, "params must be an object")

        try:
            result = await self._dispatch(method, params)
        except _MCPError as exc:
            if is_notification:
                return None
            return _error(msg_id, exc.code, exc.message, exc.data)
        except Exception as exc:  # noqa: BLE001
            log.exception("mcp dispatch failed")
            if is_notification:
                return None
            return _error(msg_id, _INTERNAL_ERROR, f"{type(exc).__name__}: {exc}")

        if is_notification:
            return None
        return {"jsonrpc": "2.0", "id": msg_id, "result": result}

    async def _dispatch(self, method: str, params: dict[str, Any]) -> Any:
        if method == "initialize":
            return self._on_initialize(params)
        if method == "notifications/initialized":
            self._initialized = True
            return None
        if method == "ping":
            return {}
        if method == "tools/list":
            return {"tools": skills_to_tools(self.agent)}
        if method == "tools/call":
            return await self._on_tool_call(params)
        raise _MCPError(_METHOD_NOT_FOUND, f"unknown method: {method}")

    def _on_initialize(self, params: dict[str, Any]) -> dict[str, Any]:
        agent_cls = type(self.agent)
        client_version = params.get("protocolVersion") or MCP_PROTOCOL_VERSION
        return {
            "protocolVersion": client_version,
            "capabilities": {
                "tools": {"listChanged": False},
                # Server may issue ``elicitation/create`` requests mid
                # ``tools/call`` (driven by skills using ``ctx.collect`` /
                # ``ctx.ask``). The HTTP transport delivers them inline
                # on the SSE response and routes client responses back to
                # the pending future via session id. See mcp/http.py.
                "elicitation": {},
            },
            "serverInfo": {
                "name": agent_cls.name,
                "version": agent_cls.version,
            },
            "instructions": agent_cls.description or "",
        }

    async def _on_tool_call(self, params: dict[str, Any]) -> dict[str, Any]:
        name = params.get("name")
        if not isinstance(name, str):
            raise _MCPError(_INVALID_PARAMS, "tools/call: missing name")
        arguments = params.get("arguments") or {}
        if not isinstance(arguments, dict):
            raise _MCPError(_INVALID_PARAMS, "tools/call: arguments must be object")
        # Optional caller context — set by MCP gateways that have a CP login
        # available (a2amcp). When present, hand the skill a workspace
        # scoped to the user's bucket and the CP JWT so agents that need
        # to act on behalf of the user (agent-builder, file tools, …) can
        # do so from an MCP client the same way they do from the platform
        # orchestrator. Unknown clients omit _meta and get the original
        # unauthenticated context.
        meta = params.get("_meta") or {}
        if not isinstance(meta, dict):
            meta = {}

        spec = self.agent.skills.get(name)
        if spec is None:
            # Per MCP, unknown tools surface as isError content, not JSON-RPC error.
            return tool_call_error(f"unknown tool: {name}")

        ctx_or_awaitable = self._build_ctx(name)
        if hasattr(ctx_or_awaitable, "__await__"):
            ctx = await ctx_or_awaitable  # type: ignore[assignment]
        else:
            ctx = ctx_or_awaitable

        _apply_caller_meta(ctx, meta)

        try:
            value = await self.agent.invoke_json(name, ctx, arguments)
        except SkillNotFound:
            return tool_call_error(f"unknown tool: {name}")
        except SkillInputError as exc:
            return tool_call_error(f"invalid arguments: {exc}")
        except MissingScopes as exc:
            return tool_call_error(f"permission denied: {exc}")
        except Exception as exc:  # noqa: BLE001
            log.exception("tool call failed: %s", name)
            return tool_call_error(f"{type(exc).__name__}: {exc}")

        return tool_call_result(spec, value)


class _MCPError(Exception):
    """Internal sentinel for JSON-RPC error returns."""

    def __init__(self, code: int, message: str, data: Any | None = None) -> None:
        super().__init__(message)
        self.code = code
        self.message = message
        self.data = data


def _error(
    msg_id: Any, code: int, message: str, data: Any | None = None
) -> dict[str, Any]:
    err: dict[str, Any] = {"code": code, "message": message}
    if data is not None:
        err["data"] = data
    return {"jsonrpc": "2.0", "id": msg_id, "error": err}


def parse_message(raw: str | bytes) -> dict[str, Any]:
    """Parse a JSON-RPC frame; raises :class:`_MCPError` with PARSE_ERROR."""
    try:
        msg = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise _MCPError(_PARSE_ERROR, f"parse error: {exc}") from exc
    if not isinstance(msg, dict):
        raise _MCPError(_INVALID_REQUEST, "request must be a JSON object")
    return msg
