"""
orchestrated workspace example - true multi-agent communication

This demonstrates the inbox-based multi-agent messaging system:
1. Agents run as microservices with inbox endpoints
2. Agents maintain their own internal conversation memory
3. MessageRouter routes messages between agents one at a time
4. No conversation history blobs - only direct message exchange
5. True agent-to-agent communication via structured actions

Architecture:
- Each agent has POST /inbox (receive messages) and GET /inbox (get messages)
- Each agent has POST /respond (process inbox and generate reply)
- Each agent maintains server-side message memory
- MessageRouter routes messages: Agent A → Router → Agent B
- Agents respond with: {"send_to": "AgentName", "content": "text"}
- No conversation history is ever concatenated or passed around

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 asyncio
import os
import json
from pathlib import Path
from dotenv import load_dotenv

# Import the new inbox-based API
import synqed
from synqed import AgentLogicContext

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


# ============================================================================
# Agent Logic Functions
# ============================================================================

async def writer_logic(context: AgentLogicContext) -> dict:
    """
    Writer agent logic.
    
    This function receives:
    - context.memory: Agent's message memory
    - context.latest_message: Latest incoming message
    - context.response: ResponseBuilder helper
    
    Returns:
        dict with "send_to" and "content" keys
    """
    import anthropic
    
    client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = (
        "You are a creative writer collaborating with an editor. "
        "You receive messages one at a time and respond naturally. "
        "When you want to send a message to the Editor, respond with JSON: "
        '{"send_to": "Editor", "content": "your message here"}. '
        "Be brief and creative."
    )
    
    # Get the latest message from memory
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("Editor", "I'm ready to start writing!")
    
    # Build conversation context from agent's own memory
    # Only include messages this agent has received
    conversation_parts = []
    for msg in context.memory.get_messages():
        if msg.from_agent == "USER":
            conversation_parts.append(f"User: {msg.content}")
        elif msg.from_agent == "Editor":
            conversation_parts.append(f"Editor: {msg.content}")
        else:
            conversation_parts.append(f"{msg.from_agent}: {msg.content}")
    
    # Add instruction to respond
    conversation_parts.append("\n\nRespond with JSON: {\"send_to\": \"Editor\", \"content\": \"your response\"}")
    
    conversation_text = "\n\n".join(conversation_parts)
    
    # Call Anthropic API
    resp = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=300,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    response_text = resp.content[0].text.strip()
    
    # Try to parse as JSON, if not valid JSON, wrap it
    try:
        parsed = json.loads(response_text)
        if "send_to" in parsed and "content" in parsed:
            return parsed
        else:
            # Invalid structure, wrap it
            return context.build_response("Editor", response_text)
    except json.JSONDecodeError:
        # Not JSON, wrap it
        return context.build_response("Editor", response_text)


async def editor_logic(context: AgentLogicContext) -> dict:
    """
    Editor agent logic.
    
    This function receives:
    - context.memory: Agent's message memory
    - context.latest_message: Latest incoming message
    - context.response: ResponseBuilder helper
    
    Returns:
        dict with "send_to" and "content" keys
    """
    import anthropic
    
    client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
    
    system_prompt = (
        "You are an editor collaborating with a writer. "
        "You receive messages one at a time and provide feedback. "
        "When you want to send a message to the Writer, respond with JSON: "
        '{"send_to": "Writer", "content": "your message here"}. '
        "Be concise and helpful."
    )
    
    # Get the latest message from memory
    latest_message = context.latest_message
    if not latest_message:
        return context.build_response("Writer", "I'm ready to edit!")
    
    # Build conversation context from agent's own memory
    # Only include messages this agent has received
    conversation_parts = []
    for msg in context.memory.get_messages():
        if msg.from_agent == "USER":
            conversation_parts.append(f"User: {msg.content}")
        elif msg.from_agent == "Writer":
            conversation_parts.append(f"Writer: {msg.content}")
        else:
            conversation_parts.append(f"{msg.from_agent}: {msg.content}")
    
    # Add instruction to respond
    conversation_parts.append("\n\nRespond with JSON: {\"send_to\": \"Writer\", \"content\": \"your response\"}")
    
    conversation_text = "\n\n".join(conversation_parts)
    
    # Call Anthropic API
    resp = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=300,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    response_text = resp.content[0].text.strip()
    
    # Try to parse as JSON, if not valid JSON, wrap it
    try:
        parsed = json.loads(response_text)
        if "send_to" in parsed and "content" in parsed:
            return parsed
        else:
            # Invalid structure, wrap it
            return context.build_response("Writer", response_text)
    except json.JSONDecodeError:
        # Not JSON, wrap it
        return context.build_response("Writer", response_text)


# ============================================================================
# Main Execution
# ============================================================================

async def main():
    print("\n" + "="*80)
    print("  Multi-Agent Collaboration - True Agent-to-Agent Communication")
    print("  Inbox-based routing | Server-side memory | No history blobs")
    print("="*80 + "\n")
    
    # Step 1: Create agents
    print("📝 Creating agents...\n")
    
    writer_agent = synqed.Agent(
        name="Writer",
        description="a creative writer specializing in storytelling and narrative",
        logic=writer_logic,
        default_target="Editor"
    )
    
    editor_agent = synqed.Agent(
        name="Editor",
        description="an editor who reviews and improves written content",
        logic=editor_logic,
        default_target="Writer"
    )
    
    print(f"✓ created agent: {writer_agent.name}")
    print(f"✓ created agent: {editor_agent.name}\n")
    
    # Step 2: Start agent servers with inbox endpoints
    print("🚀 Starting agent servers with inbox endpoints...\n")
    
    writer_server = synqed.AgentServer(writer_agent, port=8001)
    editor_server = synqed.AgentServer(editor_agent, port=8002)
    
    await writer_server.start_background()
    await editor_server.start_background()
    
    # Give servers time to start
    await asyncio.sleep(2.0)
    
    writer_url = "http://localhost:8001"
    editor_url = "http://localhost:8002"
    
    print(f"✓ writer server running on {writer_url}")
    print(f"✓ editor server running on {editor_url}")
    print(f"✓ endpoints: POST /inbox, GET /inbox, POST /respond\n")
    
    # Step 3: Create message router
    print("🔀 Creating message router...\n")
    
    router = synqed.MessageRouter()
    router.register_agent("Writer", writer_url)
    router.register_agent("Editor", editor_url)
    
    print(f"✓ registered agents: {router.list_agents()}\n")
    
    # Step 4: Run turn-based conversation
    print("="*80)
    print("  Conversation structure:")
    print("  turn 1: Writer receives user task, replies to Editor")
    print("  turn 2: Editor replies to Writer")
    print("  turn 3: Writer replies to Editor")
    print("  turn 4: Editor replies to Writer")
    print("  turn 5: Writer replies to Editor")
    print("  turn 6: Editor replies to Writer")
    print("="*80 + "\n")
    
    task = "write a short story about a robot learning to paint."
    
    print(f"📋 task: {task}\n")
    print("="*80)
    print("  starting turn-based conversation...")
    print("="*80 + "\n")
    
    # Execute turn-based conversation with real-time logging
    transcript = []
    
    # Start conversation: place initial task in Writer's inbox
    print(f"\n[Initial] Placing user task in Writer's inbox...")
    await router.send_message("USER", "Writer", task)
    transcript.append({"from": "USER", "to": "Writer", "content": task})
    
    # Track current speaker
    current_agent = "Writer"
    
    # Run turns
    for turn in range(6):
        print(f"\n[Turn {turn + 1}] {current_agent} processing inbox and responding...")
        
        # Get response from current agent
        try:
            response = await router.get_agent_response(current_agent)
            
            # Preview message content (first 100 chars)
            content_preview = response.content[:100] + "..." if len(response.content) > 100 else response.content
            print(f"[Turn {turn + 1}] {current_agent} → {response.send_to}: {content_preview}")
            
            # Add to transcript
            transcript.append({
                "from": current_agent,
                "to": response.send_to,
                "content": response.content
            })
            
            # Route message to target agent's inbox
            await router.send_message(current_agent, response.send_to, response.content)
            
            # Switch to target agent for next turn
            current_agent = response.send_to
            
        except Exception as e:
            print(f"[Turn {turn + 1}] Error: {e}")
            break
    
    # Step 5: Display results
    print("\n" + "="*80)
    print("  conversation complete")
    print("="*80 + "\n")
    
    print("✅ task completed successfully!\n")
    
    # Print final message transcript
    print("📜 full conversation transcript (from actual exchanged messages):")
    print("-" * 80)
    for i, msg in enumerate(transcript, 1):
        print(f"\n[{i}] {msg['from']} → {msg['to']}:")
        print(msg['content'])
    print("-" * 80 + "\n")
    
    # Clean up
    print("🧹 cleaning up...")
    
    await writer_server.stop()
    await editor_server.stop()
    
    print("✓ servers stopped\n")
    
    print("="*80)
    print("  demo complete!")
    print("  architecture:")
    print("  • agents run as microservices with inbox endpoints")
    print("  • each agent maintains server-side message memory")
    print("  • MessageRouter routes messages one at a time")
    print("  • no conversation history blobs")
    print("  • true agent-to-agent communication")
    print("="*80 + "\n")


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