"""
orchestrated workspace example - true turn-based multi-agent conversation

this demonstrates the fixed architecture with:
1. no fake subtask pipeline - just real turn-based conversation
2. deterministic turn order - only one agent speaks at a time
3. full conversation history - each agent sees what the other said
4. no broadcast bugs - directed messages only go to selected agent
5. clear logging - see exactly who speaks when

architecture changes:
- removed: subtask_1, subtask_2, etc. pipeline nonsense
- removed: broadcasts that trigger both agents
- added: true turn-based loop (turn 1, turn 2, ..., turn 6)
- added: conversation history with full visibility
- added: final result composed from conversation, not artifacts

the conversation flow:
- turn 1: writer proposes character and concept
- turn 2: editor responds with feedback
- turn 3: writer writes draft based on feedback
- turn 4: editor critiques draft
- turn 5: writer revises based on critique
- turn 6: editor provides final polish

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 - true turn-based conversation")
    print("  no subtask pipeline | no broadcast bugs | real agent interaction")
    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 feedback)\n")
    
    orchestrated_workspace = synqed.OrchestratedWorkspace(
        orchestrator=orchestrator,
        enable_agent_discussion=True,  # enable turn-based conversation
        collaboration_mode=CollaborationMode.COLLABORATIVE,  # supportive mode
        display_messages=True,  # show messages in real-time
    )
    
    print("✓ orchestrated workspace created")
    print("✓ turn-based conversation enabled (6 turns)")
    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("  conversation structure:")
    print("  turn 1: writer → character concept")
    print("  turn 2: editor → feedback on concept")
    print("  turn 3: writer → story draft")
    print("  turn 4: editor → critique and revision")
    print("  turn 5: writer → revised story")
    print("  turn 6: editor → final polish")
    print("="*80 + "\n")
    
    # step 6: give the agents a task
    task = """write a short story about a robot learning to paint."""
    
    print(f"📋 task: {task}\n")
    print("="*80)
    print("  starting turn-based conversation...")
    print("  watch the logs to see deterministic turn order!")
    print("="*80 + "\n")
    
    # execute the task - agents will have 6 turns of conversation
    result = await orchestrated_workspace.execute_task(task)
    
    # step 7: display results
    print("\n" + "="*80)
    print("  conversation complete")
    print("="*80 + "\n")
    
    if result.success:
        print("✅ task completed successfully!\n")
        print("📝 final result (composed from conversation history):")
        print("-" * 80)
        print(f"{result.final_result}")
        print("-" * 80 + "\n")
        
        print("✓ writer spoke on turns 1, 3, 5")
        print("✓ editor spoke on turns 2, 4, 6")
        print("✓ each agent saw the other's previous messages")
        print("✓ no broadcasts triggered both agents")
        print("✓ final result composed from conversation, not subtask artifacts")
    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("  the architecture is now fixed:")
    print("  • no fake subtask pipeline")
    print("  • no broadcast bugs")
    print("  • real turn-based conversation")
    print("  • agents see and respond to each other")
    print("="*80 + "\n")


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

