"""
Orchestrator with Two Teams - Hierarchical Multi-Workspace Demo

This demonstrates the orchestrator pattern with hierarchical workspaces:
1. Root workspace: Project Manager (orchestrator) plans and delegates
2. Research Team workspace: 3 agents (Research Lead, Data Analyst, Report Writer)
3. Development Team workspace: 3 agents (Tech Lead, Backend Dev, Frontend Dev)

Architecture:
- PlannerLLM creates a hierarchical task tree with subteams
- WorkspaceManager creates nested workspaces (1 root + 2 child workspaces)
- Each workspace has independent agents with their own memory
- Parent workspace can delegate to child workspaces
- Child workspaces return results to parent via [subteam_result] messages

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

# Import the Workspace-based API
import synqed
from synqed.workspace_manager import WorkspaceManager, AgentRuntimeRegistry
from synqed.execution_engine import WorkspaceExecutionEngine
from synqed.planner import PlannerLLM, TaskTreeNode

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

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


# ============================================================================
# Project Manager (Orchestrator) Logic
# ============================================================================

async def project_manager_logic(context: synqed.AgentLogicContext) -> dict:
    """
    Project Manager acts as an orchestrator, delegating tasks to specialized teams.
    
    NOW SIMPLIFIED with automatic API utilities!
    
    This agent:
    - Receives the overall project task from USER
    - Plans and breaks down the work
    - Delegates to Research Team and Development Team
    - Aggregates results and reports back to USER
    """
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = (
        "You are a Project Manager orchestrating a software project. "
        "You have two specialized team leads available:\n"
        "1. Research Lead - leads the research team (market research, analysis, documentation)\n"
        "2. Tech Lead - leads the development team (technical implementation, coding)\n\n"
        "Your workflow:\n"
        "1. Receive project task from USER\n"
        "2. Break it down and delegate to appropriate team leads\n"
        "3. Wait for team results (they'll send you their completed work)\n"
        "4. Synthesize results and report back to USER\n\n"
        "IMPORTANT: Send to 'Research Lead' or 'Tech Lead' (not 'Research Team' or 'Development Team').\n"
        "When delegating, explain what you need them to accomplish.\n"
        "Always respond with JSON: "
        '{"send_to": "Research Lead", "content": "..."} or {"send_to": "Tech Lead", "content": "..."} or {"send_to": "USER", "content": "..."}'
    )
    
    # Get latest message
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("USER", "I'm ready to manage your project!")
    
    # 🎯 NEW API: Get conversation history automatically!
    # Note: For orchestrators, we want to see subteam_result messages
    conversation_text = context.get_conversation_history(include_system_messages=True)
    
    # Add instruction to respond
    conversation_text += "\n\nRespond with JSON: {\"send_to\": \"target\", \"content\": \"your response\"}"
    
    # Call Anthropic API
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=600,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    # 🎯 NEW API: Just return the response! All parsing happens automatically
    return resp.content[0].text.strip()


# ============================================================================
# Research Team Agents (3 agents)
# ============================================================================

async def research_lead_logic(context: synqed.AgentLogicContext) -> dict:
    """Research Lead coordinates the research team - NOW SIMPLIFIED!"""
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = (
        "You are a Research Lead coordinating a research team. "
        "You have two teammates: Data Analyst and Report Writer.\n"
        "Workflow:\n"
        "1. Receive research task from Project Manager\n"
        "2. Delegate data gathering to Data Analyst\n"
        "3. Wait for analysis results\n"
        "4. Ask Report Writer to create final report\n"
        "5. Send completed report to 'Project Manager'\n"
        "Keep research focused and concise (2-3 rounds max).\n"
        "IMPORTANT: When research is complete, send final report to 'Project Manager'.\n"
        "Always respond with JSON: "
        '{"send_to": "Data Analyst|Report Writer|Project Manager", "content": "message"}'
    )
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("Data Analyst", "Ready to lead research!")
    
    # 🎯 NEW API: Get conversation history automatically!
    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=400,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    # 🎯 NEW API: Just return the response!
    return resp.content[0].text.strip()


async def data_analyst_logic(context: synqed.AgentLogicContext) -> dict:
    """Data Analyst gathers and analyzes data - NOW SIMPLIFIED!"""
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = (
        "You are a Data Analyst on a research team. "
        "You gather and analyze data, then report findings to the Research Lead.\n"
        "Provide concise, actionable insights (150 words max).\n"
        "Always respond with JSON: "
        '{"send_to": "Research Lead", "content": "your analysis"}'
    )
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("Research Lead", "Ready to analyze data!")
    
    # 🎯 NEW API: Simple conversation building
    conversation_text = (
        f"Research Lead asked: {latest_message.content}\n\n"
        "Provide your data analysis. Respond with JSON: {\"send_to\": \"Research Lead\", \"content\": \"your analysis\"}"
    )
    
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=300,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    # 🎯 NEW API: Just return the response!
    return resp.content[0].text.strip()


async def report_writer_logic(context: synqed.AgentLogicContext) -> dict:
    """Report Writer creates final research reports - NOW SIMPLIFIED!"""
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = (
        "You are a Report Writer on a research team. "
        "You create clear, well-structured reports from analysis data.\n"
        "Keep reports concise (200 words max) but comprehensive.\n"
        "Always respond with JSON: "
        '{"send_to": "Research Lead", "content": "your report"}'
    )
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("Research Lead", "Ready to write reports!")
    
    # 🎯 NEW API: Simple conversation building
    conversation_text = (
        f"Research Lead asked: {latest_message.content}\n\n"
        "Create the report. Respond with JSON: {\"send_to\": \"Research Lead\", \"content\": \"your report\"}"
    )
    
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=400,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    # 🎯 NEW API: Just return the response!
    return resp.content[0].text.strip()


# ============================================================================
# Development Team Agents (3 agents)
# ============================================================================

async def tech_lead_logic(context: synqed.AgentLogicContext) -> dict:
    """Tech Lead coordinates the development team - NOW SIMPLIFIED!"""
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = (
        "You are a Tech Lead coordinating a development team. "
        "You have two teammates: Backend Dev and Frontend Dev.\n"
        "Workflow:\n"
        "1. Receive development task from Project Manager\n"
        "2. Break down into backend and frontend work\n"
        "3. Delegate to Backend Dev and Frontend Dev\n"
        "4. Review their work\n"
        "5. Send completed implementation plan to 'Project Manager'\n"
        "Keep it concise (2-3 rounds max).\n"
        "IMPORTANT: When development plan is complete, send final report to 'Project Manager'.\n"
        "Always respond with JSON: "
        '{"send_to": "Backend Dev|Frontend Dev|Project Manager", "content": "message"}'
    )
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("Backend Dev", "Ready to lead development!")
    
    # 🎯 NEW API: Get conversation history automatically!
    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=400,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    # 🎯 NEW API: Just return the response!
    return resp.content[0].text.strip()


async def backend_dev_logic(context: synqed.AgentLogicContext) -> dict:
    """Backend Developer implements server-side features - NOW SIMPLIFIED!"""
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = (
        "You are a Backend Developer specializing in server-side implementation. "
        "You design APIs, databases, and backend architecture.\n"
        "Provide clear, technical implementation plans (150 words max).\n"
        "Always respond with JSON: "
        '{"send_to": "Tech Lead", "content": "your implementation plan"}'
    )
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("Tech Lead", "Ready to code backend!")
    
    conversation_text = (
        f"Tech Lead asked: {latest_message.content}\n\n"
        "Provide backend implementation plan. Respond with JSON: {\"send_to\": \"Tech Lead\", \"content\": \"your plan\"}"
    )
    
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=300,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    # 🎯 NEW API: Just return the response!
    return resp.content[0].text.strip()


async def frontend_dev_logic(context: synqed.AgentLogicContext) -> dict:
    """Frontend Developer implements client-side features - NOW SIMPLIFIED!"""
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = (
        "You are a Frontend Developer specializing in UI/UX implementation. "
        "You design user interfaces, components, and client-side logic.\n"
        "Provide clear, technical implementation plans (150 words max).\n"
        "Always respond with JSON: "
        '{"send_to": "Tech Lead", "content": "your implementation plan"}'
    )
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("Tech Lead", "Ready to code frontend!")
    
    # 🎯 NEW API: Simple conversation building
    conversation_text = (
        f"Tech Lead asked: {latest_message.content}\n\n"
        "Provide frontend implementation plan. Respond with JSON: {\"send_to\": \"Tech Lead\", \"content\": \"your plan\"}"
    )
    
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=300,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    # 🎯 NEW API: Just return the response!
    return resp.content[0].text.strip()


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

async def main():
    print("\n" + "="*80)
    print("  Multi-Team Orchestrator Demo")
    print("  1 Project Manager → 2 Teams (Research + Development, 3 agents each)")
    print("="*80 + "\n")
    
    # ========================================================================
    # Step 1: Register all agent prototypes
    # ========================================================================
    
    # Project Manager (orchestrator in root workspace)
    project_manager = synqed.Agent(
        name="Project Manager",
        description="Project manager who orchestrates teams",
        logic=project_manager_logic,
        default_target="USER"
    )
    
    # Research Team agents
    research_lead = synqed.Agent(
        name="Research Lead",
        description="Research team leader",
        logic=research_lead_logic,
        default_target="Data Analyst"
    )
    
    data_analyst = synqed.Agent(
        name="Data Analyst",
        description="Data analysis specialist",
        logic=data_analyst_logic,
        default_target="Research Lead"
    )
    
    report_writer = synqed.Agent(
        name="Report Writer",
        description="Report writing specialist",
        logic=report_writer_logic,
        default_target="Research Lead"
    )
    
    # Development Team agents
    tech_lead = synqed.Agent(
        name="Tech Lead",
        description="Development team leader",
        logic=tech_lead_logic,
        default_target="Backend Dev"
    )
    
    backend_dev = synqed.Agent(
        name="Backend Dev",
        description="Backend development specialist",
        logic=backend_dev_logic,
        default_target="Tech Lead"
    )
    
    frontend_dev = synqed.Agent(
        name="Frontend Dev",
        description="Frontend development specialist",
        logic=frontend_dev_logic,
        default_target="Tech Lead"
    )
    
    # Register all agents
    AgentRuntimeRegistry.register("Project Manager", project_manager)
    AgentRuntimeRegistry.register("Research Lead", research_lead)
    AgentRuntimeRegistry.register("Data Analyst", data_analyst)
    AgentRuntimeRegistry.register("Report Writer", report_writer)
    AgentRuntimeRegistry.register("Tech Lead", tech_lead)
    AgentRuntimeRegistry.register("Backend Dev", backend_dev)
    AgentRuntimeRegistry.register("Frontend Dev", frontend_dev)
    
    print("✅ Registered 7 agents (1 orchestrator + 2 teams of 3)\n")
    
    # ========================================================================
    # Step 2: Create workspace manager and planner
    # ========================================================================
    
    workspace_manager = WorkspaceManager(workspaces_root=Path("/tmp/synqed_orchestrator_demo"))
    
    planner = PlannerLLM(
        provider="anthropic",
        api_key=os.environ["ANTHROPIC_API_KEY"],
        model="claude-sonnet-4-5"
    )
    
    # ========================================================================
    # Step 3: Create hierarchical workspace structure
    # ========================================================================
    
    # Root workspace: Project Manager
    root_task_node = TaskTreeNode(
        id="root",
        description="Orchestrate project with specialized teams",
        required_agents=["Project Manager"],
        may_need_subteams=True,
        children=[
            TaskTreeNode(
                id="research-team",
                description="Research and analysis tasks",
                required_agents=["Research Lead", "Data Analyst", "Report Writer"],
                may_need_subteams=False,
                children=[]
            ),
            TaskTreeNode(
                id="dev-team",
                description="Development and implementation tasks",
                required_agents=["Tech Lead", "Backend Dev", "Frontend Dev"],
                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"   Agents: {list(root_workspace.agents.keys())}\n")
    
    # Create Research Team workspace (child of root)
    research_team_node = root_task_node.children[0]
    research_workspace = await workspace_manager.create_workspace(
        task_tree_node=research_team_node,
        parent_workspace_id=root_workspace.workspace_id
    )
    
    print(f"✅ Created Research Team workspace: {research_workspace.workspace_id}")
    print(f"   Agents: {list(research_workspace.agents.keys())}\n")
    
    # Create Development Team workspace (child of root)
    dev_team_node = root_task_node.children[1]
    dev_workspace = await workspace_manager.create_workspace(
        task_tree_node=dev_team_node,
        parent_workspace_id=root_workspace.workspace_id
    )
    
    print(f"✅ Created Development Team workspace: {dev_workspace.workspace_id}")
    print(f"   Agents: {list(dev_workspace.agents.keys())}\n")
    
    # ========================================================================
    # Step 4: Create execution engine
    # ========================================================================
    
    execution_engine = WorkspaceExecutionEngine(
        planner=planner,
        workspace_manager=workspace_manager,
        enable_display=True,  # Real-time message display
        max_agent_turns=25,  # Allow more turns for hierarchical execution
        max_workspace_depth=3,  # Root + 2 child workspaces
    )
    
    print("✅ Created execution engine\n")
    
    # ========================================================================
    # Step 5: Send initial task to Project Manager
    # ========================================================================
    
    task = (
        "Build a social media analytics dashboard. "
        "We need market research on competitors and user needs, "
        "plus a technical implementation plan for both backend and frontend."
    )
    
    print(f"📋 Task: {task}\n")
    print("="*80)
    print("  EXECUTION LOG (Real-time messages)")
    print("="*80 + "\n")
    
    await root_workspace.route_message("USER", "Project Manager", task, manager=workspace_manager)
    
    # ========================================================================
    # Step 6: Execute the entire workspace hierarchy
    # ========================================================================
    
    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
    
    # ========================================================================
    # Step 7: Display results - NOW USING NEW API UTILITIES!
    # ========================================================================
    
    print("\n" + "="*80)
    print("  EXECUTION SUMMARY")
    print("="*80 + "\n")
    
    # Display workspace hierarchy
    print("📊 Workspace Hierarchy:")
    print(f"   Root: {root_workspace.workspace_id} (Project Manager)")
    print(f"   ├─ Research Team: {research_workspace.workspace_id}")
    print(f"   │  └─ Agents: Research Lead, Data Analyst, Report Writer")
    print(f"   └─ Development Team: {dev_workspace.workspace_id}")
    print(f"      └─ Agents: Tech Lead, Backend Dev, Frontend Dev")
    print()
    
    # 🎯 NEW API: Get completion status for each workspace
    root_status = root_workspace.get_completion_status()
    research_status = research_workspace.get_completion_status()
    dev_status = dev_workspace.get_completion_status()
    
    print(f"📈 Message Statistics:")
    print(f"   Root workspace: {root_status['total_messages']} messages")
    print(f"   Research Team: {research_status['total_messages']} messages")
    print(f"   Development Team: {dev_status['total_messages']} messages")
    print(f"   Total: {root_status['total_messages'] + research_status['total_messages'] + dev_status['total_messages']} messages")
    print()
    
    # 🎯 NEW API: Use built-in completion status
    print(f"✅ Root Workspace Status:")
    print(f"   {root_status['status_message']}")
    print()
    
    # ========================================================================
    # Step 8: Display transcripts - NOW USING NEW API!
    # ========================================================================
    
    # 🎯 NEW API: Display transcripts using built-in utilities
    root_workspace.display_transcript(title="ROOT WORKSPACE - Project Manager")
    research_workspace.display_transcript(title="RESEARCH TEAM WORKSPACE")
    dev_workspace.display_transcript(title="DEVELOPMENT TEAM WORKSPACE")
    
    # ========================================================================
    # Step 9: Clean up
    # ========================================================================
    
    await workspace_manager.destroy_workspace(root_workspace.workspace_id)
    
    print("="*80)
    print("✅ Demo complete! All workspaces cleaned up.")
    print("="*80 + "\n")


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

