"""Message history tag restoration.

Shows how tags are restored from prior tool calls when message_history is passed.
"""

import asyncio
from enum import Enum
from unittest.mock import Mock

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

from rgltr import boundary, configure, tag
from rgltr.context import context, restore_tags_from_messages
from rgltr.pydantic_ai import RegulatedAgent, RegulatedToolset


class Tag(Enum):
    FINANCIAL = "financial"


class Boundary(Enum):
    EXTERNAL = "external"


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


async def main() -> None:
    configure(boundaries={Boundary.EXTERNAL: True})

    model = TestModel(call_tools=["get_invoice"])
    agent = Agent(model)

    @agent.tool_plain
    @tag(activates=[Tag.FINANCIAL])
    async def get_invoice(invoice_id: int) -> str:
        return f"invoice:{invoice_id}"

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

    regulated = RegulatedAgent(agent, agent_id="history")
    toolset = RegulatedToolset([get_invoice, post_to_slack])

    with context():
        tools_before = await toolset.get_tools(_run_ctx())
    print("first tools:", sorted(tools_before.keys()))

    first = await regulated.run("first")

    def _lookup_activates(name: str) -> set[Enum] | None:
        rules = toolset.rules_for_tool(name)
        return set(rules.activates) if rules else None

    with context():
        restore_tags_from_messages(first.all_messages(), _lookup_activates)
        tools_after = await toolset.get_tools(_run_ctx())
    print("second tools:", sorted(tools_after.keys()))

    second = await regulated.run("second", message_history=first.all_messages())
    print("second output:", second.output)


if __name__ == "__main__":
    asyncio.run(main())
