"""Tests for MCP client support: client, connector, resolution, and API flow."""

from __future__ import annotations

import asyncio
import json
import uuid

import httpx
import pytest
from legate.agent.tooling import ToolContext
from legate.config import get_settings
from legate.connectors import resolution
from legate.connectors.base import ConnectorConfigError
from legate.connectors.mcp import McpConnector
from legate.mcp import client as mcp_client

MCP_URL = "https://mcp.example.com/rpc"


# --- A mock MCP server ------------------------------------------------------


def _mcp_handler(*, sse: bool = False, calls: list | None = None):
    """A MockTransport handler that speaks the MCP Streamable HTTP protocol."""

    def respond(payload: dict) -> httpx.Response:
        if sse:
            body = f"event: message\ndata: {json.dumps(payload)}\n\n"
            return httpx.Response(200, text=body, headers={"content-type": "text/event-stream"})
        return httpx.Response(200, json=payload)

    def handler(request: httpx.Request) -> httpx.Response:
        body = json.loads(request.content)
        method = body.get("method")
        if method == "initialize":
            return httpx.Response(
                200,
                json={"jsonrpc": "2.0", "id": body["id"], "result": {"protocolVersion": "x"}},
                headers={"mcp-session-id": "sess-1"},
            )
        if method == "notifications/initialized":
            return httpx.Response(202)
        if method == "tools/list":
            return respond(
                {
                    "jsonrpc": "2.0",
                    "id": body["id"],
                    "result": {
                        "tools": [
                            {
                                "name": "echo",
                                "description": "Echo text back.",
                                "inputSchema": {
                                    "type": "object",
                                    "properties": {"text": {"type": "string"}},
                                    "required": ["text"],
                                },
                                "annotations": {"readOnlyHint": True},
                            },
                            {
                                "name": "write_note",
                                "description": "Persist a note.",
                                "inputSchema": {"type": "object", "properties": {}},
                            },
                        ]
                    },
                }
            )
        if method == "tools/call":
            if calls is not None:
                calls.append(body["params"])
            name = body["params"]["name"]
            args = body["params"].get("arguments", {})
            return respond(
                {
                    "jsonrpc": "2.0",
                    "id": body["id"],
                    "result": {
                        "content": [{"type": "text", "text": f"{name}:{args.get('text', '')}"}],
                        "isError": False,
                    },
                }
            )
        return httpx.Response(
            200,
            json={
                "jsonrpc": "2.0",
                "id": body.get("id"),
                "error": {"code": -32601, "message": "no"},
            },
        )

    return handler


def _ctx(db=None, workspace_id=None) -> ToolContext:
    return ToolContext(workspace_id=workspace_id or uuid.uuid4(), settings=get_settings(), db=db)


# --- Client -----------------------------------------------------------------


async def test_client_lists_and_calls_tools(set_connector_http):
    set_connector_http(_mcp_handler())
    tools = await mcp_client.list_tools(MCP_URL, "tok")
    assert {t["name"] for t in tools} == {"echo", "write_note"}

    result = await mcp_client.call_tool(MCP_URL, "tok", "echo", {"text": "hi"})
    assert mcp_client.content_text(result) == "echo:hi"


async def test_client_parses_sse_responses(set_connector_http):
    set_connector_http(_mcp_handler(sse=True))
    tools = await mcp_client.list_tools(MCP_URL, None)
    assert [t["name"] for t in tools] == ["echo", "write_note"]


async def test_client_raises_on_error(set_connector_http):
    def handler(request: httpx.Request) -> httpx.Response:
        body = json.loads(request.content)
        if body.get("method") == "initialize":
            return httpx.Response(200, json={"jsonrpc": "2.0", "id": body["id"], "result": {}})
        return httpx.Response(
            200,
            json={"jsonrpc": "2.0", "id": body.get("id"), "error": {"message": "boom"}},
        )

    set_connector_http(handler)
    with pytest.raises(mcp_client.MCPError, match="boom"):
        await mcp_client.list_tools(MCP_URL, None)


# --- Connector --------------------------------------------------------------


async def test_connector_discover_normalizes_and_flags_sensitivity(set_connector_http):
    set_connector_http(_mcp_handler())
    tools = await McpConnector().discover({"url": MCP_URL}, {"token": "t"})
    by_name = {t["tool"]: t for t in tools}
    # readOnlyHint -> not sensitive; unknown -> sensitive by default.
    assert by_name["echo"]["sensitive"] is False
    assert by_name["write_note"]["sensitive"] is True
    assert by_name["echo"]["parameters"]["required"] == ["text"]


