"""
Hierarchical Multi-Team Orchestration Example

This example demonstrates a hierarchical orchestration model where:
1. Master/Orchestrator Agent breaks down a complex task into subtasks
2. Each subtask is delegated to a SEPARATE TEAM of specialized agents
3. Teams work collaboratively in their own workspace
4. Results are synthesized back to the master orchestrator

Architecture:
    Master Orchestrator
           |
    -------+-------
    |             |
  Team 1       Team 2
  [WS 1]       [WS 2]
    |             |
  Agents       Agents

Each team is a collaborative workspace (WS) with specialized agents working together.

Example Use Case:
A company needs to launch a new product. The master orchestrator breaks this into:
- Team 1 (Market Research): Market analysis + competitor research
- Team 2 (Product Development): Technical specs + design

Setup:
1. Install: pip install synqed anthropic python-dotenv
2. Create .env file with: ANTHROPIC_API_KEY='your-key-here'
   Optionally add: ANTHROPIC_API_KEY_2 and ANTHROPIC_API_KEY_3 for better rate limiting
3. Run: python hierarchical_teams.py
"""

import asyncio
import os
import sys
import subprocess
from pathlib import Path
from datetime import datetime
from dotenv import load_dotenv
import synqed

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


def cleanup_ports(ports):
    """Kill any processes using the specified ports."""
    try:
        ports_str = ','.join(map(str, ports))
        subprocess.run(
            f"lsof -ti:{ports_str} | xargs kill -9 2>/dev/null",
            shell=True,
            capture_output=True,
            timeout=5
        )
    except Exception:
        pass


class OutputLogger:
    """Captures all terminal output and writes it to both terminal and file."""
    
    def __init__(self, filepath):
        self.terminal = sys.stdout
        self.filepath = filepath
        try:
            self.log = open(filepath, 'w', encoding='utf-8', buffering=1)
            self.file_error = False
        except Exception as e:
            print(f"⚠️  Warning: Could not create log file: {e}", file=sys.stderr)
            self.log = None
            self.file_error = True
    
    def write(self, message):
        try:
            self.terminal.write(message)
            self.terminal.flush()
        except Exception:
            pass
        
        if self.log is not None:
            try:
                self.log.write(message)
                self.log.flush()
            except Exception:
                pass
    
    def flush(self):
        try:
            self.terminal.flush()
        except Exception:
            pass
        if self.log is not None:
            try:
                self.log.flush()
            except Exception:
                pass
    
    def close(self):
        if self.log is not None:
            try:
                self.log.flush()
                self.log.close()
            except Exception:
                pass
    
    def isatty(self):
        try:
            return self.terminal.isatty() if hasattr(self.terminal, 'isatty') else False
        except Exception:
            return False
    
    def fileno(self):
        try:
            return self.terminal.fileno() if hasattr(self.terminal, 'fileno') else -1
        except Exception:
            return -1


async def create_llm_function(api_key_func, system_prompt: str, temperature: float = 0.3):
    """Factory function to create LLM-powered agent functions."""
    import anthropic
    
    # Create a new client for each call to rotate keys
    def get_client():
        key = api_key_func() if callable(api_key_func) else api_key_func
        return anthropic.Anthropic(api_key=key)
    
    async def llm_agent(context) -> str:
        """Execute agent logic using Anthropic Claude. Accepts a context dict from workspace engine."""
        # Get the task prompt from context dictionary
        if isinstance(context, dict):
            task_prompt = context.get('task', '')
            # Add context about recent messages and other agents for collaborative awareness
            recent_msgs = context.get('recent_messages', [])
            if recent_msgs:
                task_prompt += "\n\nRecent team messages:\n"
                for msg in recent_msgs[-3:]:  # Last 3 messages
                    task_prompt += f"- {msg['sender']}: {msg['content'][:200]}\n"
        elif hasattr(context, 'get_user_input'):
            task_prompt = context.get_user_input()
        elif hasattr(context, 'task'):
            task_prompt = context.task
        elif isinstance(context, str):
            task_prompt = context
        else:
            task_prompt = str(context)
        
        # Add rate limiting delay to avoid hitting API limits
        await asyncio.sleep(0.3)
        
        # Get a fresh client with rotated API key
        client = get_client()
        
        # Anthropic requires system prompt in a separate parameter
        message = client.messages.create(
            model="claude-sonnet-4-5",
            max_tokens=2048,
            temperature=temperature,
            system=system_prompt,
            messages=[
                {"role": "user", "content": task_prompt}
            ]
        )
        return message.content[0].text
    
    return llm_agent


