"""
Parallel Research Teams - TRUE Parallel Execution with Broadcast Delegation

This example demonstrates TRUE PARALLEL workspace execution using broadcast delegation:

1. Research Coordinator receives a multi-topic research request
2. Coordinator uses BROADCAST to delegate to ALL 3 teams in ONE turn
   {"send_to": ["AI Research Lead", "Climate Research Lead", "Space Research Lead"]}
3. All 3 teams receive work SIMULTANEOUSLY and execute in PARALLEL
4. Each team has 2 agents (Lead Researcher + Assistant)
5. Teams complete work independently and return results concurrently
6. Coordinator waits for all results, then synthesizes final report

🚀 BROADCAST DELEGATION:
The Coordinator sends to multiple recipients in a single response, enabling true
parallel delegation. All teams receive their tasks at the same moment, not sequentially.

⚡ PARALLEL EXECUTION:
The WorkspaceExecutionEngine runs all workspaces concurrently using asyncio.gather,
so all 3 teams work simultaneously. This provides significant speedup for independent
workloads:
- Sequential: Time = Team₁ + Team₂ + Team₃
- Parallel: Time ≈ max(Team₁, Team₂, Team₃)
- Potential 3x speedup!

Architecture:
- 1 Root workspace: Research Coordinator
- 3 Child workspaces: AI Team, Climate Team, Space Team (run in parallel)
- Each child workspace: 2 agents (Lead Researcher + Research Assistant)

Setup:
1. install: pip install synqed anthropic python-dotenv
2. create .env file with: ANTHROPIC_API_KEY='your-key-here'
3. run: python parallel_research_teams.py
"""
import asyncio
import os
import logging
from pathlib import Path
from dotenv import load_dotenv

# Import the synqed API
import synqed

# Load environment variables
load_dotenv()
load_dotenv(dotenv_path=Path(__file__).parent / '.env')

# Configure logging
logging.basicConfig(
    level=logging.WARNING,
    format='%(asctime)s | %(name)s | %(levelname)s | %(message)s'
)


# ============================================================================
# Research Coordinator (Orchestrator)
# ============================================================================

async def coordinator_logic(context: synqed.AgentLogicContext) -> dict:
    """
    Research Coordinator delegates to multiple parallel research teams using BROADCAST.
    
    Uses broadcast delegation to send to all 3 teams simultaneously in ONE turn:
    {"send_to": ["Team1", "Team2", "Team3"], "content": "..."}
    
    This achieves TRUE PARALLEL delegation - all teams receive work at the same time!
    """
    import anthropic
    
    # Get latest message
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("USER", "Ready to coordinate parallel research!")
    
    # Check delegation state by counting messages sent
    conversation = context.get_conversation_history(format="raw")
    
    teams_received_from = set()
    
    # Check if we've already delegated (look for any messages to research leads)
    has_delegated = False
    for msg in conversation:
        sender = msg.get("sender", "")
        recipient = msg.get("recipient", "")
        
        if sender == "Research Coordinator" and "Research Lead" in recipient:
            has_delegated = True
        
        if recipient == "Research Coordinator" and "Research Lead" in sender:
            teams_received_from.add(sender)
    
    # DELEGATION PHASE: Broadcast to ALL teams at once
    if not has_delegated:
        user_message = latest_message.content if latest_message else ""
        
        # Extract the 3 topics from the user's request
        topics = {
            "AI Research Lead": "Latest advances in Large Language Models and AI agents",
            "Climate Research Lead": "Recent breakthroughs in climate change mitigation technology",
            "Space Research Lead": "Current status of Mars exploration missions"
        }
        
        # Create a task message that each team can interpret
        task = (
            f"Please research and provide a brief summary (100-150 words) on your specialty topic. "
            f"Original request:\n{user_message}"
        )
        
        # 🚀 BROADCAST: Send to ALL 3 teams simultaneously!
        return {
            "send_to": ["AI Research Lead", "Climate Research Lead", "Space Research Lead"],
            "content": task
        }
    
    # SYNTHESIS PHASE: Wait for all responses, then synthesize
    if len(teams_received_from) < 3:
        # Still waiting - shouldn't be called, but return a wait message
        return context.build_response("USER", f"Coordinating parallel research across {3 - len(teams_received_from)} remaining teams...")
    
    # All responses received - synthesize
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = (
        "You are a Research Coordinator. You have received research reports from 3 specialized teams "
        "(AI, Climate, Space). Synthesize them into one comprehensive summary for the user."
    )
    
    conversation_text = context.get_conversation_history()
    conversation_text += "\n\nSynthesize all 3 reports into one comprehensive summary. Respond with JSON: {\"send_to\": \"USER\", \"content\": \"synthesis\"}"
    
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=600,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    return resp.content[0].text.strip()


