"""Tests for MCP gateway — Tier 1 safety tools mounted on proxy."""

from __future__ import annotations

import json

_MCP_HEADERS = {"Accept": "application/json, text/event-stream"}


def _parse_sse_json(resp):
    """Extract the first JSON-RPC message from an SSE response."""
    for line in resp.text.splitlines():
        if line.startswith("data: "):
            return json.loads(line[6:])
    return None


def _init_mcp_session(client):
    """Initialize an MCP session and return (session_headers, init_data)."""
    resp = client.post(
        "/mcp/mcp/",
        json={
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {"name": "test", "version": "0.1"},
            },
        },
        headers=_MCP_HEADERS,
    )
    assert resp.status_code == 200, f"MCP init returned {resp.status_code}: {resp.text[:200]}"
    data = _parse_sse_json(resp)

    session_id = resp.headers.get("mcp-session-id", "")
    headers = {**_MCP_HEADERS, "mcp-session-id": session_id} if session_id else dict(_MCP_HEADERS)

    # Send initialized notification
    client.post(
        "/mcp/mcp/",
        json={
            "jsonrpc": "2.0",
            "method": "notifications/initialized",
        },
        headers=headers,
    )

    return headers, data


def test_create_mcp_app_returns_asgi():
    """create_mcp_app returns an ASGI app when fastmcp is installed."""
    from aceteam_aep.mcp_gateway import create_mcp_app
    from aceteam_aep.proxy.app import ProxyState

    state = ProxyState()
    app = create_mcp_app(state)
    assert app is not None, "create_mcp_app returned None — is fastmcp installed?"
    assert callable(app)


def test_proxy_mounts_mcp():
    """The proxy app should have /mcp/ route when fastmcp is available."""
    from starlette.testclient import TestClient

    from aceteam_aep.proxy.app import create_proxy_app

    app = create_proxy_app()
    with TestClient(app) as client:
        # FastMCP streamable HTTP expects POST at /mcp/mcp/ for messages
        # Try POST to the MCP message endpoint
        resp = client.post(
            "/mcp/mcp/",
            json={
                "jsonrpc": "2.0",
                "id": 1,
                "method": "initialize",
                "params": {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {},
                    "clientInfo": {"name": "test", "version": "0.1"},
                },
            },
            headers=_MCP_HEADERS,
        )
        # Should get 200 with initialize response, not 404
        assert resp.status_code == 200, (
            f"MCP endpoint returned {resp.status_code}: {resp.text[:200]}"
        )
        data = _parse_sse_json(resp)
        assert data.get("result", {}).get("serverInfo", {}).get("name") == "aceteam-gateway"


def test_mcp_tools_list():
    """MCP tools/list should return all 4 Tier 1 tools."""
    from starlette.testclient import TestClient

    from aceteam_aep.proxy.app import create_proxy_app

    app = create_proxy_app()
    with TestClient(app) as client:
        headers, _ = _init_mcp_session(client)

        # List tools
        resp = client.post(
            "/mcp/mcp/",
            json={
                "jsonrpc": "2.0",
                "id": 2,
                "method": "tools/list",
                "params": {},
            },
            headers=headers,
        )
        assert resp.status_code == 200
        data = _parse_sse_json(resp)
        tool_names = [t["name"] for t in data.get("result", {}).get("tools", [])]
        assert "check_safety" in tool_names
        assert "get_safety_status" in tool_names
        assert "set_safety_policy" in tool_names
        assert "get_cost_summary" in tool_names


def test_mcp_check_safety_tool():
    """check_safety tool should return PASS for safe text."""
    from starlette.testclient import TestClient

    from aceteam_aep.proxy.app import create_proxy_app

    app = create_proxy_app()
    with TestClient(app) as client:
        headers, _ = _init_mcp_session(client)

        # Call check_safety with safe text
        resp = client.post(
            "/mcp/mcp/",
            json={
                "jsonrpc": "2.0",
                "id": 3,
                "method": "tools/call",
                "params": {
                    "name": "check_safety",
                    "arguments": {"text": "What is the capital of France?"},
                },
            },
            headers=headers,
        )
        assert resp.status_code == 200
        data = _parse_sse_json(resp)
        content = data.get("result", {}).get("content", [{}])[0].get("text", "")
        assert "PASS" in content


