"""
Advanced Orchestrated Workspace Example

This example demonstrates the full power of OrchestratedWorkspace:
- Automatic task decomposition
- Intelligent subtask assignment
- Parallel execution where possible
- Agent collaboration in temporary environment
- Final result synthesis

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_advanced.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" + "=" * 80)
    print("  ADVANCED ORCHESTRATED WORKSPACE")
    print("  Automatic Task Decomposition & Intelligent Execution")
    print("=" * 80 + "\n")

    # Step 1: Create specialized agents
    print("📦 Creating specialized agents...\n")

    async def research_executor(context):
        """Research agent with deep analysis capabilities."""
        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 research expert. Provide thorough, well-researched information with citations and sources where possible.",
                },
                {"role": "user", "content": user_message},
            ],
            temperature=0.3,  # Lower temperature for factual accuracy
        )
        return response.choices[0].message.content

    async def coding_executor(context):
        """Coding agent with implementation focus."""
        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 senior software engineer. Write clean, production-ready code with best practices, error handling, and documentation.",
                },
                {"role": "user", "content": user_message},
            ],
            temperature=0.2,  # Even lower for code
        )
        return response.choices[0].message.content

    async def writing_executor(context):
        """Writing agent with content creation focus."""
        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 professional technical writer. Create clear, engaging documentation that is accessible to various skill levels.",
                },
                {"role": "user", "content": user_message},
            ],
            temperature=0.7,  # Higher for creative writing
        )
        return response.choices[0].message.content

    async def review_executor(context):
        """Review agent for quality assurance."""
        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 meticulous reviewer. Check for accuracy, completeness, consistency, and quality. Provide constructive feedback.",
                },
                {"role": "user", "content": user_message},
            ],
            temperature=0.4,
        )
        return response.choices[0].message.content

    # Create agents with detailed skills
    research_agent = synqed.Agent(
        name="ResearchAgent",
        description="Expert researcher who gathers and analyzes information",
        skills=[
            {
                "skill_id": "research",
                "name": "Research & Analysis",
                "description": "Conduct thorough research on technical topics",
                "tags": ["research", "analysis", "information-gathering"],
            },
            {
                "skill_id": "fact_checking",
                "name": "Fact Checking",
                "description": "Verify accuracy of information and claims",
                "tags": ["accuracy", "verification", "validation"],
            },
        ],
        executor=research_executor,
    )

    coding_agent = synqed.Agent(
        name="CodingAgent",
        description="Senior software engineer specializing in implementation",
        skills=[
            {
                "skill_id": "code_implementation",
                "name": "Code Implementation",
                "description": "Write production-ready code with best practices",
                "tags": ["coding", "implementation", "software-engineering"],
            },
            {
                "skill_id": "architecture",
                "name": "Software Architecture",
                "description": "Design scalable system architectures",
                "tags": ["architecture", "design", "scalability"],
            },
        ],
        executor=coding_executor,
    )

    writing_agent = synqed.Agent(
        name="WritingAgent",
        description="Technical writer creating clear documentation",
        skills=[
            {
                "skill_id": "documentation",
                "name": "Technical Documentation",
                "description": "Write comprehensive technical documentation",
                "tags": ["documentation", "writing", "technical-writing"],
            },
            {
                "skill_id": "tutorials",
                "name": "Tutorial Creation",
                "description": "Create step-by-step tutorials and guides",
                "tags": ["tutorials", "guides", "education"],
            },
        ],
        executor=writing_executor,
    )

    review_agent = synqed.Agent(
        name="ReviewAgent",
        description="Quality assurance specialist reviewing work",
        skills=[
            {
                "skill_id": "code_review",
                "name": "Code Review",
                "description": "Review code for quality, bugs, and best practices",
                "tags": ["review", "quality-assurance", "code-quality"],
            },
            {
                "skill_id": "content_review",
                "name": "Content Review",
                "description": "Review documentation for clarity and completeness",
                "tags": ["review", "editing", "quality-control"],
            },
        ],
        executor=review_executor,
    )

    print("✓ Created ResearchAgent")
    print("✓ Created CodingAgent")
    print("✓ Created WritingAgent")
    print("✓ Created ReviewAgent\n")

    # Step 2: Start agent servers
    print("🚀 Starting agent servers...\n")

    research_server = synqed.AgentServer(research_agent, port=8001)
    coding_server = synqed.AgentServer(coding_agent, port=8002)
    writing_server = synqed.AgentServer(writing_agent, port=8003)
    review_server = synqed.AgentServer(review_agent, port=8004)

    await research_server.start_background()
    await coding_server.start_background()
    await writing_server.start_background()
    await review_server.start_background()

    print(f"✓ ResearchAgent: {research_agent.url}")
    print(f"✓ CodingAgent: {coding_agent.url}")
    print(f"✓ WritingAgent: {writing_agent.url}")
    print(f"✓ ReviewAgent: {review_agent.url}\n")

    # Step 3: Create orchestrator
    print("🎯 Creating orchestrator...\n")

    orchestrator = synqed.Orchestrator(
        provider=synqed.LLMProvider.OPENAI,
        api_key=api_key,
        model="gpt-4o",
        temperature=0.3,
    )

    print("✓ Orchestrator created\n")

    # Step 4: Create orchestrated workspace
    print("🏗️  Creating orchestrated workspace...\n")

    orchestrated = synqed.OrchestratedWorkspace(
        orchestrator=orchestrator,
        enable_agent_discussion=True,
        workspace_persistence=True,  # Save the workspace for review
    )

    # Register agents
    orchestrated.register_agent(research_agent)
    orchestrated.register_agent(coding_agent)
    orchestrated.register_agent(writing_agent)
    orchestrated.register_agent(review_agent)

    print(f"✓ Orchestrated workspace ready with {len(orchestrated._agents)} agents\n")

    # Step 5: Execute complex task
    print("=" * 80)
    print("COMPLEX TASK EXECUTION")
    print("=" * 80 + "\n")

    complex_task = """
    Create a comprehensive guide on implementing a REST API with Python FastAPI.
    The guide should include:
    - Research on REST API best practices
    - Working code example with authentication
    - Step-by-step documentation
    - Review and quality check of all components
    """

    print(f"Task:\n{complex_task.strip()}\n")
    print("🤖 Orchestrator analyzing task...")
    print("📋 Breaking into subtasks...")
    print("🎯 Assigning to best agents...")
    print("🚀 Executing in temporary workspace...\n")

    # Execute the task - orchestrator handles everything!
    result = await orchestrated.execute_task(
        task=complex_task,
        workspace_name="FastAPI Guide Creation",
    )

    # Step 6: Display results
    print("=" * 80)
    print("EXECUTION RESULTS")
    print("=" * 80 + "\n")

    if result.success:
        print("✅ Task completed successfully!\n")

        # Show the execution plan
        print("📋 EXECUTION PLAN")
        print("-" * 80)
        print(f"Subtasks: {len(result.plan.subtasks)}")
        print(f"Agents used: {len(result.plan.selected_agents)}")
        print(f"\nReasoning:\n{result.plan.reasoning}\n")

        # Show subtasks and assignments
        print("🎯 SUBTASK BREAKDOWN")
        print("-" * 80)
        for i, subtask in enumerate(result.plan.subtasks, 1):
            print(f"\n{i}. {subtask.description}")
            print(f"   Assigned to: {subtask.assigned_agent_name}")
            print(f"   Status: {subtask.status}")
            if subtask.dependencies:
                print(f"   Dependencies: {', '.join(subtask.dependencies)}")

        # Show execution stages
        print(f"\n\n⚡ EXECUTION STAGES")
        print("-" * 80)
        for i, stage in enumerate(result.plan.execution_order, 1):
            print(f"\nStage {i}: {len(stage)} subtask(s) in parallel")
            for subtask_id in stage:
                subtask = next(
                    st for st in result.plan.subtasks if st.subtask_id == subtask_id
                )
                print(f"  • {subtask_id}: {subtask.description[:60]}...")

        # Show individual subtask results
        print(f"\n\n📝 SUBTASK RESULTS")
        print("-" * 80)
        for subtask_id, result_text in result.subtask_results.items():
            subtask = next(
                st for st in result.plan.subtasks if st.subtask_id == subtask_id
            )
            print(f"\n{subtask_id} ({subtask.assigned_agent_name}):")
            print(f"Task: {subtask.description}")
            print(f"Result: {result_text[:200]}...")
            print()

        # Show final synthesized result
        print("=" * 80)
        print("🎉 FINAL RESULT (Synthesized)")
        print("=" * 80 + "\n")
        print(result.final_result)
        print()

        # Show workspace info
        print("=" * 80)
        print("WORKSPACE INFO")
        print("=" * 80 + "\n")
        print(f"Workspace ID: {result.workspace_id}")
        print(f"Subtasks executed: {result.metadata['num_subtasks']}")
        print(f"Agents collaborated: {result.metadata['num_agents']}")
        print("\n✓ Workspace state has been saved for review\n")

    else:
        print("❌ Task execution failed\n")
        print(f"Error: {result.error}\n")

    # Step 7: Cleanup
    print("=" * 80)
    print("CLEANUP")
    print("=" * 80 + "\n")

    print("🛑 Stopping servers...\n")
    await research_server.stop()
    await coding_server.stop()
    await writing_server.stop()
    await review_server.stop()

    print("✓ All servers stopped\n")

    print("=" * 80)
    print("✅ ORCHESTRATED EXECUTION COMPLETE!")
    print("=" * 80 + "\n")

    print("Summary:")
    print(f"  • Task was automatically broken into {len(result.plan.subtasks)} subtasks")
    print(f"  • {len(result.plan.selected_agents)} specialized agents collaborated")
    print(
        f"  • Executed in {len(result.plan.execution_order)} stages (parallel where possible)"
    )
    print("  • Final result synthesized from all subtask outputs")
    print("  • Workspace state saved for review\n")


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

