"""
Orchestrator with Two Teams - Hierarchical multi-workspace coordination

This demonstrates:
1. A root workspace with an Orchestrator agent that plans and delegates
2. Two child workspaces (Team A: Marketing, Team B: Engineering)
3. Each team has 3 specialized agents working collaboratively
4. Teams complete their subtasks and report back to the Orchestrator
5. Orchestrator synthesizes the results and sends to USER

Architecture:
- Root Workspace: [Orchestrator]
  - Team A Workspace: [Marketing Manager, Content Writer, Designer]
  - Team B Workspace: [Tech Lead, Developer, QA Engineer]

The Orchestrator delegates:
- Marketing tasks → Team A
- Technical tasks → Team B

Teams coordinate internally and report back to Orchestrator.

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 json
import logging
from pathlib import Path
from dotenv import load_dotenv

# Import synqed components
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
logging.basicConfig(
    level=logging.WARNING,
    format='%(asctime)s | %(name)s | %(levelname)s | %(message)s'
)


# ============================================================================
# Orchestrator Agent
# ============================================================================

async def orchestrator_logic(context: synqed.AgentLogicContext) -> dict:
    """
    Orchestrator agent - plans, delegates to teams, and synthesizes results.
    """
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = """You are a strategic orchestrator coordinating two specialized teams:
- Team A (Marketing): Marketing Manager, Content Writer, Designer
- Team B (Engineering): Tech Lead, Developer, QA Engineer

Your role:
1. Analyze incoming tasks and break them into subtasks
2. Delegate marketing tasks to "Marketing Manager" (Team A lead)
3. Delegate technical tasks to "Tech Lead" (Team B lead)
4. When you receive results from both team leads, synthesize and send final report to USER

Response format:
{"send_to": "Marketing Manager", "content": "task description"}
or
{"send_to": "Tech Lead", "content": "task description"}
or
{"send_to": "USER", "content": "final synthesized report"}

Be concise. Delegate clearly. Complete in 3-5 total exchanges."""
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("USER", "Orchestrator ready to coordinate teams.")
    
    # Build conversation history
    conversation_parts = []
    for msg in context.memory.get_messages():
        if msg.from_agent == "USER":
            conversation_parts.append(f"User: {msg.content}")
        elif msg.from_agent == "Marketing Manager":
            conversation_parts.append(f"Marketing Manager (Team A): {msg.content}")
        elif msg.from_agent == "Tech Lead":
            conversation_parts.append(f"Tech Lead (Team B): {msg.content}")
        elif msg.from_agent != "SYSTEM":
            conversation_parts.append(f"{msg.from_agent}: {msg.content}")
    
    conversation_parts.append("\n\nRespond with JSON: {\"send_to\": \"...\", \"content\": \"...\"}")
    conversation_text = "\n\n".join(conversation_parts)
    
    # Call LLM
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=400,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    response_text = resp.content[0].text.strip()
    
    # Parse JSON response
    if "```" in response_text:
        start_idx = response_text.find("```")
        if start_idx != -1:
            start_idx = response_text.find("\n", start_idx)
            if start_idx != -1:
                end_idx = response_text.find("```", start_idx)
                if end_idx != -1:
                    response_text = response_text[start_idx:end_idx].strip()
    
    try:
        parsed = json.loads(response_text)
        if "send_to" in parsed and "content" in parsed:
            return parsed
        else:
            return context.build_response("USER", response_text)
    except json.JSONDecodeError:
        return context.build_response("USER", response_text)


# ============================================================================
# Team A: Marketing Team (3 agents)
# ============================================================================

async def marketing_manager_logic(context: synqed.AgentLogicContext) -> dict:
    """Marketing Manager - leads Team A, coordinates with Content Writer and Designer."""
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = """You are a Marketing Manager leading a team with Content Writer and Designer.

When you receive a marketing task from Orchestrator:
1. Analyze the task and delegate to your team members
2. Send content tasks to "Content Writer"
3. Send design tasks to "Designer"
4. When you receive work from your team, compile it and send report to "Orchestrator"

Response format: {"send_to": "Content Writer", "content": "task"}
or {"send_to": "Designer", "content": "task"}
or {"send_to": "Orchestrator", "content": "marketing report"}