def test_mcp_set_safety_policy():
    """set_safety_policy tool should toggle safety off and back on."""
    from starlette.testclient import TestClient

    from aceteam_aep.proxy.app import create_proxy_app

    app = create_proxy_app()
    with TestClient(app) as client:
        headers, _ = _init_mcp_session(client)

        # Disable safety
        resp = client.post(
            "/mcp/mcp/",
            json={
                "jsonrpc": "2.0",
                "id": 4,
                "method": "tools/call",
                "params": {
                    "name": "set_safety_policy",
                    "arguments": {"enabled": False},
                },
            },
            headers=headers,
        )
        assert resp.status_code == 200
        data = _parse_sse_json(resp)
        content = data.get("result", {}).get("content", [{}])[0].get("text", "")
        assert '"safety_enabled": false' in content

        # Re-enable safety
        resp = client.post(
            "/mcp/mcp/",
            json={
                "jsonrpc": "2.0",
                "id": 5,
                "method": "tools/call",
                "params": {
                    "name": "set_safety_policy",
                    "arguments": {"enabled": True},
                },
            },
            headers=headers,
        )
        assert resp.status_code == 200
        data = _parse_sse_json(resp)
        content = data.get("result", {}).get("content", [{}])[0].get("text", "")
        assert '"safety_enabled": true' in content


def test_mcp_get_cost_summary():
    """get_cost_summary tool should return cost data."""
    from starlette.testclient import TestClient

    from aceteam_aep.proxy.app import create_proxy_app

    app = create_proxy_app()
    with TestClient(app) as client:
        headers, _ = _init_mcp_session(client)

        resp = client.post(
            "/mcp/mcp/",
            json={
                "jsonrpc": "2.0",
                "id": 6,
                "method": "tools/call",
                "params": {
                    "name": "get_cost_summary",
                    "arguments": {},
                },
            },
            headers=headers,
        )
        assert resp.status_code == 200
        data = _parse_sse_json(resp)
        content = data.get("result", {}).get("content", [{}])[0].get("text", "")
        assert "total_cost_usd" in content
        assert "total_calls" in content


def test_fastmcp_importable():
    """fastmcp should be installed in dev dependencies."""
    try:
        import fastmcp  # noqa: F401

        has_fastmcp = True
    except ImportError:
        has_fastmcp = False
    assert has_fastmcp


def test_mcp_check_safety_increments_call_count():
    """Each check_safety call should increment the session call counter."""
    import json as _json

    from starlette.testclient import TestClient

    from aceteam_aep.proxy.app import create_proxy_app

    app = create_proxy_app()
    with TestClient(app) as client:
        headers, _ = _init_mcp_session(client)

        # Make two check_safety calls
        for i in range(2):
            resp = client.post(
                "/mcp/mcp/",
                json={
                    "jsonrpc": "2.0",
                    "id": 10 + i,
                    "method": "tools/call",
                    "params": {
                        "name": "check_safety",
                        "arguments": {"text": f"Safe message number {i}"},
                    },
                },
                headers=headers,
            )
            assert resp.status_code == 200

        # Verify call count via get_safety_status
        resp = client.post(
            "/mcp/mcp/",
            json={
                "jsonrpc": "2.0",
                "id": 20,
                "method": "tools/call",
                "params": {
                    "name": "get_safety_status",
                    "arguments": {},
                },
            },
            headers=headers,
        )
        assert resp.status_code == 200
        data = _parse_sse_json(resp)
        content = data.get("result", {}).get("content", [{}])[0].get("text", "")
        status = _json.loads(content)
        assert status["calls"] >= 2, f"Expected at least 2 calls, got {status['calls']}"


def test_mcp_set_policy_with_detectors():
    """set_safety_policy should accept detector configuration and report updated status."""
    import json as _json

    from starlette.testclient import TestClient

    from aceteam_aep.proxy.app import create_proxy_app

    app = create_proxy_app()
    with TestClient(app) as client:
        headers, _ = _init_mcp_session(client)

        resp = client.post(
            "/mcp/mcp/",
            json={
                "jsonrpc": "2.0",
                "id": 30,
                "method": "tools/call",
                "params": {
                    "name": "set_safety_policy",
                    "arguments": {
                        "detectors": {
                            "pii": {"enabled": False},
                            "agent_threat": {"enabled": True, "action": "block"},
                        }
                    },
                },
            },
            headers=headers,
        )
        assert resp.status_code == 200
        data = _parse_sse_json(resp)
        content = data.get("result", {}).get("content", [{}])[0].get("text", "")
        result = _json.loads(content)
        assert result["status"] == "updated"
        detectors = result.get("policy", {}).get("detectors", {})
        assert detectors.get("pii", {}).get("enabled") is False
        assert detectors.get("agent_threat", {}).get("action") == "block"
