"""
Level 3: Two Agents Collaborating

This demonstrates TRUE AI-to-AI collaboration:
- ResearchAgent: Gathers information
- WriterAgent: Creates content from research

The orchestrator automatically coordinates them!

Setup:
1. Make sure you have OPENAI_API_KEY in .env
2. Run: python level3_two_agents.py
"""

import asyncio
import os
from pathlib import Path
from dotenv import load_dotenv
import synqed

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


async def create_agent_executor(system_prompt: str, temperature: float = 0.7):
    """
    Factory function to create agent executors.
    This pattern is used throughout the Qualtrics example.
    """
    async def executor(context):
        user_input = context.get_user_input()
        
        from openai import AsyncOpenAI
        client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))
        
        response = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_input}
            ],
            temperature=temperature
        )
        
        return response.choices[0].message.content
    
    return executor


async def main():
    # Check for API key
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        print("❌ Please set OPENAI_API_KEY in your .env file")
        return
    
    print("=" * 80)
    print("  LEVEL 3: TWO-AGENT COLLABORATION")
    print("  Demonstrating AI-to-AI Autonomous Collaboration")
    print("=" * 80 + "\n")
    
    print("🏗️  Creating specialized agents...\n")
    
    # Agent 1: Research Specialist
    print("🔍 Creating ResearchAgent...")
    researcher = synqed.Agent(
        name="ResearchAgent",
        description="Conducts thorough research on any topic",
        skills=[
            {
                "skill_id": "research",
                "name": "Topic Research",
                "description": "Researches any topic thoroughly, gathering key facts and insights",
                "tags": ["research", "information", "analysis", "facts"]
            }
        ],
        executor=await create_agent_executor(
            system_prompt="""You are a research specialist.

Your role:
- Conduct thorough research on the given topic
- Gather key facts, statistics, and insights
- Organize information clearly
- Provide actionable findings

Format your research with:
1. KEY FACTS
2. IMPORTANT STATISTICS
3. INSIGHTS & TRENDS
4. SUMMARY

Be thorough and factual.""",
            temperature=0.3  # Lower = more focused
        )
    )
    print("   ✓ Research specialist ready\n")
    
    # Agent 2: Content Writer
    print("📝 Creating WriterAgent...")
    writer = synqed.Agent(
        name="WriterAgent",
        description="Creates engaging written content",
        skills=[
            {
                "skill_id": "content_writing",
                "name": "Content Writing",
                "description": "Writes engaging articles, blog posts, and other content",
                "tags": ["writing", "content", "articles", "creative"]
            }
        ],
        executor=await create_agent_executor(
            system_prompt="""You are a professional content writer.

Your role:
- Create engaging, well-structured content
- Write in a clear, accessible style
- Use compelling narratives
- Make information interesting and readable

Format your writing with:
- Engaging introduction
- Well-organized body sections
- Strong conclusion
- Clear takeaways

Be creative but accurate.""",
            temperature=0.7  # Higher = more creative
        )
    )
    print("   ✓ Content writer ready\n")
    
    print("─" * 80)
    print("✅ Both agents created successfully!")
    print("─" * 80 + "\n")
    
    # Start both agent servers
    print("🚀 Starting agent servers...\n")
    
    research_server = synqed.AgentServer(researcher, port=8001)
    writer_server = synqed.AgentServer(writer, port=8002)
    
    await research_server.start_background()
    print(f"   ✓ ResearchAgent: {researcher.url}")
    
    await writer_server.start_background()
    print(f"   ✓ WriterAgent: {writer.url}\n")
    
    print("─" * 80)
    print("✅ Both agents are running and ready!")
    print("─" * 80 + "\n")
    
    # Create orchestrator
    print("🤖 Creating Orchestrator (the AI coordinator)...\n")
    orchestrator = synqed.Orchestrator(
        provider=synqed.LLMProvider.OPENAI,
        api_key=api_key,
        model="gpt-4o",  # Powerful model for coordination
        temperature=0.3
    )
    print("   ✓ Orchestrator ready (using GPT-4o for planning)\n")
    
    # Create collaborative workspace
    print("🏗️  Creating Collaborative Workspace...\n")
    
    # NEW: Built-in colored message display - no custom callback needed!
    workspace = synqed.OrchestratedWorkspace(
        orchestrator=orchestrator,
        enable_agent_discussion=True,  # Enable collaboration
        workspace_persistence=True,    # Track messages
        display_messages=True,         # Built-in colored display (NEW!)
    )
    
    # Register both agents
    workspace.register_agent(researcher)
    workspace.register_agent(writer)
    
    print("   ✓ Workspace created")
    print("   ✓ Both agents registered")
    print("   ✓ Real-time message display enabled\n")
    
    print("─" * 80)
    print("✅ Everything ready for AI-to-AI collaboration!")
    print("─" * 80 + "\n")
    
    # Execute collaborative task
    print("=" * 80)
    print("  COLLABORATIVE TASK: Research + Writing")
    print("=" * 80 + "\n")
    
    task = """
    Research the topic of "quantum computing" and write a brief,
    engaging article about it for a general audience.
    
    The article should explain what quantum computing is, why it matters,
    and what the future might hold. Make it interesting and accessible.
    """
    
    print("📋 Task:")
    print(task.strip())
    print("\n" + "─" * 80 + "\n")
    
    print("🤖 Orchestrator is analyzing the task...")
    print("   └─ Identifying needed capabilities")
    print("   └─ Creating execution plan")
    print("   └─ Selecting best agents for each subtask\n")
    
    print("🔍 AI-to-AI Discovery:")
    print("   └─ Orchestrator querying agent capabilities")
    print("   └─ Learning what each AI can do\n")
    
    print("📤 AI-to-AI Delegation:")
    print("   └─ Orchestrator assigning subtasks to specialist AIs")
    print("   └─ Watch the real-time messages below!\n")
    
    print("┌" + "─" * 78 + "┐")
    print("│" + " LIVE: AI AGENTS WORKING TOGETHER ".center(78) + "│")
    print("└" + "─" * 78 + "┘\n")
    
    # Execute!
    result = await workspace.execute_task(
        task=task,
        workspace_name="Quantum Computing Article"
    )
    
    print("\n" + "=" * 80)
    print("  🎊 COLLABORATION COMPLETE!")
    print("=" * 80 + "\n")
    
    if result.success:
        print("✅ Two AI agents successfully collaborated!\n")
        
        # Show the plan
        print("📋 ORCHESTRATION DETAILS:")
        print("-" * 80)
        print(f"Subtasks created: {len(result.plan.subtasks)}")
        print(f"Agents used: {', '.join(result.plan.selected_agents)}")
        print(f"Execution stages: {len(result.plan.execution_order)}")
        print(f"\nOrchestrator's reasoning:")
        print(result.plan.reasoning)
        print("-" * 80 + "\n")
        
        # Show how work was divided
        print("🎯 WORK DISTRIBUTION:")
        print("-" * 80)
        for i, subtask in enumerate(result.plan.subtasks, 1):
            deps = f" (depends on: {', '.join(subtask.dependencies)})" if subtask.dependencies else ""
            print(f"{i}. {subtask.assigned_agent_name}")
            print(f"   Task: {subtask.description}")
            print(f"   Status: {subtask.status}{deps}")
            print()
        print("-" * 80 + "\n")
        
        # Show the final result
        print("=" * 80)
        print("  📄 FINAL ARTICLE (Synthesized by Orchestrator)")
        print("=" * 80 + "\n")
        print(result.final_result)
        print("\n" + "=" * 80 + "\n")
        
        # Show collaboration stats
        print("📊 COLLABORATION STATISTICS:")
        print("-" * 80)
        print(f"Total messages exchanged: {len(result.workspace_messages)}")
        print(f"Workspace ID: {result.workspace_id}")
        
        # Count messages by agent
        agent_messages = {}
        for msg in result.workspace_messages:
            sender = msg.get('sender_name', 'Unknown')
            agent_messages[sender] = agent_messages.get(sender, 0) + 1
        
        print("\nMessages by agent:")
        for agent, count in sorted(agent_messages.items()):
            print(f"  • {agent}: {count} messages")
        
        print("-" * 80 + "\n")
        
        print("💡 WHAT JUST HAPPENED:")
        print("-" * 80)
        print("1. Orchestrator AI analyzed your task")
        print("2. It broke it into subtasks: [Research] → [Writing]")
        print("3. It selected the best AI for each subtask")
        print("4. ResearchAgent AI gathered information")
        print("5. WriterAgent AI created content using the research")
        print("6. Orchestrator AI synthesized the final result")
        print("7. All autonomously, without human intervention!")
        print("-" * 80 + "\n")
        
    else:
        print(f"❌ Task failed: {result.error}\n")
    
    # Cleanup
    print("🧹 Cleaning up...\n")
    await research_server.stop()
    print("   ✓ Research agent stopped")
    await writer_server.stop()
    print("   ✓ Writer agent stopped\n")
    
    print("─" * 80)
    print("✅ All servers stopped cleanly")
    print("─" * 80 + "\n")
    
    print("=" * 80)
    print("  🎓 LEVEL 3 COMPLETE!")
    print("=" * 80 + "\n")
    
    print("🎉 Congratulations! You just built a multi-agent AI system!")
    print("\n📚 What you learned:")
    print("   ✓ Creating multiple specialized agents")
    print("   ✓ Running agents on different ports")
    print("   ✓ Using an orchestrator for coordination")
    print("   ✓ Building collaborative workspaces")
    print("   ✓ Real-time message monitoring")
    print("   ✓ Managing agent lifecycle (start/stop)")
    print("\n🚀 Next steps:")
    print("   1. Read QUALTRICS_MASTERY_GUIDE.md")
    print("   2. Try adding a 3rd agent (Editor?)")
    print("   3. Tackle Level 4: Build a complete system")
    print("   4. Study the Qualtrics example (5 agents!)")
    print("\n" + "=" * 80 + "\n")


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