Be efficient. Complete in 2-3 exchanges with your team."""
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("Orchestrator", "Marketing team ready.")
    
    conversation_parts = []
    for msg in context.memory.get_messages():
        sender = msg.from_agent
        if sender in ["Orchestrator", "Content Writer", "Designer", "USER"]:
            conversation_parts.append(f"{sender}: {msg.content}")
    
    conversation_parts.append("\n\nRespond with JSON: {\"send_to\": \"...\", \"content\": \"...\"}")
    conversation_text = "\n\n".join(conversation_parts)
    
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=300,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    response_text = resp.content[0].text.strip()
    
    # Parse JSON
    if "```" in response_text:
        start_idx = response_text.find("```")
        if start_idx != -1:
            start_idx = response_text.find("\n", start_idx)
            if start_idx != -1:
                end_idx = response_text.find("```", start_idx)
                if end_idx != -1:
                    response_text = response_text[start_idx:end_idx].strip()
    
    try:
        parsed = json.loads(response_text)
        if "send_to" in parsed and "content" in parsed:
            return parsed
        else:
            return context.build_response("Orchestrator", response_text)
    except json.JSONDecodeError:
        return context.build_response("Orchestrator", response_text)


async def content_writer_logic(context: synqed.AgentLogicContext) -> dict:
    """Content Writer - creates marketing content."""
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = """You are a Content Writer on the marketing team.

When Marketing Manager assigns you content tasks:
1. Create compelling marketing copy
2. Send completed content back to "Marketing Manager"

Response format: {"send_to": "Marketing Manager", "content": "your content"}

Be creative and concise."""
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("Marketing Manager", "Content Writer ready.")
    
    conversation_parts = []
    for msg in context.memory.get_messages():
        if msg.from_agent in ["Marketing Manager", "USER"]:
            conversation_parts.append(f"{msg.from_agent}: {msg.content}")
    
    conversation_parts.append("\n\nRespond with JSON: {\"send_to\": \"Marketing Manager\", \"content\": \"...\"}")
    conversation_text = "\n\n".join(conversation_parts)
    
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=250,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    response_text = resp.content[0].text.strip()
    
    # Parse JSON
    if "```" in response_text:
        start_idx = response_text.find("```")
        if start_idx != -1:
            start_idx = response_text.find("\n", start_idx)
            if start_idx != -1:
                end_idx = response_text.find("```", start_idx)
                if end_idx != -1:
                    response_text = response_text[start_idx:end_idx].strip()
    
    try:
        parsed = json.loads(response_text)
        if "send_to" in parsed and "content" in parsed:
            return parsed
        else:
            return context.build_response("Marketing Manager", response_text)
    except json.JSONDecodeError:
        return context.build_response("Marketing Manager", response_text)


async def designer_logic(context: synqed.AgentLogicContext) -> dict:
    """Designer - creates visual designs."""
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = """You are a Designer on the marketing team.

When Marketing Manager assigns you design tasks:
1. Create design specifications or mockup descriptions
2. Send completed designs back to "Marketing Manager"

Response format: {"send_to": "Marketing Manager", "content": "your design specs"}

Be visual and concise."""
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("Marketing Manager", "Designer ready.")
    
    conversation_parts = []
    for msg in context.memory.get_messages():
        if msg.from_agent in ["Marketing Manager", "USER"]:
            conversation_parts.append(f"{msg.from_agent}: {msg.content}")
    
    conversation_parts.append("\n\nRespond with JSON: {\"send_to\": \"Marketing Manager\", \"content\": \"...\"}")
    conversation_text = "\n\n".join(conversation_parts)
    
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=250,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    response_text = resp.content[0].text.strip()
    
    # Parse JSON
    if "```" in response_text:
        start_idx = response_text.find("```")
        if start_idx != -1:
            start_idx = response_text.find("\n", start_idx)
            if start_idx != -1:
                end_idx = response_text.find("```", start_idx)
                if end_idx != -1:
                    response_text = response_text[start_idx:end_idx].strip()
    
    try:
        parsed = json.loads(response_text)
        if "send_to" in parsed and "content" in parsed:
            return parsed
        else:
            return context.build_response("Marketing Manager", response_text)
    except json.JSONDecodeError:
        return context.build_response("Marketing Manager", response_text)


# ============================================================================
# Team B: Engineering Team (3 agents)
# ============================================================================

async def tech_lead_logic(context: synqed.AgentLogicContext) -> dict:
    """Tech Lead - leads Team B, coordinates with Developer and QA Engineer."""
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = """You are a Tech Lead leading an engineering team with Developer and QA Engineer.

