"""
Two-Agent Collaboration Example - Inbox-based messaging

This example demonstrates basic two-agent collaboration using the inbox-based system:
- Two agents (Writer and Editor) collaborate on a task
- Messages are routed via MessageRouter
- Each agent maintains its own memory
- Agents respond with structured JSON messages

Setup:
1. install: pip install synqed anthropic python-dotenv
2. create .env file with: ANTHROPIC_API_KEY='your-key-here'
3. run: python two_agent_collaboration.py
"""
import asyncio
import os
import json
from pathlib import Path
from dotenv import load_dotenv

import synqed
from synqed import AgentLogicContext

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


async def writer_logic(context: AgentLogicContext) -> dict:
    """Writer agent logic - creates content."""
    import anthropic
    
    client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    latest = context.latest_message
    if not latest:
        return context.build_response("Editor", "Hello! I'm ready to write.")
    
    # Build conversation from memory
    conversation = []
    for msg in context.memory.get_messages():
        conversation.append(f"{msg.from_agent}: {msg.content}")
    
    prompt = "\n".join(conversation) + "\n\nRespond as a creative writer. Return JSON: {\"send_to\": \"Editor\", \"content\": \"your message\"}"
    
    resp = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=200,
        system="You are a creative writer collaborating with an editor.",
        messages=[{"role": "user", "content": prompt}],
    )
    
    response_text = resp.content[0].text.strip()
    
    # Parse or wrap response
    try:
        parsed = json.loads(response_text)
        if "send_to" in parsed and "content" in parsed:
            return parsed
    except json.JSONDecodeError:
        pass
    
    return context.build_response("Editor", response_text)


async def editor_logic(context: AgentLogicContext) -> dict:
    """Editor agent logic - reviews and provides feedback."""
    import anthropic
    
    client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    latest = context.latest_message
    if not latest:
        return context.build_response("Writer", "Hello! I'm ready to edit.")
    
    # Build conversation from memory
    conversation = []
    for msg in context.memory.get_messages():
        conversation.append(f"{msg.from_agent}: {msg.content}")
    
    prompt = "\n".join(conversation) + "\n\nRespond as an editor providing feedback. Return JSON: {\"send_to\": \"Writer\", \"content\": \"your message\"}"
    
    resp = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=200,
        system="You are an editor collaborating with a writer.",
        messages=[{"role": "user", "content": prompt}],
    )
    
    response_text = resp.content[0].text.strip()
    
    # Parse or wrap response
    try:
        parsed = json.loads(response_text)
        if "send_to" in parsed and "content" in parsed:
            return parsed
    except json.JSONDecodeError:
        pass
    
    return context.build_response("Writer", response_text)


async def main():
    print("\n" + "="*80)
    print("  Two-Agent Collaboration Example")
    print("  Inbox-based messaging system")
    print("="*80 + "\n")
    
    # Create agents
    print("📝 Creating agents...")
    writer = synqed.Agent(
        name="Writer",
        description="Creative writer",
        logic=writer_logic,
        default_target="Editor"
    )
    
    editor = synqed.Agent(
        name="Editor",
        description="Content editor",
        logic=editor_logic,
        default_target="Writer"
    )
    
    print(f"✓ Created: {writer.name}")
    print(f"✓ Created: {editor.name}\n")
    
    # Start servers
    print("🚀 Starting servers...")
    writer_server = synqed.AgentServer(writer, port=8001)
    editor_server = synqed.AgentServer(editor, port=8002)
    
    await writer_server.start_background()
    await editor_server.start_background()
    await asyncio.sleep(2.0)
    
    print(f"✓ Writer: http://localhost:8001")
    print(f"✓ Editor: http://localhost:8002\n")
    
    # Create router
    print("🔀 Setting up message router...")
    router = synqed.MessageRouter()
    router.register_agent("Writer", "http://localhost:8001")
    router.register_agent("Editor", "http://localhost:8002")
    print(f"✓ Registered: {router.list_agents()}\n")
    
    # Run collaboration
    print("="*80)
    print("  Starting collaboration...")
    print("="*80 + "\n")
    
    task = "Write a short poem about artificial intelligence."
    print(f"📋 Task: {task}\n")
    
    transcript = await router.start_turn_cycle(
        initial_agent="Writer",
        initial_message=task,
        num_turns=4,
        initial_from="USER"
    )
    
    # Display results
    print("\n" + "="*80)
    print("  Collaboration Complete")
    print("="*80 + "\n")
    
    print("📜 Conversation transcript:")
    print("-" * 80)
    for i, msg in enumerate(transcript, 1):
        print(f"\n[{i}] {msg['from']} → {msg['to']}:")
        print(f"    {msg['content'][:200]}...")
    print("-" * 80 + "\n")
    
    # Cleanup
    print("🧹 Cleaning up...")
    await writer_server.stop()
    await editor_server.stop()
    print("✓ Done\n")


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