from enum import Enum
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.tools import Tool
from pydantic_ai.usage import RunUsage

from rgltr.config import configure
from rgltr.context import context
from rgltr.decorators import boundary, governed, policy, tag
from rgltr.pydantic_ai import RegulatedToolset
from rgltr.types import PolicyDecision, PolicyRequest


class Tag(Enum):
    CUSTOMER = "customer"


class Boundary(Enum):
    EXTERNAL = "external"


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


@tag(activates=[Tag.CUSTOMER])
async def get_customer() -> str:
    return "ok"


@boundary(Boundary.EXTERNAL)
async def post_to_slack(message: str) -> str:
    return message


@policy(require=[deny_policy])
async def restricted() -> str:
    return "nope"


@boundary(Boundary.EXTERNAL)
async def gated_slack(message: str) -> str:
    return message


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


@pytest.mark.anyio
async def test_visibility_changes_after_tag_activation():
    configure(boundaries={Boundary.EXTERNAL: True})
    toolset = RegulatedToolset([get_customer, post_to_slack])

    with context():
        tools = await toolset.get_tools(_run_ctx())
        assert "get_customer" in tools
        assert "post_to_slack" in tools

        tool = tools["get_customer"]
        await toolset.call_tool("get_customer", {}, _run_ctx(), tool)

        tools_after = await toolset.get_tools(_run_ctx())
        assert "post_to_slack" not in tools_after


@pytest.mark.anyio
async def test_policy_denial_raises_model_retry():
    toolset = RegulatedToolset([restricted])
    tools = await toolset.get_tools(_run_ctx())
    tool = tools["restricted"]
    with context():
        with pytest.raises(ModelRetry):
            await toolset.call_tool("restricted", {}, _run_ctx(), tool)


@pytest.mark.anyio
async def test_call_tool_blocks_hidden_tools():
    configure(boundaries={Boundary.EXTERNAL: True})
    toolset = RegulatedToolset([gated_slack, get_customer])
    tools = await toolset.wrapped.get_tools(_run_ctx())
    tool = tools["gated_slack"]

    with context() as ctx:
        ctx.activate([Tag.CUSTOMER])
        with pytest.raises(ModelRetry):
            await toolset.call_tool("gated_slack", {"message": "hi"}, _run_ctx(), tool)


@pytest.mark.anyio
async def test_monitor_mode_allows_denied_policy():
    configure(mode="monitor")
    toolset = RegulatedToolset([restricted])
    tools = await toolset.get_tools(_run_ctx())
    tool = tools["restricted"]
    with context():
        await toolset.call_tool("restricted", {}, _run_ctx(), tool)


@pytest.mark.anyio
async def test_tool_object_metadata_and_rules():
    configure(boundaries={Boundary.EXTERNAL: True})
    tool_obj = Tool(post_to_slack, name="send_slack", description="custom", max_retries=5)
    toolset = RegulatedToolset([tool_obj, get_customer])

    with context():
        tools = await toolset.get_tools(_run_ctx())
        assert "send_slack" in tools
        assert tools["send_slack"].tool_def.description == "custom"
        assert tools["send_slack"].max_retries == 5

        customer_tool = tools["get_customer"]
        await toolset.call_tool("get_customer", {}, _run_ctx(), customer_tool)

        tools_after = await toolset.get_tools(_run_ctx())
        assert "send_slack" not in tools_after


@pytest.mark.anyio
async def test_toolset_decorator_metadata_and_rules():
    configure(boundaries={Boundary.EXTERNAL: True})
    toolset_decorator = FunctionToolset()

    @toolset_decorator.tool(name="send_slack", description="custom", retries=3)
    @boundary(Boundary.EXTERNAL)
    async def post_to_slack_v2(message: str) -> str:
        return message

    toolset_decorator.add_function(get_customer)
    toolset = RegulatedToolset(toolset=toolset_decorator)

    with context():
        tools = await toolset.get_tools(_run_ctx())
        assert "send_slack" in tools
        assert tools["send_slack"].tool_def.description == "custom"
        assert tools["send_slack"].max_retries == 3

        customer_tool = tools["get_customer"]
        await toolset.call_tool("get_customer", {}, _run_ctx(), customer_tool)

        tools_after = await toolset.get_tools(_run_ctx())
        assert "send_slack" not in tools_after


@pytest.mark.anyio
async def test_decorator_order_preserves_rules():
    configure(boundaries={Boundary.EXTERNAL: True})
    toolset_decorator = FunctionToolset()

    @boundary(Boundary.EXTERNAL)
    @toolset_decorator.tool(name="send_slack_ordered", description="custom")
    async def post_to_slack_v3(message: str) -> str:
        return message

    toolset_decorator.add_function(get_customer)
    toolset = RegulatedToolset(toolset=toolset_decorator)

    with context():
        tools = await toolset.get_tools(_run_ctx())
        assert "send_slack_ordered" in tools

        customer_tool = tools["get_customer"]
        await toolset.call_tool("get_customer", {}, _run_ctx(), customer_tool)

        tools_after = await toolset.get_tools(_run_ctx())
        assert "send_slack_ordered" not in tools_after


@pytest.mark.anyio
async def test_governed_annotation_counts_for_require_annotations():
    configure(require_annotations=True)

    @governed()
    async def audited_tool() -> str:
        return "ok"

    toolset = RegulatedToolset([audited_tool])

    with context():
        tools = await toolset.get_tools(_run_ctx())
        assert "audited_tool" in tools