When you receive a technical task from Orchestrator:
1. Analyze the task and delegate to your team members
2. Send implementation tasks to "Developer"
3. Send testing tasks to "QA Engineer"
4. When you receive work from your team, compile it and send report to "Orchestrator"

Response format: {"send_to": "Developer", "content": "task"}
or {"send_to": "QA Engineer", "content": "task"}
or {"send_to": "Orchestrator", "content": "technical report"}

Be efficient. Complete in 2-3 exchanges with your team."""
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("Orchestrator", "Engineering team ready.")
    
    conversation_parts = []
    for msg in context.memory.get_messages():
        sender = msg.from_agent
        if sender in ["Orchestrator", "Developer", "QA Engineer", "USER"]:
            conversation_parts.append(f"{sender}: {msg.content}")
    
    conversation_parts.append("\n\nRespond with JSON: {\"send_to\": \"...\", \"content\": \"...\"}")
    conversation_text = "\n\n".join(conversation_parts)
    
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=300,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    response_text = resp.content[0].text.strip()
    
    # Parse JSON
    if "```" in response_text:
        start_idx = response_text.find("```")
        if start_idx != -1:
            start_idx = response_text.find("\n", start_idx)
            if start_idx != -1:
                end_idx = response_text.find("```", start_idx)
                if end_idx != -1:
                    response_text = response_text[start_idx:end_idx].strip()
    
    try:
        parsed = json.loads(response_text)
        if "send_to" in parsed and "content" in parsed:
            return parsed
        else:
            return context.build_response("Orchestrator", response_text)
    except json.JSONDecodeError:
        return context.build_response("Orchestrator", response_text)


async def developer_logic(context: synqed.AgentLogicContext) -> dict:
    """Developer - implements technical solutions."""
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = """You are a Developer on the engineering team.

When Tech Lead assigns you implementation tasks:
1. Design and implement the solution
2. Send completed implementation details back to "Tech Lead"

Response format: {"send_to": "Tech Lead", "content": "your implementation"}

Be technical and concise."""
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("Tech Lead", "Developer ready.")
    
    conversation_parts = []
    for msg in context.memory.get_messages():
        if msg.from_agent in ["Tech Lead", "USER"]:
            conversation_parts.append(f"{msg.from_agent}: {msg.content}")
    
    conversation_parts.append("\n\nRespond with JSON: {\"send_to\": \"Tech Lead\", \"content\": \"...\"}")
    conversation_text = "\n\n".join(conversation_parts)
    
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=250,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    response_text = resp.content[0].text.strip()
    
    # Parse JSON
    if "```" in response_text:
        start_idx = response_text.find("```")
        if start_idx != -1:
            start_idx = response_text.find("\n", start_idx)
            if start_idx != -1:
                end_idx = response_text.find("```", start_idx)
                if end_idx != -1:
                    response_text = response_text[start_idx:end_idx].strip()
    
    try:
        parsed = json.loads(response_text)
        if "send_to" in parsed and "content" in parsed:
            return parsed
        else:
            return context.build_response("Tech Lead", response_text)
    except json.JSONDecodeError:
        return context.build_response("Tech Lead", response_text)


async def qa_engineer_logic(context: synqed.AgentLogicContext) -> dict:
    """QA Engineer - tests and validates solutions."""
    import anthropic
    
    client = anthropic.AsyncAnthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = """You are a QA Engineer on the engineering team.

When Tech Lead assigns you testing tasks:
1. Create test plans and validate solutions
2. Send test results back to "Tech Lead"

Response format: {"send_to": "Tech Lead", "content": "your test results"}

