from __future__ import annotations

from pathlib import Path

import pytest
from mcp import McpError
from mcp.types import ErrorData

from wf_mcp.broker import WfMcpService
from wf_mcp.capabilities import DiscoveredTool
from wf_mcp.models import ConnectionConfig
from wf_mcp.sdk import McpSdkAdapter
from wf_mcp.storage import FileStore
from wf_sources_mcp.connections import mcp_source_connection_from_connection_config

from .test_support import (
    everything_server_connection,
    fixture_server_path,
    sys,
)


class _ToolsOnlyAdapter:
    async def list_tools(self, connection, auth):
        return [
            DiscoveredTool(
                name="echo_tool",
                title="Echo",
                description="Echo text.",
                input_schema={"type": "object", "properties": {}},
                output_schema={"type": "object", "properties": {}},
            )
        ]

    async def list_resources(self, connection, auth):
        raise McpError(ErrorData(code=-32601, message="Method not found"))

    async def list_prompts(self, connection, auth):
        raise McpError(ErrorData(code=-32601, message="Method not found"))

    async def get_connection_metadata(self, connection, auth):
        return {
            "server": getattr(
                connection, "provider", getattr(connection, "server", None)
            )
        }

    async def read_resource(self, connection, auth, uri):
        raise NotImplementedError

    async def get_prompt(self, connection, auth, prompt_name, arguments=None):
        raise NotImplementedError

    async def invoke_method(self, connection, auth, method, params=None):
        raise NotImplementedError

    async def send_notification(self, connection, auth, method, params=None):
        raise NotImplementedError

    async def call_tool(self, connection, auth, tool_name, payload):
        raise NotImplementedError


class _WrappedToolsOnlyAdapter(_ToolsOnlyAdapter):
    async def list_resources(self, connection, auth):
        raise ExceptionGroup(
            "unhandled errors in a TaskGroup",
            [McpError(ErrorData(code=-32601, message="Method not found"))],
        )


@pytest.mark.asyncio
async def test_mcp_sdk_adapter_lists_and_calls_stdio_tool(tmp_path: Path) -> None:
    service = WfMcpService(store=FileStore(tmp_path / "sdk_adapter_store"))
    service.register_connection(
        ConnectionConfig(
            id="fixture.personal",
            server="fixture",
            account="personal",
            metadata={
                "transport": "stdio",
                "command": sys.executable,
                "args": [fixture_server_path()],
            },
        )
    )
    service.register_adapter("fixture", McpSdkAdapter())

    try:
        await service.refresh_connection_catalog("fixture.personal")
    except PermissionError as exc:
        pytest.skip(f"stdio MCP transport is not permitted in this environment: {exc}")

    payload = service.get_catalog().as_payload()
    assert payload["nodes"][0]["qualified_name"] == "fixture.personal.echo_tool"
    assert payload["resources"] == [
        {
            "qualified_name": "fixture.personal.resource.welcome",
            "connection_id": "fixture.personal",
            "local_name": "resource.welcome",
            "uri": "fixture://docs/welcome",
            "title": "Resource Welcome",
            "description": "Welcome text resource for fixture tests.",
            "mime_type": "text/plain",
            "metadata": payload["resources"][0]["metadata"],
        }
    ]
    assert payload["prompts"] == [
        {
            "qualified_name": "fixture.personal.prompt.summarize",
            "connection_id": "fixture.personal",
            "local_name": "prompt.summarize",
            "title": "Prompt Summarize",
            "description": "Summarize an input text for fixture tests.",
            "arguments": payload["prompts"][0]["arguments"],
            "metadata": payload["prompts"][0]["metadata"],
        }
    ]

    resource_result = await service.read_resource("fixture.personal.resource.welcome")
    prompt_result = await service.render_prompt(
        "fixture.personal.prompt.summarize",
        arguments={"text": "hello"},
    )
    assert (
        resource_result["contents"][0]["text"] == "Welcome from the fixture MCP server."
    )
    assert (
        prompt_result["messages"][0]["content"]["text"]
        == "Summarize this text:\n\nhello"
    )
    ping_result = await service.invoke_method("fixture.personal", "ping")
    assert ping_result == {}

    adapter = McpSdkAdapter()
    try:
        source_connection = mcp_source_connection_from_connection_config(
            service.connections.get("fixture.personal")
        )
        result = await adapter.call_tool(
            connection=source_connection,
            auth=None,
            tool_name="echo_tool",
            payload={"text": "hello"},
        )
    except PermissionError as exc:
        pytest.skip(f"stdio MCP transport is not permitted in this environment: {exc}")
    assert result.outcome == "ok"
    assert result.output == {"echoed": "hello"}


async def test_mcp_sdk_adapter_can_probe_everything_server(tmp_path: Path) -> None:
    connection = everything_server_connection()
    if connection is None:
        pytest.skip(
            "set MCP_EVERYTHING_COMMAND to enable the live everything-server integration test"
        )

    service = WfMcpService(store=FileStore(tmp_path / "everything_server_store"))
    service.register_connection(connection)
    service.register_adapter("everything", McpSdkAdapter())

    try:
        await service.refresh_connection_catalog("everything.default")
    except PermissionError as exc:
        pytest.skip(f"live MCP transport is not permitted in this environment: {exc}")

    payload = service.get_catalog().as_payload()
    assert payload["nodes"], "everything-server should expose at least one tool"
    assert all(
        node["qualified_name"].startswith("everything.default.")
        for node in payload["nodes"]
    )
    assert "resources" in payload
    assert "prompts" in payload


async def test_refresh_catalog_keeps_tools_when_optional_lists_are_unsupported(
    tmp_path: Path,
) -> None:
    service = WfMcpService(store=FileStore(tmp_path / "tools_only_server_store"))
    service.register_connection(
        ConnectionConfig(
            id="tools_only.personal",
            server="tools_only",
            account="personal",
            metadata={"transport": "stdio", "command": "fake-tools-only"},
        )
    )
    service.register_adapter("tools_only", _ToolsOnlyAdapter())

    await service.refresh_connection_catalog("tools_only.personal")

    payload = service.get_catalog().as_payload()
    assert payload["nodes"][0]["qualified_name"] == "tools_only.personal.echo_tool"
    assert payload["resources"] == []
    assert payload["prompts"] == []


async def test_refresh_catalog_unwraps_taskgroup_method_not_found(
    tmp_path: Path,
) -> None:
    service = WfMcpService(
        store=FileStore(tmp_path / "wrapped_tools_only_server_store")
    )
    service.register_connection(
        ConnectionConfig(
            id="wrapped_tools_only.personal",
            server="wrapped_tools_only",
            account="personal",
            metadata={"transport": "stdio", "command": "fake-tools-only"},
        )
    )
    service.register_adapter("wrapped_tools_only", _WrappedToolsOnlyAdapter())

    await service.refresh_connection_catalog("wrapped_tools_only.personal")

    payload = service.get_catalog().as_payload()
    assert payload["nodes"][0]["qualified_name"] == (
        "wrapped_tools_only.personal.echo_tool"
    )