from __future__ import annotations

from enum import Enum

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 pydantic_ai.toolsets.function import FunctionToolset

from rgltr.config import configure
from rgltr.decorators import boundary, tag
from rgltr.pydantic_ai import RegulatedAgent


class Tag(Enum):
    CUSTOMER = "customer"


class Boundary(Enum):
    EXTERNAL = "external"


def _tool_names(model: TestModel) -> list[str]:
    assert model.last_model_request_parameters is not None
    return [tool.name for tool in model.last_model_request_parameters.function_tools]


@pytest.mark.anyio
async def test_regulated_agent_wraps_agent_tool_decorator():
    configure(boundaries={Boundary.EXTERNAL: True})
    model = TestModel(call_tools="all")
    agent = Agent(model)

    @agent.tool
    @tag(activates=[Tag.CUSTOMER])
    async def get_customer(ctx: RunContext[None], customer_id: int) -> str:
        return f"customer:{customer_id}"

    @agent.tool
    @boundary(Boundary.EXTERNAL)
    async def post_to_slack(ctx: RunContext[None], message: str) -> str:
        return message

    regulated = RegulatedAgent(agent, agent_id="test")

    first = await regulated.run("hi")
    messages = first.all_messages()

    await regulated.run("hi again", message_history=messages)
    tool_names = _tool_names(model)

    assert "get_customer" in tool_names
    assert "post_to_slack" not in tool_names


@pytest.mark.anyio
async def test_regulated_agent_wraps_agent_tool_plain():
    configure(boundaries={Boundary.EXTERNAL: True})
    model = TestModel(call_tools="all")
    agent = Agent(model)

    @agent.tool_plain
    @tag(activates=[Tag.CUSTOMER])
    async def get_customer(customer_id: int) -> str:
        return f"customer:{customer_id}"

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

    regulated = RegulatedAgent(agent, agent_id="test")

    first = await regulated.run("hi")
    messages = first.all_messages()

    await regulated.run("hi again", message_history=messages)
    tool_names = _tool_names(model)

    assert "get_customer" in tool_names
    assert "post_to_slack" not in tool_names


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

    @tag(activates=[Tag.CUSTOMER])
    async def get_customer(customer_id: int) -> str:
        return f"customer:{customer_id}"

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

    model = TestModel(call_tools="all")
    agent = Agent(model, tools=[get_customer, post_to_slack])
    regulated = RegulatedAgent(agent, agent_id="test")

    first = await regulated.run("hi")
    messages = first.all_messages()

    await regulated.run("hi again", message_history=messages)
    tool_names = _tool_names(model)

    assert "get_customer" in tool_names
    assert "post_to_slack" not in tool_names


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

    @tag(activates=[Tag.CUSTOMER])
    async def get_customer(customer_id: int) -> str:
        return f"customer:{customer_id}"

    toolset = FunctionToolset()

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

    async def toolset_factory(_: RunContext[None]) -> FunctionToolset:
        return toolset

    model = TestModel(call_tools="all")
    agent = Agent(model, tools=[get_customer], toolsets=[toolset_factory])
    regulated = RegulatedAgent(agent, agent_id="test")

    first = await regulated.run("hi")
    messages = first.all_messages()

    await regulated.run("hi again", message_history=messages)
    tool_names = _tool_names(model)

    assert "get_customer" in tool_names
    assert "post_to_slack" not in tool_names