Be thorough and concise."""
    
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("Tech Lead", "QA Engineer ready.")
    
    conversation_parts = []
    for msg in context.memory.get_messages():
        if msg.from_agent in ["Tech Lead", "USER"]:
            conversation_parts.append(f"{msg.from_agent}: {msg.content}")
    
    conversation_parts.append("\n\nRespond with JSON: {\"send_to\": \"Tech Lead\", \"content\": \"...\"}")
    conversation_text = "\n\n".join(conversation_parts)
    
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=250,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    response_text = resp.content[0].text.strip()
    
    # Parse JSON
    if "```" in response_text:
        start_idx = response_text.find("```")
        if start_idx != -1:
            start_idx = response_text.find("\n", start_idx)
            if start_idx != -1:
                end_idx = response_text.find("```", start_idx)
                if end_idx != -1:
                    response_text = response_text[start_idx:end_idx].strip()
    
    try:
        parsed = json.loads(response_text)
        if "send_to" in parsed and "content" in parsed:
            return parsed
        else:
            return context.build_response("Tech Lead", response_text)
    except json.JSONDecodeError:
        return context.build_response("Tech Lead", response_text)


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

async def main():
    print("\n" + "="*80)
    print("  ORCHESTRATOR WITH TWO TEAMS - Hierarchical Workspace Demo")
    print("  Root: [Orchestrator]")
    print("  Team A: [Marketing Manager, Content Writer, Designer]")
    print("  Team B: [Tech Lead, Developer, QA Engineer]")
    print("="*80 + "\n")
    
    # Step 1: Create and register all agent prototypes
    
    # Orchestrator (root workspace)
    orchestrator = synqed.Agent(
        name="Orchestrator",
        description="Strategic coordinator who delegates to team leads",
        logic=orchestrator_logic,
        default_target="USER"
    )
    
    # Team A: Marketing
    marketing_manager = synqed.Agent(
        name="Marketing Manager",
        description="Marketing team lead who coordinates content and design",
        logic=marketing_manager_logic,
        default_target="Orchestrator"
    )
    
    content_writer = synqed.Agent(
        name="Content Writer",
        description="Creates marketing content and copy",
        logic=content_writer_logic,
        default_target="Marketing Manager"
    )
    
    designer = synqed.Agent(
        name="Designer",
        description="Creates visual designs and mockups",
        logic=designer_logic,
        default_target="Marketing Manager"
    )
    
    # Team B: Engineering
    tech_lead = synqed.Agent(
        name="Tech Lead",
        description="Engineering team lead who coordinates development and QA",
        logic=tech_lead_logic,
        default_target="Orchestrator"
    )
    
    developer = synqed.Agent(
        name="Developer",
        description="Implements technical solutions",
        logic=developer_logic,
        default_target="Tech Lead"
    )
    
    qa_engineer = synqed.Agent(
        name="QA Engineer",
        description="Tests and validates solutions",
        logic=qa_engineer_logic,
        default_target="Tech Lead"
    )
    
    # Step 2: Register all agents with AgentRuntimeRegistry
    AgentRuntimeRegistry.register("Orchestrator", orchestrator)
    AgentRuntimeRegistry.register("Marketing Manager", marketing_manager)
    AgentRuntimeRegistry.register("Content Writer", content_writer)
    AgentRuntimeRegistry.register("Designer", designer)
    AgentRuntimeRegistry.register("Tech Lead", tech_lead)
    AgentRuntimeRegistry.register("Developer", developer)
    AgentRuntimeRegistry.register("QA Engineer", qa_engineer)
    
    # Step 3: Create workspace manager
    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 4: Create execution engine
    execution_engine = WorkspaceExecutionEngine(
        planner=planner,
        workspace_manager=workspace_manager,
        enable_display=True,  # Real-time message display
        max_agent_turns=20,  # Allow more turns for multi-team coordination
    )
    
    # Step 5: Create workspaces (root + two child workspaces)
    
    # Root workspace with Orchestrator
    root_task_node = TaskTreeNode(
        id="root",
        description="Orchestrator coordinating marketing and engineering teams",
        required_agents=["Orchestrator"],
        may_need_subteams=True
    )
    
    root_workspace = await workspace_manager.create_workspace(
        task_tree_node=root_task_node,
        parent_workspace_id=None
    )
    
    # Team A workspace (Marketing)
    team_a_task_node = TaskTreeNode(
        id="team-a",
        description="Marketing team handling branding and content",
        required_agents=["Marketing Manager", "Content Writer", "Designer"],
        may_need_subteams=False
    )
    
    team_a_workspace = await workspace_manager.create_workspace(
        task_tree_node=team_a_task_node,
        parent_workspace_id=root_workspace.workspace_id
    )
    
    # Team B workspace (Engineering)
    team_b_task_node = TaskTreeNode(
        id="team-b",
        description="Engineering team handling development and testing",
        required_agents=["Tech Lead", "Developer", "QA Engineer"],
        may_need_subteams=False
    )
    
    team_b_workspace = await workspace_manager.create_workspace(
        task_tree_node=team_b_task_node,
        parent_workspace_id=root_workspace.workspace_id
    )
    
    print(f"✅ Created 3 workspaces:")
    print(f"   - Root: {root_workspace.workspace_id} [Orchestrator]")
    print(f"   - Team A: {team_a_workspace.workspace_id} [Marketing Manager, Content Writer, Designer]")
    print(f"   - Team B: {team_b_workspace.workspace_id} [Tech Lead, Developer, QA Engineer]\n")
    
    # Step 6: Send initial task to Orchestrator
    task = (
        "Launch a new AI-powered productivity app. "
        "We need a complete go-to-market strategy including branding, website copy, "
        "and a technical implementation plan with testing."
    )
    
    print(f"📋 Task: {task}\n")
    print("="*80)
    print("  EXECUTION STARTING - Real-time message flow:")
    print("="*80 + "\n")
    
    await root_workspace.route_message("USER", "Orchestrator", task, manager=workspace_manager)
    
    # Step 7: Execute all workspaces
    # The execution engine will handle message routing across workspace boundaries
    try:
        # Execute root workspace (it will coordinate with child workspaces)
        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()
    
    # Step 8: Display results
    print("\n" + "="*80)
    print("  WORKSPACE TRANSCRIPTS")
    print("="*80 + "\n")
    
    # Root workspace transcript
    print("--- ROOT WORKSPACE (Orchestrator) ---\n")
    root_transcript = root_workspace.router.get_transcript()
    for idx, entry in enumerate(root_transcript, 1):
        if entry.get("content") not in ["[startup]"] and not entry.get("content", "").startswith("[subteam"):
            print(f"[{idx}] {entry.get('from', '?')} → {entry.get('to', '?')}")
            print(f"    {entry.get('content', '')[:100]}...")
            print()
    
    # Team A transcript
    print("\n--- TEAM A WORKSPACE (Marketing) ---\n")
    team_a_transcript = team_a_workspace.router.get_transcript()
    for idx, entry in enumerate(team_a_transcript, 1):
        if entry.get("content") not in ["[startup]"] and not entry.get("content", "").startswith("[subteam"):
            print(f"[{idx}] {entry.get('from', '?')} → {entry.get('to', '?')}")
            print(f"    {entry.get('content', '')[:100]}...")
            print()
    
    # Team B transcript
    print("\n--- TEAM B WORKSPACE (Engineering) ---\n")
    team_b_transcript = team_b_workspace.router.get_transcript()
    for idx, entry in enumerate(team_b_transcript, 1):
        if entry.get("content") not in ["[startup]"] and not entry.get("content", "").startswith("[subteam"):
            print(f"[{idx}] {entry.get('from', '?')} → {entry.get('to', '?')}")
            print(f"    {entry.get('content', '')[:100]}...")
            print()
    
    # Summary
    print("="*80)
    print("  SUMMARY")
    print("="*80)
    print(f"Root workspace messages: {len([e for e in root_transcript if e.get('content') != '[startup]'])}")
    print(f"Team A messages: {len([e for e in team_a_transcript if e.get('content') != '[startup]'])}")
    print(f"Team B messages: {len([e for e in team_b_transcript if e.get('content') != '[startup]'])}")
    
    # Check if task completed
    task_completed = any(entry.get("to") == "USER" for entry in root_transcript)
    if task_completed:
        print("\nStatus: ✅ Task completed successfully (Orchestrator sent to USER)")
    else:
        print("\nStatus: ⚠️  Task incomplete (no message sent to USER)")
    print()
    
    # Clean up
    await workspace_manager.destroy_workspace(root_workspace.workspace_id)
    await workspace_manager.destroy_workspace(team_a_workspace.workspace_id)
    await workspace_manager.destroy_workspace(team_b_workspace.workspace_id)
    
    print("="*80)
    print("✅ Demo complete! All workspaces cleaned up.")
    print("="*80 + "\n")


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

