from unittest.mock import Mock

import pytest

pytest.importorskip("pydantic_ai")

from pydantic_ai._run_context import RunContext
from pydantic_ai.exceptions import ModelRetry
from pydantic_ai.toolsets.function import FunctionToolset
from pydantic_ai.usage import RunUsage

from enum import Enum

from rgltr.config import configure
from rgltr.context import context
from rgltr.decorators import boundary, policy
from rgltr.pydantic_ai import RegulatedMCPAdapter
from rgltr.types import PolicyDecision, PolicyRequest


class Tag(Enum):
    CUSTOMER = "customer"


class Boundary(Enum):
    EXTERNAL = "external"


async def tool_a() -> str:
    return "a"


async def tool_b() -> str:
    return "b"


async def deny_policy(_: PolicyRequest) -> PolicyDecision:
    return PolicyDecision.deny("blocked")


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


@pytest.mark.anyio
async def test_mcp_adapter_filters_tools():
    base = FunctionToolset([tool_a, tool_b])
    adapter = RegulatedMCPAdapter(
        wrapped=base,
        tools={"tool_a": True},
    )

    tools = await adapter.get_tools(_run_ctx())
    assert "tool_a" in tools
    assert "tool_b" not in tools


@pytest.mark.anyio
async def test_mcp_adapter_policy_denial():
    base = FunctionToolset([tool_a])
    adapter = RegulatedMCPAdapter(
        wrapped=base,
        tools={"tool_a": [policy(require=[deny_policy])]},
    )
    tools = await adapter.get_tools(_run_ctx())
    tool = tools["tool_a"]
    with pytest.raises(ModelRetry):
        await adapter.call_tool("tool_a", {}, _run_ctx(), tool)


@pytest.mark.anyio
async def test_mcp_adapter_blocks_hidden_tools():
    configure(boundaries={Boundary.EXTERNAL: True})
    base = FunctionToolset([tool_a])
    adapter = RegulatedMCPAdapter(
        wrapped=base,
        tools={"tool_a": [boundary(Boundary.EXTERNAL)]},
    )
    tools = await base.get_tools(_run_ctx())
    tool = tools["tool_a"]
    with context() as ctx:
        ctx.activate([Tag.CUSTOMER])
        with pytest.raises(ModelRetry):
            await adapter.call_tool("tool_a", {}, _run_ctx(), tool)
