import briefcase_ai as bf
from typing import Any, Dict, List, Optional

def wrap_swarm(client):
    """
    Instruments OpenAI Swarm to capture handoffs and execution traces.
    """
    original_run = client.run

    def tracked_run(*args, **kwargs):
        # 1. Capture initial intent
        agent = kwargs.get("agent")
        messages = kwargs.get("messages", [])
        
        main_decision = bf.Decision(f"swarm_execution_{agent.name}")
        main_decision.add_tag("framework", "openai_swarm")
        main_decision.add_tag("initial_agent", agent.name)
        
        # 2. Execute original run
        response = original_run(*args, **kwargs)
        
        # 3. Process trajectory from response
        # Swarm response contains messages and the last active agent
        if hasattr(response, 'messages'):
            for i, msg in enumerate(response.messages):
                if msg.get("tool_calls"):
                    for tc in msg["tool_calls"]:
                        # Capture tool calls as sub-decisions or tags
                        main_decision.add_tag(f"step_{i}_tool", tc["function"]["name"])

        if hasattr(response, 'agent'):
            main_decision.add_tag("final_agent", response.agent.name)
            
        bf.save(main_decision)
        return response

    client.run = tracked_run
    return client
