"""
orchestrated workspace example - two agents having a real multi-turn conversation

this demonstrates:
1. create two ai agents with different specialties
2. set up an orchestrated workspace with real turn-based interaction
3. watch agents have genuine back-and-forth dialogue with controlled turn order!

the orchestrated workspace now features:
- real conversation history that agents can see and respond to
- deterministic turn order controlled by the orchestrator
- agents receive only directed messages when it's their turn
- genuine agent-to-agent interaction (not fake subtask pipeline)
- built-in message display showing clear turn order

key fix implemented:
- the orchestrator announces who goes first
- only that designated agent receives the first turn
- agents alternate in a controlled manner
- each agent sees full conversation history from previous turns

setup:
1. install: pip install synqed anthropic python-dotenv
2. create .env file with: anthropic_api_key='your-key-here'
3. run: python orchestrated_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
    from a2a.utils import get_message_text
    
    client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))
    
    # Build conversation history from the task history
    # This lets the agent see what OTHER agents have said!
    messages = []
    
    if context.current_task and context.current_task.history:
        for msg in context.current_task.history:
            text = get_message_text(msg)
            if msg.role.value == "user":
                messages.append({"role": "user", "content": text})
            elif msg.role.value == "agent":
                # Messages from other agents appear as "assistant" in the conversation
                messages.append({"role": "assistant", "content": text})
    
    # If no history, use the current input
    if not messages:
        user_message = context.get_user_input()
        messages = [{"role": "user", "content": user_message}]
    
    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=messages
    )
    
    return response.content[0].text


async def editor_logic(context):
    """An editor agent that reviews and improves content."""
    import anthropic
    from a2a.utils import get_message_text
    
    client = anthropic.Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))
    
    # Build conversation history from the task history
    # This lets the agent see what OTHER agents have said!
    messages = []
    
    if context.current_task and context.current_task.history:
        for msg in context.current_task.history:
            text = get_message_text(msg)
            if msg.role.value == "user":
                messages.append({"role": "user", "content": text})
            elif msg.role.value == "agent":
                # Messages from other agents appear as "assistant" in the conversation
                messages.append({"role": "assistant", "content": text})
    
    # If no history, use the current input
    if not messages:
        user_message = context.get_user_input()
        messages = [{"role": "user", "content": user_message}]
    
    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=messages
    )
    
    return response.content[0].text


async def main():
    print("\n" + "="*80)
    print("  orchestrated workspace demo - real turn-based agent interaction")
    print("  fixed: deterministic turn order, genuine conversation history")
    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 real turn-based conversation")
    print("  orchestrator controls turn order - watch who speaks first!")
    print("="*80 + "\n")
    
    # Step 6: Give the agents a collaborative task
    task = """Write a story about a robot learning to paint."""
    
    print(f"📋 Task:\n{task}\n")
    print("="*80)
    print("  watching real agent-to-agent dialogue with controlled turns")
    print("  note: logs show actual turn order matching orchestrator's plan")
    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("  turn-based collaboration complete")
    print("="*80 + "\n")
    
    if result.success:
        print("✅ task completed successfully!\n")
        print("📝 final result:")
        print(f"{result.final_result}\n")
        
        if result.subtask_results:
            print("💡 agent deliverables:")
            for subtask_id, subtask_result in result.subtask_results.items():
                print(f"\n{subtask_id}:")
                print(f"  {subtask_result[:200]}..." if len(subtask_result) > 200 else f"  {subtask_result}")
        
        print("\n✓ turn order was respected throughout the dialogue")
        print("✓ agents saw and responded to each other's messages")
        print("✓ conversation history was maintained and used")
    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 - real multi-agent interaction achieved!")
    print("="*80 + "\n")


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

