import pytest
from unittest.mock import Mock

pytest.importorskip("pydantic_ai")

from pydantic_ai._run_context import RunContext
from pydantic_ai.mcp import MCPServerStreamableHTTP
from pydantic_ai.usage import RunUsage

from rgltr.config import configure
from rgltr.context import context
from rgltr.decorators import boundary, tag
from rgltr.pydantic_ai import RegulatedMCPAdapter


class Tag:
    DOCS = "docs"


class Boundary:
    EXTERNAL = "external"


def _run_ctx() -> RunContext:
    return RunContext(deps=None, model=Mock(), usage=RunUsage())


@pytest.mark.integration
@pytest.mark.anyio
async def test_remote_filters_tools():
    sse_url = "https://mcp.context7.com/mcp"
    mcp_server = MCPServerStreamableHTTP(url=sse_url)
    adapter = RegulatedMCPAdapter(
        wrapped=mcp_server,
        tools={"resolve-library-id": True},
    )

    async with adapter:
        tools = await adapter.get_tools(_run_ctx())
        assert "resolve-library-id" in tools
        assert len(tools) == 1


@pytest.mark.integration
@pytest.mark.anyio
async def test_remote_call_activates_tag():
    sse_url = "https://mcp.context7.com/mcp"
    mcp_server = MCPServerStreamableHTTP(url=sse_url)
    adapter = RegulatedMCPAdapter(
        wrapped=mcp_server,
        tools={"resolve-library-id": [tag(activates=[Tag.DOCS])]},
    )

    async with adapter:
        tools = await adapter.get_tools(_run_ctx())
        tool = tools["resolve-library-id"]
        with context() as ctx:
            await adapter.call_tool(
                "resolve-library-id",
                {"libraryName": "react", "query": "react hooks"},
                _run_ctx(),
                tool,
            )
            assert Tag.DOCS in ctx.active_tags


@pytest.mark.integration
@pytest.mark.anyio
async def test_remote_visibility_after_tag_activation():
    configure(boundaries={Boundary.EXTERNAL: [Tag.DOCS]})
    sse_url = "https://mcp.context7.com/mcp"
    mcp_server = MCPServerStreamableHTTP(url=sse_url)
    adapter = RegulatedMCPAdapter(
        wrapped=mcp_server,
        tools={
            "resolve-library-id": [tag(activates=[Tag.DOCS])],
            "query-docs": [boundary(Boundary.EXTERNAL)],
        },
    )

    async with adapter:
        with context():
            tools = await adapter.get_tools(_run_ctx())
            assert "query-docs" in tools

            tool = tools["resolve-library-id"]
            await adapter.call_tool(
                "resolve-library-id",
                {"libraryName": "react", "query": "react hooks"},
                _run_ctx(),
                tool,
            )

            tools_after = await adapter.get_tools(_run_ctx())
            assert "query-docs" not in tools_after
