"""
Collaborative Agents Demo - True Agent Collaboration in Action

This example demonstrates TRUE agent collaboration where agents:
1. See what other agents are working on
2. Propose their initial approaches
3. Provide feedback to each other
4. Refine their work based on peer input
5. Create a cohesive solution together

This is NOT just task delegation - agents actually interact and collaborate!

Setup:
1. Install: pip install synqed openai python-dotenv
2. Create .env file with: OPENAI_API_KEY='your-key-here'
3. Run: python collaborative_agents_demo.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 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("\n" + "═" * 80)
    print("  🤝 COLLABORATIVE AGENTS DEMO 🤝")
    print("  Watch Agents Truly Collaborate!")
    print("═" * 80 + "\n")

    # Create specialized agent executors
    async def researcher_executor(context):
        """Research specialist."""
        user_message = context.get_user_input()
        from openai import AsyncOpenAI

        client = AsyncOpenAI(api_key=api_key)
        
        system_prompt = """You are a research specialist. You gather information, analyze data, 
        and provide evidence-based insights. You're thorough, detail-oriented, and always cite your reasoning.
        When collaborating, you focus on facts and research findings."""
        
        response = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_message},
            ],
        )
        return response.choices[0].message.content

    async def designer_executor(context):
        """Design specialist."""
        user_message = context.get_user_input()
        from openai import AsyncOpenAI

        client = AsyncOpenAI(api_key=api_key)
        
        system_prompt = """You are a UX/UI design specialist. You think about user experience, 
        aesthetics, accessibility, and usability. You're creative and user-focused.
        When collaborating, you provide design perspectives and user-centric insights."""
        
        response = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_message},
            ],
        )
        return response.choices[0].message.content

    async def developer_executor(context):
        """Development specialist."""
        user_message = context.get_user_input()
        from openai import AsyncOpenAI

        client = AsyncOpenAI(api_key=api_key)
        
        system_prompt = """You are a software development specialist. You think about implementation, 
        architecture, performance, and maintainability. You're practical and solution-oriented.
        When collaborating, you provide technical feasibility insights and implementation suggestions."""
        
        response = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_message},
            ],
        )
        return response.choices[0].message.content

    # Create agents
    print("✨ Creating specialized agents...\n")
    
    researcher = synqed.Agent(
        name="Researcher",
        description="Research specialist who gathers and analyzes information",
        skills=[
            {
                "skill_id": "research",
                "name": "Research & Analysis",
                "description": "Gather information and provide evidence-based insights",
                "tags": ["research", "analysis", "data", "evidence"],
            }
        ],
        executor=researcher_executor,
    )

    designer = synqed.Agent(
        name="Designer",
        description="UX/UI design specialist focused on user experience",
        skills=[
            {
                "skill_id": "design",
                "name": "User Experience Design",
                "description": "Create user-centric designs and interfaces",
                "tags": ["design", "ux", "ui", "user-experience"],
            }
        ],
        executor=designer_executor,
    )

    developer = synqed.Agent(
        name="Developer",
        description="Software development specialist focused on implementation",
        skills=[
            {
                "skill_id": "development",
                "name": "Software Development",
                "description": "Build robust and maintainable solutions",
                "tags": ["development", "coding", "architecture", "implementation"],
            }
        ],
        executor=developer_executor,
    )

    print("  ✓ Researcher - Research & Analysis")
    print("  ✓ Designer - User Experience Design")
    print("  ✓ Developer - Software Development\n")

    # Start servers
    print("🚀 Starting agent servers...\n")
    
    researcher_server = synqed.AgentServer(researcher, port=8001)
    designer_server = synqed.AgentServer(designer, port=8002)
    developer_server = synqed.AgentServer(developer, port=8003)

    await researcher_server.start_background()
    await designer_server.start_background()
    await developer_server.start_background()

    print(f"  ✓ Researcher: {researcher.url}")
    print(f"  ✓ Designer: {designer.url}")
    print(f"  ✓ Developer: {developer.url}\n")

    # Create orchestrator
    print("🎭 Creating orchestrator...\n")
    
    orchestrator = synqed.Orchestrator(
        provider=synqed.LLMProvider.OPENAI,
        api_key=api_key,
        model="gpt-4o",
    )

    # Set up real-time message display FIRST
    print("📡 Setting up real-time collaboration display...\n")
    
    # Track phases for display
    class CollaborationDisplay:
        def __init__(self):
            self.current_phase = None
            self.phase_names = {
                'broadcast': '🤝 PHASE 1: COLLABORATIVE KICKOFF',
                'proposal': '📝 PHASE 2: INITIAL PROPOSALS',
                'feedback': '💬 PHASE 3: PEER FEEDBACK & DISCUSSION',
                'refinement': '✨ PHASE 4: REFINED & INTEGRATED RESULTS',
            }
            self.icons = {
                'Researcher': '🔬',
                'Designer': '🎨',
                'Developer': '💻',
                'Orchestrator': '🎭',
                'System': '⚙️',
            }
        
        def display_message(self, msg):
            """Display a message in real-time."""
            msg_type = msg.get('message_type', 'unknown')
            sender = msg.get('sender_name', 'Unknown')
            content = msg.get('content', '')
            
            # Detect phase changes
            new_phase = None
            if 'COLLABORATIVE WORKSPACE SESSION' in content:
                new_phase = 'broadcast'
            elif 'INITIAL PROPOSAL PHASE' in content:
                new_phase = 'proposal'
            elif 'COLLABORATIVE FEEDBACK PHASE' in content:
                new_phase = 'feedback'
            elif 'REFINEMENT PHASE' in content:
                new_phase = 'refinement'
            
            # Show phase header
            if new_phase and new_phase != self.current_phase:
                if self.current_phase is not None:
                    print("\n" + "─" * 80 + "\n")
                print("╔" + "═" * 78 + "╗")
                print(f"║  {self.phase_names.get(new_phase, 'PHASE').ljust(76)}║")
                print("╚" + "═" * 78 + "╝\n")
                self.current_phase = new_phase
            
            # Skip pure system messages
            if msg_type == 'system' and sender == 'System' and 'joined' not in content.lower():
                return
            
            # Show agent messages
            if msg_type == 'agent':
                icon = self.icons.get(sender, '💬')
                print(f"\n{icon} {sender} is responding...")
                print("─" * 76)
                
                # Show content in real-time
                lines = content.split('\n')
                for line in lines:
                    if line.strip():
                        print(f"  {line}")
                print()
    
    display = CollaborationDisplay()
    
    # Create orchestrated workspace WITH collaboration enabled AND real-time display
    print("🏛️  Creating collaborative workspace...\n")
    
    orchestrated = synqed.OrchestratedWorkspace(
        orchestrator=orchestrator,
        enable_agent_discussion=True,  # Enables true collaboration
        message_callback=display.display_message,  # Real-time message display!
    )

    orchestrated.register_agent(researcher)
    orchestrated.register_agent(designer)
    orchestrated.register_agent(developer)

    print("✓ Workspace ready with collaboration enabled\n")
    print("✓ Real-time message streaming enabled\n")
    
    # The collaborative task
    print("═" * 80)
    print("  🎯 THE COLLABORATIVE TASK 🎯")
    print("═" * 80 + "\n")

    task = """Design a new mobile app feature for habit tracking. The feature should help users 
    build and maintain daily habits. Consider user research, design principles, and technical 
    implementation. Each team member should contribute their expertise and work together to 
    create a cohesive solution."""

    print("📜 TASK:\n")
    print(task)
    print("\n" + "─" * 80 + "\n")
    
    print("🎬 Starting collaborative work...\n")
    print("Watch the agents collaborate IN REAL-TIME:")
    print("  1. 🤝 Kickoff - Understanding the task together")
    print("  2. 📝 Proposals - Each agent shares their approach")
    print("  3. 💬 Feedback - Agents review and comment on each other's work")
    print("  4. ✨ Refinement - Agents incorporate peer insights\n")
    print("─" * 80 + "\n")

    # Execute with collaboration - WATCH IT HAPPEN LIVE!
    result = await orchestrated.execute_task(task)

    # Summary after completion
    if result.success:

        # Show execution summary
        print("\n" + "═" * 80)
        print("  📊 EXECUTION SUMMARY 📊")
        print("═" * 80 + "\n")
        
        print(f"Subtasks Completed: {len(result.subtask_results)}")
        print(f"Collaboration Enabled: Yes ✅")
        print(f"Total Workspace Messages: {len(result.workspace_messages)}")
        
        # Count feedback artifacts
        feedback_artifacts = [a for a in result.metadata.get('artifacts', []) 
                             if isinstance(a, dict) and a.get('artifact_type') == 'feedback']
        print(f"Peer Feedback Exchanges: {len(feedback_artifacts)}")
        
        print("\n" + "─" * 80 + "\n")

        # Show plan details
        print("ORCHESTRATION PLAN:")
        print(f"  Reasoning: {result.plan.reasoning[:200]}...")
        print(f"\n  Subtask Breakdown:")
        for i, subtask in enumerate(result.plan.subtasks, 1):
            print(f"    {i}. {subtask.assigned_agent_name}: {subtask.description}")

        # Show final result
        print("\n" + "═" * 80)
        print("  🎉 FINAL INTEGRATED SOLUTION 🎉")
        print("═" * 80 + "\n")
        
        print(result.final_result)
        print()

        # Highlight the collaboration benefits
        print("\n" + "═" * 80)
        print("  🌟 COLLABORATION HIGHLIGHTS 🌟")
        print("═" * 80 + "\n")
        
        print("✅ Agents shared their initial approaches")
        print("✅ Agents reviewed each other's proposals")
        print("✅ Agents provided constructive feedback")
        print("✅ Agents refined their work based on peer input")
        print("✅ Final solution integrates all perspectives\n")
        
        print("This is TRUE collaboration, not just parallel task execution!")
        print()

    else:
        print(f"❌ Collaboration failed: {result.error}\n")

    # Cleanup
    print("🌙 Closing collaborative workspace...\n")
    
    await researcher_server.stop()
    await designer_server.stop()
    await developer_server.stop()

    print("═" * 80)
    print("  ✨ Collaboration Complete! ✨")
    print("═" * 80 + "\n")


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