# ============================================================================
# Lead Researcher Logic (Generic for all teams)
# ============================================================================

async def lead_researcher_logic(context: synqed.AgentLogicContext) -> dict:
    """Generic lead researcher that delegates to assistant and reports back."""
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    # Get agent name to determine specialty
    agent_name = context.agent_name or "Lead Researcher"
    
    if "AI" in agent_name:
        specialty = "artificial intelligence and machine learning"
        topic_focus = "Latest advances in Large Language Models and AI agents"
        assistant = "AI Research Assistant"
    elif "Climate" in agent_name:
        specialty = "climate change and environmental science"
        topic_focus = "Recent breakthroughs in climate change mitigation technology"
        assistant = "Climate Research Assistant"
    elif "Space" in agent_name:
        specialty = "space exploration and astronomy"
        topic_focus = "Current status of Mars exploration missions"
        assistant = "Space Research Assistant"
    else:
        specialty = "general research"
        topic_focus = "general research topic"
        assistant = "Research Assistant"
    
    system_prompt = (
        f"You are a Lead Researcher specializing in {specialty}. "
        f"Your research focus: {topic_focus}\n"
        f"You have one assistant: {assistant}.\n\n"
        "Workflow:\n"
        "1. Receive research task from Research Coordinator\n"
        "2. Delegate focused research to your Assistant\n"
        "3. Review Assistant's findings\n"
        "4. Send comprehensive report (100-150 words) to Research Coordinator\n\n"
        "Keep research focused (2-3 rounds max).\n"
        "Always respond with JSON: "
        f'{{"send_to": "{assistant}|Research Coordinator", "content": "message"}}'
    )
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response(assistant, "Ready to lead research!")
    
    conversation_text = context.get_conversation_history()
    conversation_text += "\n\nRespond with JSON: {\"send_to\": \"target\", \"content\": \"message\"}"
    
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=300,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    return resp.content[0].text.strip()


# ============================================================================
# Research Assistant Logic (Generic for all teams)
# ============================================================================

async def research_assistant_logic(context: synqed.AgentLogicContext) -> dict:
    """Generic research assistant that conducts research and reports to lead."""
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    # Get agent name to determine specialty
    agent_name = context.agent_name or "Research Assistant"
    
    if "AI" in agent_name:
        specialty = "artificial intelligence and machine learning"
        lead = "AI Research Lead"
    elif "Climate" in agent_name:
        specialty = "climate change and environmental science"
        lead = "Climate Research Lead"
    elif "Space" in agent_name:
        specialty = "space exploration and astronomy"
        lead = "Space Research Lead"
    else:
        specialty = "general research"
        lead = "Lead Researcher"
    
    system_prompt = (
        f"You are a Research Assistant specializing in {specialty}. "
        f"You conduct research and report findings to {lead}.\n"
        "Provide concise, factual research (150 words max).\n"
        "Always respond with JSON: "
        f'{{"send_to": "{lead}", "content": "your research findings"}}'
    )
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response(lead, "Ready to conduct research!")
    
    conversation_text = f"{lead} asked: {latest_message.content}\n\n"
    conversation_text += "Provide research findings. Respond with JSON: {\"send_to\": \"" + lead + "\", \"content\": \"findings\"}"
    
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=250,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    return resp.content[0].text.strip()


# ============================================================================
# Main Execution
# ============================================================================

