import asyncio
import re

from langgraph.checkpoint.memory import MemorySaver
from rich import print

from universal_mcp.agentr.agent_builder_registry import AgentBuilderRegistry
from universal_mcp.agentr.client import AgentrClient
from universal_mcp.agentr.registry import AgentrRegistry
from universal_mcp.agents.codeact0.agent import CodeActPlaybookAgent
from universal_mcp.agents.utils import messages_to_list


async def main():
    memory = MemorySaver()

    # Initialize AgentrClient and AgentBuilderRegistry with specific agent_id
    client = AgentrClient()
    await client.ainit()

    agent_id = "851208ef-6a6e-4bee-a89d-4f69cd535e63"
    agent_builder_registry = AgentBuilderRegistry(client=client, default_agent_id=agent_id)

    # First, fetch and print the agent script to see what functions are defined
    agent_data = await agent_builder_registry.get_agent(agent_id)
    print("=" * 80)
    print("AGENT SCRIPT:")
    print("=" * 80)
    if agent_data and hasattr(agent_data, "instructions"):
        script = agent_data.instructions.get("script", "No script found")
        print(script)
        print("=" * 80)

        # Count function definitions
        functions = re.findall(r"^(?:async\s+)?def\s+(\w+)\s*\(", script, re.MULTILINE)
        print(f"\nDefined functions ({len(functions)}):")
        for func in functions:
            print(f"  - {func}")
        print("=" * 80)

    # Create AgentrRegistry
    registry = AgentrRegistry(client=client)

    print("\nStarting agent...")

    # Test case 1: Test if helper functions are accessible
    async with CodeActPlaybookAgent(
        name="CodeAct Agent",
        instructions="Be very concise in your answers.",
        model="anthropic:claude-haiku-4-5",
        registry=registry,
        memory=memory,
        agent_builder_registry=agent_builder_registry,
    ) as agent:
        result = await agent.invoke(
            user_input="Hi, what can you do?",
            thread_id="debug-123",
        )
        print("\n" + "=" * 80)
        print("TEST 1: Agent capabilities")
        print("=" * 80)
        print(messages_to_list(result["messages"]))

    # Test case 2: Actually call the function with a new agent instance
    async with CodeActPlaybookAgent(
        name="CodeAct Agent",
        instructions="Be very concise in your answers.",
        model="anthropic:claude-haiku-4-5",
        registry=registry,
        memory=memory,
        agent_builder_registry=agent_builder_registry,
    ) as agent:
        result2 = await agent.invoke(
            user_input="search scenes for portrait",
            thread_id="debug-123",
        )
        print("\n" + "=" * 80)
        print("TEST 2: Search for portrait scenes")
        print("=" * 80)
        print(messages_to_list(result2["messages"]))


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