async def test_connector_test_connection_reports_tool_count(set_connector_http):
    set_connector_http(_mcp_handler())
    res = await McpConnector().test_connection({"url": MCP_URL}, {})
    assert res.ok and "2 tools" in res.message


async def test_connector_execute_maps_content(set_connector_http):
    set_connector_http(_mcp_handler())
    res = await McpConnector().execute("echo", {"text": "yo"}, {"url": MCP_URL}, {}, _ctx())
    assert res.ok and res.data["text"] == "echo:yo"


def test_connector_validates_url():
    conn = McpConnector()
    conn.validate_config({"url": "https://ok.example.com/rpc"})
    with pytest.raises(ConnectorConfigError):
        conn.validate_config({"url": "not-a-url"})
    with pytest.raises(ConnectorConfigError):
        conn.validate_config({})  # url required


def test_connector_has_no_static_tools():
    assert McpConnector().list_tools() == []


# --- Resolution -------------------------------------------------------------


class _Row:
    def __init__(self, rid, config):
        self.id = rid
        self.type = "mcp"
        self.config_json = config


def test_mcp_specs_namespacing_and_sensitivity():
    rid = uuid.uuid4()
    row = _Row(
        rid,
        {
            "url": MCP_URL,
            "mcp_tools": [
                {"tool": "echo", "description": "e", "parameters": {}, "sensitive": False},
                {"tool": "write_note", "description": "w", "parameters": {}, "sensitive": True},
            ],
        },
    )
    specs = resolution.mcp_specs_for_row(row)
    names = {spec.name: (spec, original) for spec, original in specs}
    key = f"mcp.{str(rid)[:8]}.echo"
    assert key in names
    assert names[key][1] == "echo"  # original tool name preserved for the server call
    assert names[key][0].is_sensitive is False
    assert names[f"mcp.{str(rid)[:8]}.write_note"][0].is_sensitive is True


# --- Full API flow ----------------------------------------------------------


def test_install_discovers_caches_and_exposes_mcp_tools(client, helpers, set_connector_http):
    set_connector_http(_mcp_handler())
    owner = helpers.bootstrap_owner(client)
    token = owner["access_token"]

    created = helpers.install_connector(
        client, token, type="mcp", name="Demo MCP", config={"url": MCP_URL}, secrets={"token": "t"}
    )
    # Tools were discovered and cached on the connector row.
    cached = created["config_json"]["mcp_tools"]
    assert {t["tool"] for t in cached} == {"echo", "write_note"}

    # They appear in the workspace tool list, namespaced under "mcp.".
    tools = client.get("/tools", headers=helpers.headers(token)).json()
    mcp_names = [t["name"] for t in tools if t["name"].startswith("mcp.")]
    assert any(n.endswith(".echo") for n in mcp_names)
    assert any(n.endswith(".write_note") for n in mcp_names)

    # An agent can be created with an MCP tool in its allow-list (validation passes).
    echo_tool = next(n for n in mcp_names if n.endswith(".echo"))
    agent = helpers.create_agent(client, token, name="MCP Agent", tools=[echo_tool])
    assert echo_tool in agent["tools"]


def test_registry_registers_mcp_tools_that_call_the_server(client, helpers, db, set_connector_http):
    calls: list = []
    set_connector_http(_mcp_handler(calls=calls))
    owner = helpers.bootstrap_owner(client)
    created = helpers.install_connector(
        client,
        owner["access_token"],
        type="mcp",
        name="Demo MCP",
        config={"url": MCP_URL},
        secrets={"token": "t"},
    )
    ws_id = uuid.UUID(owner["workspace"]["id"])

    registry = resolution.build_workspace_tool_registry(db, ws_id)
    tool_name = f"mcp.{created['id'][:8]}.echo"
    registered = registry.get(tool_name)
    assert registered is not None

    result = asyncio.run(registered.handler({"text": "hey"}, _ctx(db=db, workspace_id=ws_id)))
    assert result.ok and result.data["text"] == "echo:hey"
    assert calls and calls[-1]["name"] == "echo"  # server was called with the original name
