"""
Two Agents Collaborating - Workspace-based multi-agent communication

This demonstrates the new Workspace-based multi-agent messaging system:
1. Agents run in a Workspace (logical routing domain)
2. Agents maintain their own internal conversation memory
3. Workspace routes messages between agents using MessageRouter
4. WorkspaceExecutionEngine executes agents with event-driven scheduling
5. True agent-to-agent communication via structured actions

Architecture:
- Agents are created with logic functions
- Agents are registered with AgentRuntimeRegistry
- WorkspaceManager creates workspaces with agent instances
- WorkspaceExecutionEngine executes agents in workspaces
- Agents respond with: {"send_to": "AgentName", "content": "text"}

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

# Import the new Workspace-based API
import synqed
from synqed.workspace_manager import WorkspaceManager, AgentRuntimeRegistry
from synqed.execution_engine import WorkspaceExecutionEngine
from synqed.planner import PlannerLLM, TaskTreeNode

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

# Configure logging (set to WARNING to reduce noise, display handles output)
logging.basicConfig(
    level=logging.WARNING,
    format='%(asctime)s | %(name)s | %(levelname)s | %(message)s'
)


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

async def writer_logic(context: synqed.AgentLogicContext) -> dict:
    """
    Writer agent logic.
    
    This function receives:
    - context.memory: Agent's message memory
    - context.latest_message: Latest incoming message
    - context.build_response(): Helper to build structured responses
    
    Returns:
        dict with "send_to" and "content" keys
    """
    import anthropic
    
    client = anthropic.AsyncAnthropic(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. "
        "IMPORTANT: After receiving feedback from the Editor, make ONE revision, "
        "then send the final story directly to USER. "
        "When sending to USER, use: "
        '{"send_to": "USER", "content": "your final story here"}. '
        "For drafts to the editor, use: "
        '{"send_to": "Editor", "content": "your draft or question here"}. '
        "Aim to complete in 2-3 exchanges total. Do not endlessly revise."
    )
    
    # 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}")
        elif msg.from_agent == "SYSTEM":
            continue  # Skip system messages
        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 (async)
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=200,
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    response_text = resp.content[0].text.strip()
    
    # Strip markdown code blocks if present (handle various formats)
    if "```" in response_text:
        # Find the JSON content between code blocks
        start_idx = response_text.find("```")
        if start_idx != -1:
            # Skip the opening ``` and optional language identifier
            start_idx = response_text.find("\n", start_idx)
            if start_idx != -1:
                end_idx = response_text.find("```", start_idx)
                if end_idx != -1:
                    response_text = response_text[start_idx:end_idx].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:
        # JSON parsing failed - check if response was INTENDED for USER
        # This handles truncated responses like: {"send_to": "USER", "content": "story text...
        if '"send_to": "USER"' in response_text or '"send_to":"USER"' in response_text:
            # Extract whatever content we can and send to USER
            # Try to find the content field
            content_match = response_text.find('"content":')
            if content_match != -1:
                # Extract content after "content": (may be incomplete)
                content_start = content_match + len('"content":')
                content_text = response_text[content_start:].strip()
                # Remove leading quote if present
                if content_text.startswith('"'):
                    content_text = content_text[1:]
                # Try to clean up trailing incomplete JSON
                content_text = content_text.rstrip('"}')
                return {"send_to": "USER", "content": content_text}
        
        # Not intended for USER, wrap and send to Editor
        return context.build_response("Editor", response_text)


async def editor_logic(context: synqed.AgentLogicContext) -> dict:
    """
    Editor agent logic.
    
    This function receives:
    - context.memory: Agent's message memory
    - context.latest_message: Latest incoming message
    - context.build_response(): Helper to build structured responses
    
    Returns:
        dict with "send_to" and "content" keys
    """
    import anthropic
    
    client = anthropic.AsyncAnthropic(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 focused feedback. "
        "IMPORTANT: Give brief, actionable feedback on the first draft, then tell the writer to finalize and send to USER. "
        "After seeing a revision, approve it and tell the writer to send it to USER immediately. "
        "Do not request multiple rounds of revisions. Keep it to 1-2 exchanges. "
        "Always respond with JSON: "
        '{"send_to": "Writer", "content": "your feedback here"}'
    )
    
    # 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}")
        elif msg.from_agent == "SYSTEM":
            continue  # Skip system messages
        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 (async)
    resp = await client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=500,  # Increased to allow complete stories
        system=system_prompt,
        messages=[{"role": "user", "content": conversation_text}],
    )
    
    response_text = resp.content[0].text.strip()
    
    # Strip markdown code blocks if present (handle various formats)
    if "```" in response_text:
        # Find the JSON content between code blocks
        start_idx = response_text.find("```")
        if start_idx != -1:
            # Skip the opening ``` and optional language identifier
            start_idx = response_text.find("\n", start_idx)
            if start_idx != -1:
                end_idx = response_text.find("```", start_idx)
                if end_idx != -1:
                    response_text = response_text[start_idx:end_idx].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 - Real-Time Message Display")
    print("  Writer and Editor collaborating on a story")
    print("="*80 + "\n")
    
    # Step 1: Create agent prototypes
    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"
    )
    
    # Step 2: Register agents with AgentRuntimeRegistry
    AgentRuntimeRegistry.register("Writer", writer_agent)
    AgentRuntimeRegistry.register("Editor", editor_agent)
    
    # Step 3: Create workspace manager and execution engine
    workspace_manager = WorkspaceManager(workspaces_root=Path("/tmp/synqed_workspaces"))
    
    planner = PlannerLLM(
        provider="anthropic",
        api_key=os.environ["ANTHROPIC_API_KEY"],
        model="claude-sonnet-4-5"
    )
    
    # Create execution engine
    # SIMPLE: Just set max_agent_turns to control how many times agents can respond
    execution_engine = WorkspaceExecutionEngine(
        planner=planner,
        workspace_manager=workspace_manager,
        enable_display=True,  # Real-time message display
        max_agent_turns=10,  # Maximum 10 agent responses (Writer should send to USER before this)
        # Leave other parameters at defaults - they're safety rails
    )
    
    # Step 4: Create workspace
    task_node = TaskTreeNode(
        id="collaboration-task",
        description="Writer and Editor collaborating on a story",
        required_agents=["Writer", "Editor"],
        may_need_subteams=False
    )
    
    workspace = await workspace_manager.create_workspace(
        task_tree_node=task_node,
        parent_workspace_id=None
    )
    
    # Step 5: Send initial task to Writer
    task = "write a short story about a robot learning to paint."
    
    print(f"📋 Task: {task}\n")
    
    await workspace.route_message("USER", "Writer", task, manager=workspace_manager)
    
    # Step 6: Execute the workspace (messages will display in real-time)
    try:
        await execution_engine.run(workspace.workspace_id)
    except Exception as e:
        print(f"\n❌ Error during execution: {e}\n")
        import traceback
        traceback.print_exc()
        raise
    
    # Display complete conversation transcript
    print("\n" + "="*80)
    print("  COMPLETE CONVERSATION TRANSCRIPT")
    print("="*80 + "\n")
    
    transcript = workspace.router.get_transcript()
    message_num = 0
    
    for entry in transcript:
        sender = entry.get("from", "?")
        recipient = entry.get("to", "?")
        content = entry.get("content", "")
        
        # Skip system startup messages
        if content == "[startup]":
            continue
        
        # Skip subteam result messages
        if content.startswith("[subteam_result]"):
            continue
        
        message_num += 1
        
        print(f"[Message {message_num}] {sender} → {recipient}")
        print("-" * 80)
        
        # Try to parse and pretty-print JSON content
        try:
            parsed = json.loads(content)
            if isinstance(parsed, dict) and "send_to" in parsed and "content" in parsed:
                # It's a structured message, show the actual content
                print(parsed["content"])
            else:
                print(content)
        except (json.JSONDecodeError, TypeError):
            # Not JSON, show as-is
            print(content)
        
        print()
    
    # Display summary
    print("="*80)
    print("  SUMMARY")
    print("="*80)
    print(f"Total messages: {message_num}")
    
    # Check if task completed (any message to USER)
    task_completed = any(entry.get("to") == "USER" for entry in transcript)
    
    # Also check for malformed attempts that were intended for USER
    attempted_user_messages = 0
    for entry in transcript:
        content = entry.get("content", "")
        if '"send_to": "USER"' in content or '"send_to":"USER"' in content:
            attempted_user_messages += 1
    
    if task_completed:
        print("Status: ✅ Task completed successfully (message sent to USER)")
    elif attempted_user_messages > 0:
        print(f"Status: ⚠️  Task incomplete ({attempted_user_messages} attempt(s) to send to USER, but JSON was malformed/truncated)")
        print("        Suggestion: Increase max_tokens to allow complete responses")
    else:
        print("Status: ⚠️  Task incomplete (no message sent to USER)")
    print()
    
    # Clean up
    await workspace_manager.destroy_workspace(workspace.workspace_id)
    
    print("="*80)
    print("✅ Demo complete!")
    print("="*80 + "\n")


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