"""
Orchestrated Workspace Example - Two agents having a natural, multi-turn conversation

This shows how to:
1. Create two AI agents with different specialties
2. Set up an orchestrated workspace with collaborative mode
3. Watch them have a natural, multi-turn conversation in real-time!

The OrchestratedWorkspace enables:
- Multi-phase collaboration (proposal → discussion → refinement → final)
- Agents communicate with brief, 1-2 sentence responses like humans
- Built-in colored message display in real-time
- Agents naturally build on each other's ideas through multiple turns

Setup:
1. Install: pip install synqed anthropic python-dotenv
2. Create .env file with: ANTHROPIC_API_KEY='your-key-here'
3. Run: python workspace.py
"""
import synqed
from synqed import CollaborationMode
import asyncio
import os
from pathlib import Path
from dotenv import load_dotenv

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

# Create agent logic functions
async def writer_logic(context):
    """A creative writer agent."""
    import anthropic
    
    client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))
    user_message = context.get_user_input()
    
    response = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=150,
        system="""You are a creative writer having a natural conversation with an editor colleague. 
Respond conversationally in 1-2 short sentences, like a human would in a casual discussion. 
Be friendly, collaborative, and build on what others say. No need to be formal or verbose.""",
        messages=[{"role": "user", "content": user_message}]
    )
    
    return response.content[0].text


async def editor_logic(context):
    """An editor agent that reviews and improves content."""
    import anthropic
    
    client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))
    user_message = context.get_user_input()
    
    response = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=150,
        system="""You are a professional editor having a natural conversation with a writer colleague.
Respond conversationally in 1-2 short sentences, like a human would in a casual discussion.
Be supportive, give brief feedback, and build on the conversation naturally.""",
        messages=[{"role": "user", "content": user_message}]
    )
    
    return response.content[0].text


async def main():
    print("\n" + "="*80)
    print("  ORCHESTRATED WORKSPACE CONVERSATION DEMO")
    print("  Two agents having a multi-turn, natural conversation")
    print("="*80 + "\n")
    
    # Step 1: Create the agents
    print("📝 Creating agents...\n")
    
    writer = synqed.Agent(
        name="Creative Writer",
        description="A creative writer specializing in storytelling and narrative",
        skills=["creative_writing", "storytelling", "character_development"],
        executor=writer_logic
    )
    
    editor = synqed.Agent(
        name="Professional Editor",
        description="An editor who reviews and improves written content",
        skills=["editing", "proofreading", "content_improvement"],
        executor=editor_logic
    )
    
    print(f"✓ Created agent: {writer.name}")
    print(f"✓ Created agent: {editor.name}\n")
    
    # Step 2: Start agent servers
    print("🚀 Starting agent servers...\n")
    
    writer_server = synqed.AgentServer(writer, port=8001)
    editor_server = synqed.AgentServer(editor, port=8002)
    
    await writer_server.start_background()
    await editor_server.start_background()
    
    # Give servers a moment to fully start
    await asyncio.sleep(0.5)
    
    print(f"✓ Writer server running on port 8001")
    print(f"✓ Editor server running on port 8002\n")
    
    # Step 3: Create orchestrator with Anthropic
    print("🎯 Creating orchestrator...\n")
    
    orchestrator = synqed.Orchestrator(
        provider=synqed.LLMProvider.ANTHROPIC,
        api_key=os.environ.get("ANTHROPIC_API_KEY"),
        model="claude-sonnet-4-5",
        temperature=0.7,
    )
    
    print("✓ Orchestrator created with Claude Sonnet 4-5\n")
    
    # Step 4: Create orchestrated workspace with collaborative mode
    print("🏗️  Creating orchestrated workspace...\n")
    print("Collaboration mode: COLLABORATIVE (supportive, natural conversation)\n")
    
    orchestrated_workspace = synqed.OrchestratedWorkspace(
        orchestrator=orchestrator,
        enable_agent_discussion=True,  # Enable multi-turn conversation
        collaboration_mode=CollaborationMode.COLLABORATIVE,  # Natural, supportive mode
        display_messages=True,  # Show colored messages in real-time
    )
    
    print("✓ Orchestrated workspace created")
    print("✓ Multi-turn conversation enabled")
    print("✓ Real-time message display enabled\n")
    
    # Step 5: Register agents with the orchestrated workspace
    print("👥 Registering agents...\n")
    
    orchestrated_workspace.register_agent(writer)
    orchestrated_workspace.register_agent(editor)
    
    print(f"✓ Registered {writer.name}")
    print(f"✓ Registered {editor.name}\n")
    
    print("="*80)
    print("  AGENTS WILL NOW HAVE A MULTI-TURN CONVERSATION")
    print("  Phases: Proposal → Discussion → Refinement → Final Result")
    print("="*80 + "\n")
    
    # Step 6: Give the agents a collaborative task
    task = """Let's brainstorm a story about a robot learning to paint. 
Create an interesting opening scene together."""
    
    print(f"📋 Task:\n{task}\n")
    print("="*80)
    print("  WATCHING AGENTS COMMUNICATE IN REAL-TIME")
    print("="*80 + "\n")
    
    # Execute the collaborative task - agents will have multiple turns!
    result = await orchestrated_workspace.execute_task(task)
    
    # Step 7: Display results
    print("\n" + "="*80)
    print("  COLLABORATION COMPLETE")
    print("="*80 + "\n")
    
    if result.success:
        print("✅ Task completed successfully!\n")
        print("📝 Final Result:")
        print(f"{result.final_result}\n")
        
        if result.agent_results:
            print("💡 Individual Agent Contributions:")
            for agent_name, agent_result in result.agent_results.items():
                print(f"\n{agent_name}:")
                print(f"  {agent_result[:200]}..." if len(agent_result) > 200 else f"  {agent_result}")
    else:
        print("❌ Task failed")
        print(f"Error: {result.error}\n")
    
    # Clean up
    print("\n🧹 Cleaning up...")
    
    # Stop servers
    await writer_server.stop()
    await editor_server.stop()
    
    print("✓ Servers stopped\n")
    
    print("="*80)
    print("  DEMO COMPLETE")
    print("="*80 + "\n")


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