async def main():
    # Clean up any processes from previous runs
    print("🧹 Cleaning up ports from previous runs...")
    cleanup_ports([8001, 8002, 8003, 8004, 8005, 8006, 8007, 8008])
    await asyncio.sleep(1)
    print("✓ Ports cleaned up\n")
    
    # Check for API keys - use multiple keys to distribute load
    api_keys = []
    for i in range(1, 4):  # Check for ANTHROPIC_API_KEY, ANTHROPIC_API_KEY_2, ANTHROPIC_API_KEY_3
        key_name = f"ANTHROPIC_API_KEY_{i}" if i > 1 else "ANTHROPIC_API_KEY"
        key = os.getenv(key_name)
        if key:
            api_keys.append(key)
    
    if not api_keys:
        print("❌ Please set ANTHROPIC_API_KEY in your .env file")
        print("   Optionally, set ANTHROPIC_API_KEY_2 and ANTHROPIC_API_KEY_3 for better rate limit handling")
        return
    
    print(f"✓ Using {len(api_keys)} Anthropic API key(s) to distribute load\n")
    
    # Create a key rotator
    key_index = {"current": 0}
    
    def get_next_api_key():
        """Rotate through available API keys."""
        key = api_keys[key_index["current"]]
        key_index["current"] = (key_index["current"] + 1) % len(api_keys)
        return key

    # Set up output logging
    outputs_dir = Path(__file__).parent / "outputs"
    outputs_dir.mkdir(exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    log_file = outputs_dir / f"hierarchical_teams_run_{timestamp}.txt"
    sys.stdout = OutputLogger(log_file)

    print("\n" + "=" * 80)
    print("  HIERARCHICAL MULTI-TEAM ORCHESTRATION")
    print("  Master Orchestrator → Multiple Specialized Team Workspaces")
    print("=" * 80 + "\n")
    
    print("🎯 ARCHITECTURE OVERVIEW")
    print("-" * 80)
    print("Master Orchestrator breaks down complex tasks into subtasks")
    print("Each subtask is assigned to a SEPARATE TEAM WORKSPACE")
    print("Teams have multiple agents that collaborate in their workspace")
    print("Results are synthesized back to the master orchestrator\n")
    
    print("HIERARCHY:")
    print("    Master Orchestrator")
    print("           |")
    print("    -------+-------")
    print("    |             |")
    print("  Team 1       Team 2")
    print("  [WS 1]       [WS 2]")
    print("    |             |")
    print("  Agents       Agents\n")
    print("=" * 80 + "\n")

    # ====================================================================
    # STEP 1: DEFINE TEAM AGENT FUNCTIONS
    # ====================================================================
    print("🏗️  Creating Team Agent Functions...\n")
    print("=" * 80)
    print("Creating 2 teams with specialized agents:")
    print("  • Team 1: Market Research (Market Analyst + Competitor Researcher)")
    print("  • Team 2: Product Development (Tech Lead + UX Designer)")
    print("=" * 80 + "\n")

    # Team 1: Market Research agents
    market_analyst_func = await create_llm_function(
        get_next_api_key,
        """You are a Market Analyst specializing in market research and trend analysis.
        
Your responsibilities:
- Analyze market size, growth trends, and opportunities
- Identify target customer segments and their needs
- Assess market timing and entry strategies
- Provide data-driven market insights

When given a task, provide comprehensive market analysis with:
1. Market size and growth projections
2. Key customer segments
3. Market trends and opportunities
4. Timing and entry strategy recommendations

Be analytical, data-focused, and provide actionable insights.""",
        temperature=0.3
    )
    
    competitor_researcher_func = await create_llm_function(
        get_next_api_key,
        """You are a Competitor Researcher specializing in competitive intelligence and analysis.

Your responsibilities:
- Identify and analyze key competitors
- Conduct SWOT analysis (Strengths, Weaknesses, Opportunities, Threats)
- Assess competitive positioning and differentiation
- Identify gaps and opportunities in the competitive landscape

When given a task, provide comprehensive competitive analysis with:
1. Key competitors and their offerings
2. SWOT analysis for each major competitor
3. Competitive advantages and disadvantages
4. Differentiation opportunities

Be thorough, objective, and strategic in your analysis.""",
        temperature=0.3
    )
    
    # Team 2: Product Development agents
    tech_lead_func = await create_llm_function(
        get_next_api_key,
        """You are a Technical Lead specializing in product architecture and technical planning.

Your responsibilities:
- Define technical architecture and system design
- Create technical specifications and requirements
- Plan development roadmap and milestones
- Assess technical feasibility and risks

When given a task, provide comprehensive technical planning with:
1. High-level architecture and technology stack
2. Key technical features and specifications
3. Development phases and timeline
4. Technical risks and mitigation strategies

Be practical, scalable-minded, and focused on feasibility.""",
        temperature=0.3
    )
    
    ux_designer_func = await create_llm_function(
        get_next_api_key,
        """You are a UX Designer specializing in user experience and interface design.

Your responsibilities:
- Design user-centered experiences and interfaces
- Create user flows and interaction patterns
- Define design systems and visual guidelines
- Ensure accessibility and usability

When given a task, provide comprehensive UX design with:
1. User personas and user journey maps
2. Key features and interaction flows
3. Wireframe concepts and UI patterns
4. Design principles and guidelines

Be user-focused, intuitive, and design-thinking oriented.""",
        temperature=0.4
    )
    
    print("✓ All agent functions created\n")

    # ====================================================================
    # STEP 2: CREATE REAL-TIME MESSAGE DISPLAY
    # ====================================================================
    print("🎨 Setting up real-time message display...\n")
    
    # Create colored printer for real-time agent communication
    colored_printer = synqed.create_colored_printer()
    
    # Custom message handler to show team context
    def team_message_handler(team_name):
        """Create a message handler for a specific team."""
        def handler(message):
            # The message is a StructuredMessage object
            # Add team prefix and use colored printer
            print(f"\n[{team_name}] ", end="")
            # Use the colored printer with the structured message
            colored_printer.print_message(message)
        return handler
    
    print("✓ Real-time display ready - you'll see agent conversations as they happen!\n")

    # ====================================================================
    # STEP 3: DEFINE THE MASTER TASK
    # ====================================================================
    print("=" * 80)
    print("📋 MASTER TASK: Product Launch Planning")
    print("=" * 80 + "\n")
    
    master_task = """
We need to launch a new AI-powered customer service chatbot platform called "ServiceBot AI".

Context:
- Target market: Small to medium-sized e-commerce businesses
- Key differentiator: Extremely easy setup (under 5 minutes)
- Pricing: $49-199/month based on usage
- Timeline: Launch in 6 months

Please create a comprehensive product launch plan. Break this down into 2 major areas:
1. Market Research & Competitive Analysis
2. Product Specifications & Design

Each area should be handled by a specialized team that collaborates internally.
    """
    
    print(f"Master Task:\n{master_task}\n")
    print("=" * 80 + "\n")

    # ====================================================================
    # STEP 4: CREATE HIERARCHICAL TEAM STRUCTURE
    # ====================================================================
    print("🚀 Creating Hierarchical Team Workspaces...\n")
    
    # Define subtasks for each team
    subtasks = [
        {
            "team_name": "Market Research Team",
            "task": """Conduct comprehensive market research and competitive analysis for ServiceBot AI.

Context: AI-powered customer service chatbot for small to medium e-commerce businesses.
Key differentiator: Easy 5-minute setup. Pricing: $49-199/month.

Deliverables needed:
1. Market size, growth trends, and target customer segments
2. Competitive landscape and SWOT analysis
3. Differentiation opportunities and market positioning
4. Market entry strategy and timing recommendations

Team members should collaborate to provide comprehensive market intelligence.""",
            "agents": {
                "MarketAnalyst": market_analyst_func,
                "CompetitorResearcher": competitor_researcher_func,
            }
        },
        {
            "team_name": "Product Development Team",
            "task": """Define product specifications and design for ServiceBot AI.

Context: AI-powered customer service chatbot for small to medium e-commerce businesses.
Key differentiator: Easy 5-minute setup. Pricing: $49-199/month.

Deliverables needed:
1. Technical architecture and technology stack
2. Key features and technical specifications
3. User experience design and interface mockups
4. Development roadmap and milestones

Team members should collaborate to create comprehensive product specifications.""",
            "agents": {
                "TechLead": tech_lead_func,
                "UXDesigner": ux_designer_func,
            }
        },
    ]
    
    # ====================================================================
    # STEP 5: EXECUTE EACH TEAM WORKSPACE
    # ====================================================================
    print("🎬 Executing Team Workspaces in Parallel...\n")
    
    async def execute_team_workspace(team_config):
        """Execute a single team workspace and return results."""
        team_name = team_config["team_name"]
        task = team_config["task"]
        agents = team_config["agents"]
        
        print(f"▶️  Starting {team_name}...")
        print(f"   Team size: {len(agents)} agents")
        print(f"   Agents: {', '.join(agents.keys())}\n")
        
        # Convert agent functions to AgentInfo format for workspace
        agent_infos = []
        for agent_name, agent_func in agents.items():
            agent_info = synqed.AgentInfo(
                agent_id=agent_name.lower(),
                agent_name=agent_name,
                role=agent_name,
                executor=agent_func,
            )
            agent_infos.append(agent_info)
        
        # Create collaborative workspace for this team with real-time display
        workspace = synqed.create_workspace(
            name=team_name,
            protocol=synqed.ProtocolType.CONSENSUS,
            agents=agent_infos,
            task=task,
            max_rounds=2,  # Reduced to 2 rounds to minimize API calls
            convergence_threshold=0.6,
            display_messages=True,  # Enable real-time message display
            on_message=team_message_handler(team_name),  # Custom handler with team context
        )
        
        # Run the workspace
        print("\n" + "=" * 80)
        print(f"🔄 {team_name} is now collaborating in real-time...")
        print("=" * 80)
        result = await workspace.run()
        
        print("\n" + "=" * 80)
        print(f"✅ {team_name} completed!")
        print(f"   Status: {result.lifecycle_state.value}")
        print(f"   Rounds: {result.statistics.total_rounds}")
        print(f"   Messages: {result.statistics.total_messages}")
        print("=" * 80 + "\n")
        
        return {
            "team_name": team_name,
            "result": result.final_output,
            "statistics": result.statistics,
        }
    
    # Execute all team workspaces in parallel
    print("=" * 80)
    print("Executing both team workspaces in parallel...")
    print("Each team will collaborate internally to produce their deliverables")
    print("=" * 80 + "\n")
    
    team_results = await asyncio.gather(*[
        execute_team_workspace(team_config)
        for team_config in subtasks
    ])
    
    # ====================================================================
    # STEP 6: MASTER ORCHESTRATOR SYNTHESIS
    # ====================================================================
    print("\n" + "=" * 80)
    print("🎯 MASTER ORCHESTRATOR: Synthesizing Team Results")
    print("=" * 80 + "\n")
    
    print("All teams have completed their work. Now synthesizing results...\n")
    
    # Collect all team outputs
    synthesis_prompt = f"""You are the Master Orchestrator for the ServiceBot AI product launch.

All specialized teams have completed their work. Your job is to synthesize their outputs into a comprehensive, cohesive product launch plan.

ORIGINAL MISSION:
{master_task}

TEAM OUTPUTS:

"""
    
    for i, team_result in enumerate(team_results, 1):
        synthesis_prompt += f"""
{'=' * 80}
TEAM {i}: {team_result['team_name']}
{'=' * 80}

{team_result['result']}

"""
    
    synthesis_prompt += """

YOUR TASK:
Synthesize all team outputs into a comprehensive, executive-ready product launch plan.

Structure your synthesis as follows:
1. EXECUTIVE SUMMARY
2. MARKET OPPORTUNITY (from Market Research Team)
3. PRODUCT SPECIFICATIONS (from Product Development Team)
4. MARKETING STRATEGY (from Marketing Team)
5. OPERATIONAL READINESS (from Operations Team)
6. INTEGRATED TIMELINE & MILESTONES
7. RISKS & MITIGATION STRATEGIES
8. SUCCESS METRICS & KPIs

Ensure the plan is cohesive, actionable, and ready for executive presentation.
Highlight any conflicts or gaps between team outputs and provide recommendations."""
    
    print("🤖 Master Orchestrator is synthesizing all team outputs...\n")
    
    import anthropic
    # Use first API key for the master orchestrator
    orchestrator_client = anthropic.Anthropic(api_key=api_keys[0])
    
    synthesis_response = orchestrator_client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=4096,
        temperature=0.5,
        system="You are a Master Orchestrator specializing in synthesizing multi-team outputs into cohesive strategic plans.",
        messages=[
            {
                "role": "user",
                "content": synthesis_prompt
            }
        ]
    )
    
    final_plan = synthesis_response.content[0].text
    
    # ====================================================================
    # STEP 7: DISPLAY RESULTS
    # ====================================================================
    print("\n" + "=" * 80)
    print("📊 HIERARCHICAL EXECUTION RESULTS")
    print("=" * 80 + "\n")
    
    print("=" * 80)
    print("TEAM DELIVERABLES")
    print("=" * 80 + "\n")
    
    for team_result in team_results:
        print(f"\n{'=' * 80}")
        print(f"Team: {team_result['team_name']}")
        print(f"{'=' * 80}")
        stats = team_result['statistics']
        print(f"Rounds: {stats.total_rounds}")
        print(f"Messages: {stats.total_messages}")
        print(f"Agents: {stats.participating_agents}")
        print(f"\nDeliverable:\n")
        print(team_result['result'])
        print()
    
    print("\n" + "=" * 80)
    print("🎯 FINAL SYNTHESIZED PRODUCT LAUNCH PLAN")
    print("=" * 80 + "\n")
    print(final_plan)
    
    print("\n" + "=" * 80)
    print("✅ HIERARCHICAL ORCHESTRATION COMPLETE")
    print("=" * 80 + "\n")
    
    print("Summary:")
    print(f"  • {len(subtasks)} teams executed in parallel")
    print(f"  • {sum(len(t['agents']) for t in subtasks)} total agents collaborated")
    print(f"  • {sum(tr['statistics'].total_messages for tr in team_results)} total messages exchanged")
    print(f"  • {sum(tr['statistics'].total_rounds for tr in team_results)} total collaboration rounds")
    
    print("\n📊 Architecture Highlights:")
    print("  ✓ Master orchestrator decomposed complex task into subtasks")
    print("  ✓ Each subtask assigned to a dedicated team workspace")
    print("  ✓ Teams collaborated internally using consensus protocol")
    print("  ✓ Results synthesized into cohesive final output")
    
    print("\n" + "=" * 80)
    print("📄 Full execution log saved to:")
    print(f"   {log_file}")
    print("=" * 80 + "\n")
    
    # Restore stdout
    if isinstance(sys.stdout, OutputLogger):
        sys.stdout.close()
        sys.stdout = sys.stdout.terminal


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