from __future__ import annotations

import asyncio
import json
from importlib.util import module_from_spec, spec_from_file_location
from pathlib import Path
from typing import Any
import logging


SERVER_PATH = Path(__file__).with_name("server.py")

spec = spec_from_file_location("osvx_mcp_server_under_test", SERVER_PATH)
if spec is None or spec.loader is None:
    raise RuntimeError(f"Cannot load server module from {SERVER_PATH}")

server = module_from_spec(spec)

# Temporarily suppress INFO logs from the server module during import
target_logger = logging.getLogger("osvx_mcp_server_under_test")
original_level = target_logger.level
target_logger.setLevel(logging.WARNING)
try:
    spec.loader.exec_module(server)
finally:
    target_logger.setLevel(original_level)


class FakeRequest:
    def __init__(self, payload: dict[str, Any]) -> None:
        self._payload = payload

    async def json(self) -> dict[str, Any]:
        return self._payload


async def rpc(method: str, request_id: str) -> dict[str, Any]:
    response = await server.mcp_handler(
        FakeRequest(
            {
                "jsonrpc": "2.0",
                "method": method,
                "id": request_id,
            }
        )
    )

    assert response.status_code == 200
    payload = json.loads(response.body.decode("utf-8"))
    assert payload["jsonrpc"] == "2.0"
    assert payload["id"] == request_id
    assert "result" in payload
    return payload["result"]


async def verify() -> None:
    initialize = await rpc("initialize", "initialize-test")

    assert initialize["protocolVersion"] == "2024-11-05"
    assert initialize["capabilities"]["tools"] == {}
    assert initialize["serverInfo"]["name"] == "opax-mcp"
    assert isinstance(initialize["serverInfo"]["version"], str)
    assert initialize["serverInfo"]["version"]

    discovery = await rpc("tools/list", "tools-list-test")
    tools = discovery["tools"]
    discovered_names = {tool["name"] for tool in tools}

    assert len(server.TOOLS) == 32
    assert len(tools) == 32
    assert len(discovered_names) == 32
    assert discovered_names == set(server.TOOLS)

    for tool in tools:
        assert isinstance(tool["name"], str) and tool["name"]
        assert isinstance(tool["description"], str) and tool["description"]
        assert isinstance(tool["inputSchema"], dict)
        assert tool["inputSchema"].get("type") == "object"

    print("OSVX_MCP_DISCOVERY_VERIFIED")
    print("TOOLS_COUNT=32")
    print("INITIALIZE_TOOLS_CAPABILITY=YES")
    print("TOOLS_LIST_MATCHES_REGISTRY=YES")


if __name__ == "__main__":
    asyncio.run(verify())