"""
Simple Orchestrated Workspace Example

This example shows the simplest way to use OrchestratedWorkspace
for automatic task decomposition and execution.

Setup:
1. Install: pip install synqed openai python-dotenv
2. Create .env file with: OPENAI_API_KEY='your-key-here'
3. Run: python orchestrated_workspace_simple.py
"""

import asyncio
import os
from pathlib import Path

from dotenv import load_dotenv

import synqed

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


async def main():
    # Check for API key
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        print("❌ Please set OPENAI_API_KEY in your .env file")
        return

    print("\n" + "=" * 70)
    print("  SIMPLE ORCHESTRATED WORKSPACE")
    print("=" * 70 + "\n")

    # Step 1: Create simple agents
    async def generic_executor(context):
        """Generic agent executor using GPT-4."""
        user_message = context.get_user_input()
        from openai import AsyncOpenAI

        client = AsyncOpenAI(api_key=api_key)
        response = await client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": "You are a helpful AI assistant."},
                {"role": "user", "content": user_message},
            ],
        )
        return response.choices[0].message.content

    # Create agents with different specializations
    research_agent = synqed.Agent(
        name="ResearchAgent",
        description="Gathers and analyzes information",
        skills=["research", "analysis", "fact-checking"],
        executor=generic_executor,
    )

    writer_agent = synqed.Agent(
        name="WriterAgent",
        description="Creates written content",
        skills=["writing", "editing", "content-creation"],
        executor=generic_executor,
    )

    print("✓ Created agents\n")

    # Step 2: Start servers
    print("Starting servers...")
    research_server = synqed.AgentServer(research_agent, port=8001)
    writer_server = synqed.AgentServer(writer_agent, port=8002)

    await research_server.start_background()
    await writer_server.start_background()
    print("✓ Servers running\n")

    # Step 3: Create orchestrator
    orchestrator = synqed.Orchestrator(
        provider=synqed.LLMProvider.OPENAI,
        api_key=api_key,
        model="gpt-4o",
    )

    # Step 4: Create orchestrated workspace
    orchestrated = synqed.OrchestratedWorkspace(
        orchestrator=orchestrator,
        enable_agent_discussion=True,
    )

    # Register agents
    orchestrated.register_agent(research_agent)
    orchestrated.register_agent(writer_agent)

    print("✓ Orchestrated workspace ready\n")

    # Step 5: Execute task - the orchestrator handles everything!
    print("=" * 70)
    print("TASK EXECUTION")
    print("=" * 70 + "\n")

    task = "Research the benefits of meditation and write a short article about it"

    print(f"Task: {task}\n")
    print("🤖 Working...\n")

    # This single call:
    # 1. Breaks the task into subtasks
    # 2. Assigns each subtask to the best agent
    # 3. Creates a temporary workspace
    # 4. Executes subtasks (in parallel where possible)
    # 5. Synthesizes the final result
    result = await orchestrated.execute_task(task)

    # Step 6: Show detailed plan
    if result.success:
        print("✅ Success!\n")

        # Show the complete execution plan
        print("=" * 70)
        print("📋 EXECUTION PLAN CREATED BY ORCHESTRATOR")
        print("=" * 70 + "\n")
        
        print(f"Original Task: {result.plan.original_task}\n")
        print(f"Number of Subtasks: {len(result.plan.subtasks)}")
        print(f"Agents Selected: {', '.join(result.plan.selected_agents)}\n")
        
        print("Orchestrator's Reasoning:")
        print(result.plan.reasoning)
        print("\n" + "-" * 70 + "\n")
        
        # Show each subtask in detail
        print("SUBTASK BREAKDOWN:\n")
        for i, subtask in enumerate(result.plan.subtasks, 1):
            print(f"Subtask {i}:")
            print(f"  ID: {subtask.subtask_id}")
            print(f"  Description: {subtask.description}")
            print(f"  Assigned to: {subtask.assigned_agent_name}")
            print(f"  Dependencies: {subtask.dependencies if subtask.dependencies else 'None'}")
            print(f"  Order: {subtask.order}")
            print(f"  Status: {subtask.status}")
            print()
        
        # Show execution stages
        print("-" * 70 + "\n")
        print("EXECUTION STAGES (Parallel where possible):\n")
        for i, stage in enumerate(result.plan.execution_order, 1):
            print(f"Stage {i}: {stage}")
            print(f"  ({len(stage)} subtask(s) executed in parallel)")
            print()
        
        # Show workspace interactions
        print("=" * 70)
        print("💬 COMPLETE WORKSPACE MESSAGE HISTORY")
        print("=" * 70 + "\n")
        
        print(f"Total messages in workspace: {len(result.workspace_messages)}\n")
        
        for i, msg in enumerate(result.workspace_messages, 1):
            msg_type = msg.get('message_type', 'unknown')
            sender = msg.get('sender_name', 'Unknown')
            content = msg.get('content', '')
            timestamp = msg.get('timestamp', '')
            
            # Color code by message type
            icon = {
                'system': '⚙️',
                'agent': '🤖',
                'user': '👤',
                'status': '📊',
                'artifact': '📦'
            }.get(msg_type, '💬')
            
            print(f"[{i}] {icon} {msg_type.upper()} - {sender}")
            print(f"    Time: {timestamp}")
            print(f"    Message: {content[:200]}{'...' if len(content) > 200 else ''}")
            print()
        
        print("-" * 70 + "\n")
        
        print("Individual Subtask Results:\n")
        for subtask in result.plan.subtasks:
            subtask_result = result.subtask_results.get(subtask.subtask_id, "No result")
            print(f"┌─ {subtask.subtask_id} ({subtask.assigned_agent_name})")
            print(f"│  Task: {subtask.description}")
            print(f"│")
            print(f"│  Result:")
            # Format the result with indentation
            for line in subtask_result.split('\n'):
                print(f"│  {line}")
            print(f"└─ End of {subtask.subtask_id}\n")

        # Show final synthesized result
        print("=" * 70)
        print("🎉 FINAL SYNTHESIZED RESULT")
        print("=" * 70 + "\n")
        print(result.final_result)
        print()
    else:
        print(f"❌ Failed: {result.error}\n")

    # Cleanup
    await research_server.stop()
    await writer_server.stop()

    print("=" * 70)
    print("✓ Done!")
    print("=" * 70 + "\n")


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

