from __future__ import annotations

from enum import Enum
from unittest.mock import Mock

import pytest

pytest.importorskip("pydantic_ai")

from pydantic_ai import Agent
from pydantic_ai._run_context import RunContext
from pydantic_ai.models.test import TestModel

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


class Tag(Enum):
    SECRET = "secret"


class Boundary(Enum):
    EXTERNAL = "external"


class DummyBuiltin:
    kind = "dummy_builtin"

    def __init__(self, name: str = "dummy") -> None:
        self.name = name


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


def _run_ctx() -> RunContext:
    return RunContext(deps=None, model=Mock(), usage=None)  # type: ignore[arg-type]


@pytest.mark.anyio
async def test_builtin_tools_governed_by_boundary():
    configure(boundaries={Boundary.EXTERNAL: [Tag.SECRET]})
    agent = Agent(TestModel())
    regulated = RegulatedAgent(agent, agent_id="test", govern_builtin_tools=True)

    @boundary(Boundary.EXTERNAL)
    async def dynamic_builtin(_: RunContext[None]) -> DummyBuiltin:
        return DummyBuiltin("blocked")

    wrapped = regulated._wrap_builtin_tool(dynamic_builtin)

    with context() as ctx:
        ctx.activate({Tag.SECRET})
        result = await wrapped(_run_ctx())

    assert result is None


@pytest.mark.anyio
async def test_builtin_tools_policy_denial_blocks():
    agent = Agent(TestModel())
    regulated = RegulatedAgent(agent, agent_id="test", govern_builtin_tools=True)

    @policy(require=[deny_policy])
    async def dynamic_builtin(_: RunContext[None]) -> DummyBuiltin:
        return DummyBuiltin("blocked")

    wrapped = regulated._wrap_builtin_tool(dynamic_builtin)

    with context():
        result = await wrapped(_run_ctx())

    assert result is None


@pytest.mark.anyio
async def test_builtin_tools_require_annotations():
    configure(require_annotations=True)
    agent = Agent(TestModel())
    regulated = RegulatedAgent(agent, agent_id="test", govern_builtin_tools=True)

    async def dynamic_builtin(_: RunContext[None]) -> DummyBuiltin:
        return DummyBuiltin("missing")

    wrapped = regulated._wrap_builtin_tool(dynamic_builtin)

    with context():
        with pytest.raises(ValueError, match="lacks rgltr annotations"):
            await wrapped(_run_ctx())


@pytest.mark.anyio
async def test_builtin_tools_governed_annotation_counts():
    configure(require_annotations=True)
    agent = Agent(TestModel())
    regulated = RegulatedAgent(agent, agent_id="test", govern_builtin_tools=True)

    @governed()
    async def dynamic_builtin(_: RunContext[None]) -> DummyBuiltin:
        return DummyBuiltin("ok")

    wrapped = regulated._wrap_builtin_tool(dynamic_builtin)

    with context():
        result = await wrapped(_run_ctx())

    assert result is not None