async def main():
    print("\n" + "="*80)
    print("  🚀 Parallel Research Teams Demo - TRUE PARALLEL EXECUTION")
    print("  Coordinator BROADCASTS to 3 Teams → All work SIMULTANEOUSLY")
    print("="*80 + "\n")
    
    # Step 1: Register all agent prototypes
    coordinator = synqed.Agent(
        name="Research Coordinator",
        description="Coordinates parallel research across multiple teams",
        logic=coordinator_logic,
        default_target="USER"
    )
    
    # AI Team
    ai_lead = synqed.Agent(
        name="AI Research Lead",
        description="Lead researcher for AI topics",
        logic=lead_researcher_logic,
        default_target="AI Research Assistant"
    )
    
    ai_assistant = synqed.Agent(
        name="AI Research Assistant",
        description="Research assistant for AI topics",
        logic=research_assistant_logic,
        default_target="AI Research Lead"
    )
    
    # Climate Team
    climate_lead = synqed.Agent(
        name="Climate Research Lead",
        description="Lead researcher for climate topics",
        logic=lead_researcher_logic,
        default_target="Climate Research Assistant"
    )
    
    climate_assistant = synqed.Agent(
        name="Climate Research Assistant",
        description="Research assistant for climate topics",
        logic=research_assistant_logic,
        default_target="Climate Research Lead"
    )
    
    # Space Team
    space_lead = synqed.Agent(
        name="Space Research Lead",
        description="Lead researcher for space topics",
        logic=lead_researcher_logic,
        default_target="Space Research Assistant"
    )
    
    space_assistant = synqed.Agent(
        name="Space Research Assistant",
        description="Research assistant for space topics",
        logic=research_assistant_logic,
        default_target="Space Research Lead"
    )
    
    # Register all agents (including coordinator to itself for self-messaging)
    synqed.AgentRuntimeRegistry.register("Research Coordinator", coordinator)
    synqed.AgentRuntimeRegistry.register("AI Research Lead", ai_lead)
    synqed.AgentRuntimeRegistry.register("AI Research Assistant", ai_assistant)
    synqed.AgentRuntimeRegistry.register("Climate Research Lead", climate_lead)
    synqed.AgentRuntimeRegistry.register("Climate Research Assistant", climate_assistant)
    synqed.AgentRuntimeRegistry.register("Space Research Lead", space_lead)
    synqed.AgentRuntimeRegistry.register("Space Research Assistant", space_assistant)
    
    print("✅ Registered 7 agents (1 coordinator + 3 teams of 2)\n")
    
    # Step 2: Create workspace manager and planner
    workspace_manager = synqed.WorkspaceManager(workspaces_root=Path("/tmp/synqed_parallel_research"))
    
    planner = synqed.PlannerLLM(
        provider="anthropic",
        api_key=os.environ["ANTHROPIC_API_KEY"],
        model="claude-sonnet-4-5"
    )
    
    # Step 3: Create workspace hierarchy
    root_task_node = synqed.TaskTreeNode(
        id="root",
        description="Coordinate parallel research",
        required_agents=["Research Coordinator"],
        may_need_subteams=True,
        children=[
            synqed.TaskTreeNode(
                id="ai-team",
                description="AI research",
                required_agents=["AI Research Lead", "AI Research Assistant"],
                may_need_subteams=False,
                children=[]
            ),
            synqed.TaskTreeNode(
                id="climate-team",
                description="Climate research",
                required_agents=["Climate Research Lead", "Climate Research Assistant"],
                may_need_subteams=False,
                children=[]
            ),
            synqed.TaskTreeNode(
                id="space-team",
                description="Space research",
                required_agents=["Space Research Lead", "Space Research Assistant"],
                may_need_subteams=False,
                children=[]
            )
        ]
    )
    
    # Create root workspace
    root_workspace = await workspace_manager.create_workspace(
        task_tree_node=root_task_node,
        parent_workspace_id=None
    )
    
    print(f"✅ Created root workspace: {root_workspace.workspace_id}")
    print(f"   Agent: {list(root_workspace.agents.keys())}\n")
    
    # Create AI Team workspace
    ai_team_node = root_task_node.children[0]
    ai_workspace = await workspace_manager.create_workspace(
        task_tree_node=ai_team_node,
        parent_workspace_id=root_workspace.workspace_id
    )
    
    print(f"✅ Created AI Team workspace: {ai_workspace.workspace_id}")
    print(f"   Agents: {list(ai_workspace.agents.keys())}\n")
    
    # Create Climate Team workspace
    climate_team_node = root_task_node.children[1]
    climate_workspace = await workspace_manager.create_workspace(
        task_tree_node=climate_team_node,
        parent_workspace_id=root_workspace.workspace_id
    )
    
    print(f"✅ Created Climate Team workspace: {climate_workspace.workspace_id}")
    print(f"   Agents: {list(climate_workspace.agents.keys())}\n")
    
    # Create Space Team workspace
    space_team_node = root_task_node.children[2]
    space_workspace = await workspace_manager.create_workspace(
        task_tree_node=space_team_node,
        parent_workspace_id=root_workspace.workspace_id
    )
    
    print(f"✅ Created Space Team workspace: {space_workspace.workspace_id}")
    print(f"   Agents: {list(space_workspace.agents.keys())}\n")
    
    # Step 4: Create execution engine
    execution_engine = synqed.WorkspaceExecutionEngine(
        planner=planner,
        workspace_manager=workspace_manager,
        enable_display=True,
        max_agent_turns=30,  # Allow enough turns for 3 parallel teams
        max_workspace_depth=3,
    )
    
    print("✅ Created execution engine with parallel workspace support\n")
    
    # Step 5: Send initial task
    task = (
        "Research three cutting-edge topics:\n"
        "1. Latest advances in Large Language Models and AI agents\n"
        "2. Recent breakthroughs in climate change mitigation technology\n"
        "3. Current status of Mars exploration missions\n\n"
        "Each team should provide a brief summary (100-150 words) of the most important recent developments."
    )
    
    print(f"📋 Task: {task}\n")
    print("="*80)
    print("  EXECUTION LOG (Watch for PARALLEL execution)")
    print("="*80 + "\n")
    
    await root_workspace.route_message("USER", "Research Coordinator", task, manager=workspace_manager)
    
    # Step 6: Execute (all teams will run in parallel)
    import time
    start_time = time.time()
    
    try:
        await execution_engine.run(root_workspace.workspace_id)
    except Exception as e:
        print(f"\n❌ Error during execution: {e}\n")
        import traceback
        traceback.print_exc()
        raise
    
    elapsed_time = time.time() - start_time
    
    # Step 7: Display results
    print("\n" + "="*80)
    print("  EXECUTION SUMMARY")
    print("="*80 + "\n")
    
    print(f"⏱️  Execution time: {elapsed_time:.2f} seconds\n")
    
    print("📊 Workspace Hierarchy:")
    print(f"   Root: {root_workspace.workspace_id} (Coordinator)")
    print(f"   ├─ AI Team: {ai_workspace.workspace_id}")
    print(f"   ├─ Climate Team: {climate_workspace.workspace_id}")
    print(f"   └─ Space Team: {space_workspace.workspace_id}")
    print()
    
    # Get completion status for each workspace
    root_status = root_workspace.get_completion_status()
    ai_status = ai_workspace.get_completion_status()
    climate_status = climate_workspace.get_completion_status()
    space_status = space_workspace.get_completion_status()
    
    print("📈 Message Statistics:")
    print(f"   Root: {root_status['total_messages']} messages")
    print(f"   AI Team: {ai_status['total_messages']} messages")
    print(f"   Climate Team: {climate_status['total_messages']} messages")
    print(f"   Space Team: {space_status['total_messages']} messages")
    print(f"   Total: {root_status['total_messages'] + ai_status['total_messages'] + climate_status['total_messages'] + space_status['total_messages']} messages")
    print()
    
    print(f"✅ Status: {root_status['status_message']}\n")
    
    # Display transcripts
    root_workspace.display_transcript(title="ROOT - Research Coordinator")
    ai_workspace.display_transcript(title="AI TEAM")
    climate_workspace.display_transcript(title="CLIMATE TEAM")
    space_workspace.display_transcript(title="SPACE TEAM")
    
    # Clean up
    await workspace_manager.destroy_workspace(root_workspace.workspace_id)
    
    print("="*80)
    print("✅ Parallel execution demo complete!")
    print("="*80 + "\n")


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

