"""
Intelligent Agent Orchestration Example

This example demonstrates how to use Synq's Orchestrator to intelligently route
tasks to the most appropriate agents using LLM-based decision making.

Key Features:
- LLM-powered agent selection
- Support for multiple LLM providers (OpenAI, Anthropic, Google)
- Automatic registration of agents with the orchestrator
- Integration with TaskDelegator for seamless execution

Setup:
1. Install dependencies:
   pip install openai anthropic google-generativeai python-dotenv

2. Set your API key in .env file:
   OPENAI_API_KEY='your-key-here'
   # Or for Anthropic:
   ANTHROPIC_API_KEY='your-key-here'
   # Or for Google:
   GOOGLE_API_KEY='your-key-here'

3. Run: python orchestrator_example.py

The script will demonstrate:
- Creating an orchestrator with your chosen LLM provider
- Registering multiple specialized agents
- Using the orchestrator to intelligently route tasks
- Viewing the orchestrator's reasoning and confidence scores
"""

import asyncio
import os
import logging
from pathlib import Path

from dotenv import load_dotenv

import synqed

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

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# ============================================================================
# AGENT DEFINITIONS
# ============================================================================

def create_recipe_agent() -> synqed.Agent:
    """Create a recipe search agent."""
    
    async def handle_recipe_request(context) -> str:
        user_input = context.get_user_input()
        
        # Simulate recipe search
        return f"""
🍳 Recipe Search Results for: {user_input}

Recipe: Quick Pasta Carbonara
Ingredients:
- 400g spaghetti
- 200g pancetta or bacon
- 4 large eggs
- 100g Parmesan cheese
- Black pepper
- Salt

Instructions:
1. Cook pasta in salted boiling water
2. Fry pancetta until crispy
3. Mix eggs with grated Parmesan
4. Drain pasta, mix with pancetta
5. Remove from heat, stir in egg mixture
6. Season with black pepper

Cooking time: 20 minutes
Difficulty: Easy
"""
    
    return synqed.Agent(
        name="RecipeAgent",
        description="Search and recommend recipes based on ingredients, cuisine, or dietary preferences",
        skills=[
            {
                "skill_id": "recipe_search",
                "name": "Recipe Search",
                "description": "Search for recipes by ingredient, cuisine, dietary restriction, or meal type",
                "tags": ["cooking", "food", "recipes", "ingredients"],
                "examples": [
                    "Find me a pasta recipe",
                    "What can I make with chicken and rice?",
                    "Show me vegetarian dinner ideas"
                ]
            },
            {
                "skill_id": "ingredient_substitution",
                "name": "Ingredient Substitution",
                "description": "Suggest alternatives for recipe ingredients",
                "tags": ["cooking", "substitution", "ingredients"],
                "examples": [
                    "What can I use instead of eggs?",
                    "Substitute for milk in baking"
                ]
            }
        ],
        executor=handle_recipe_request
    )


def create_shopping_agent() -> synqed.Agent:
    """Create a shopping list agent."""
    
    async def handle_shopping_request(context) -> str:
        user_input = context.get_user_input()
        
        # Simulate shopping list creation
        return f"""
🛒 Shopping List for: {user_input}

Pantry Items:
□ Spaghetti (1 pack, 400g)
□ Parmesan cheese (100g)
□ Black pepper

Fresh Items:
□ Eggs (1 dozen)
□ Pancetta or bacon (200g)

Estimated Cost: $15-20
Estimated Shopping Time: 15 minutes

Store Suggestions:
- Whole Foods (0.5 miles)
- Trader Joe's (0.8 miles)
- Local Italian Market (1.2 miles)
"""
    
    return synqed.Agent(
        name="ShoppingAgent",
        description="Create and manage shopping lists, compare prices, and suggest stores",
        skills=[
            {
                "skill_id": "create_shopping_list",
                "name": "Create Shopping List",
                "description": "Generate shopping lists based on recipes or meal plans",
                "tags": ["shopping", "grocery", "list", "planning"],
                "examples": [
                    "Create a shopping list for pasta carbonara",
                    "What groceries do I need for the week?",
                    "Make a list for a dinner party"
                ]
            },
            {
                "skill_id": "price_comparison",
                "name": "Price Comparison",
                "description": "Compare prices across different stores",
                "tags": ["shopping", "price", "comparison", "savings"],
                "examples": [
                    "Where can I buy groceries cheapest?",
                    "Compare prices for organic vegetables"
                ]
            }
        ],
        executor=handle_shopping_request
    )


def create_nutrition_agent() -> synqed.Agent:
    """Create a nutrition analysis agent."""
    
    async def handle_nutrition_request(context) -> str:
        user_input = context.get_user_input()
        
        # Simulate nutrition analysis
        return f"""
📊 Nutrition Analysis for: {user_input}

Per Serving:
- Calories: 650 kcal
- Protein: 28g
- Carbohydrates: 75g
- Fat: 24g
- Fiber: 3g
- Sugar: 2g

Nutritional Highlights:
✓ High in protein from eggs and cheese
✓ Good source of complex carbohydrates
⚠ Moderate in saturated fat
⚠ High in sodium from pancetta

Health Recommendations:
- Add a side salad for more vegetables
- Consider whole wheat pasta for more fiber
- Reduce portion size if watching calories
- Good post-workout meal due to protein content

Dietary Compatibility:
✓ Vegetarian (if using vegetarian bacon)
✗ Vegan
✗ Dairy-free
✗ Gluten-free (unless using GF pasta)
"""
    
    return synqed.Agent(
        name="NutritionAgent",
        description="Analyze nutrition content, suggest healthy alternatives, and provide dietary advice",
        skills=[
            {
                "skill_id": "nutrition_analysis",
                "name": "Nutrition Analysis",
                "description": "Calculate and analyze nutritional content of meals and recipes",
                "tags": ["nutrition", "health", "calories", "diet"],
                "examples": [
                    "How many calories in pasta carbonara?",
                    "Analyze the nutrition of my breakfast",
                    "Is this recipe healthy?"
                ]
            },
            {
                "skill_id": "dietary_advice",
                "name": "Dietary Advice",
                "description": "Provide personalized dietary recommendations",
                "tags": ["nutrition", "diet", "health", "wellness"],
                "examples": [
                    "How can I eat more protein?",
                    "What's a balanced meal plan?",
                    "Foods for weight loss"
                ]
            }
        ],
        executor=handle_nutrition_request
    )


def create_meal_planner_agent() -> synqed.Agent:
    """Create a meal planning agent."""
    
    async def handle_meal_plan_request(context) -> str:
        user_input = context.get_user_input()
        
        # Simulate meal planning
        return f"""
📅 Weekly Meal Plan: {user_input}

MONDAY
Breakfast: Greek yogurt with berries
Lunch: Chicken Caesar salad
Dinner: Pasta carbonara
Snack: Apple with almond butter

TUESDAY
Breakfast: Scrambled eggs with toast
Lunch: Quinoa bowl with vegetables
Dinner: Grilled salmon with asparagus
Snack: Protein smoothie

WEDNESDAY
Breakfast: Oatmeal with banana
Lunch: Turkey sandwich with soup
Dinner: Stir-fry chicken with rice
Snack: Hummus with carrots

THURSDAY
Breakfast: Avocado toast
Lunch: Mediterranean wrap
Dinner: Baked chicken thighs with potatoes
Snack: Trail mix

FRIDAY
Breakfast: Smoothie bowl
Lunch: Tuna salad
Dinner: Homemade pizza night
Snack: Cheese and crackers

Weekly Summary:
- Total recipes: 15 meals
- Prep time: ~6 hours
- Estimated cost: $150-180
- Dietary balance: ✓ Protein ✓ Veggies ✓ Variety

Prep Tips:
1. Batch cook rice and quinoa on Sunday
2. Pre-chop vegetables for the week
3. Marinate proteins in advance
"""
    
    return synqed.Agent(
        name="MealPlannerAgent",
        description="Create comprehensive meal plans for individuals or families",
        skills=[
            {
                "skill_id": "weekly_meal_plan",
                "name": "Weekly Meal Planning",
                "description": "Generate balanced meal plans for the week with variety and nutrition",
                "tags": ["meal-planning", "weekly", "organization", "nutrition"],
                "examples": [
                    "Create a meal plan for the week",
                    "Plan healthy meals for my family",
                    "What should I eat this week?"
                ]
            },
            {
                "skill_id": "batch_cooking",
                "name": "Batch Cooking Guide",
                "description": "Provide strategies for meal prep and batch cooking",
                "tags": ["meal-prep", "batch-cooking", "efficiency"],
                "examples": [
                    "How to meal prep for the week?",
                    "Batch cooking ideas",
                    "Make-ahead meals"
                ]
            }
        ],
        executor=handle_meal_plan_request
    )


# ============================================================================
# ORCHESTRATION EXAMPLES
# ============================================================================

async def example_1_standalone_orchestrator():
    """Example 1: Using the orchestrator standalone (without TaskDelegator)."""
    print("\n" + "="*70)
    print("EXAMPLE 1: Standalone Orchestrator")
    print("="*70 + "\n")
    
    # Get API key and choose provider
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        print("⚠ No OPENAI_API_KEY found. Skipping this example.")
        print("Set OPENAI_API_KEY in your .env file to run this example.")
        return
    
    # Create orchestrator
    orchestrator = synqed.Orchestrator(
        provider=synqed.LLMProvider.OPENAI,
        api_key=api_key,
        model="gpt-4o",
        temperature=0.3,  # Lower temperature for more consistent routing
    )
    
    # Create agents
    recipe_agent = create_recipe_agent()
    shopping_agent = create_shopping_agent()
    nutrition_agent = create_nutrition_agent()
    meal_planner = create_meal_planner_agent()
    
    # Register agents with orchestrator
    orchestrator.register_agent(recipe_agent.card, recipe_agent.url, "RecipeAgent")
    orchestrator.register_agent(shopping_agent.card, shopping_agent.url, "ShoppingAgent")
    orchestrator.register_agent(nutrition_agent.card, nutrition_agent.url, "NutritionAgent")
    orchestrator.register_agent(meal_planner.card, meal_planner.url, "MealPlannerAgent")
    
    print(f"✓ Registered {len(orchestrator.list_agents())} agents\n")
    
    # Test different tasks
    test_tasks = [
        "I want to make pasta carbonara for dinner tonight",
        "Create a shopping list for a week of healthy meals",
        "Is pizza healthy? How many calories does it have?",
        "Plan my meals for the entire week with a focus on protein",
    ]
    
    for i, task in enumerate(test_tasks, 1):
        print(f"\n{'─'*70}")
        print(f"Task {i}: {task}")
        print('─'*70)
        
        try:
            # Get orchestration result
            result = await orchestrator.orchestrate(task, max_agents=2)
            
            print(f"\n🎯 Selected Agent(s):")
            for selection in result.selected_agents:
                print(f"\n  Agent: {selection.agent_name}")
                print(f"  Confidence: {selection.confidence:.1%}")
                print(f"  Skills: {', '.join(selection.recommended_skills)}")
                print(f"  Reasoning: {selection.reasoning}")
            
            print(f"\n📋 Execution Plan:")
            print(f"  {result.execution_plan}")
            
            if result.alternative_agents:
                print(f"\n💡 Alternative Options:")
                for alt in result.alternative_agents:
                    print(f"  - {alt.agent_name} (confidence: {alt.confidence:.1%})")
        
        except Exception as e:
            print(f"❌ Error: {e}")
    
    print("\n" + "="*70 + "\n")


async def example_2_orchestrator_with_delegator():
    """Example 2: Using orchestrator integrated with TaskDelegator."""
    print("\n" + "="*70)
    print("EXAMPLE 2: Orchestrator + TaskDelegator Integration")
    print("="*70 + "\n")
    
    # Get API key
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        print("⚠ No OPENAI_API_KEY found. Skipping this example.")
        return
    
    # Create orchestrator
    orchestrator = synqed.Orchestrator(
        provider=synqed.LLMProvider.OPENAI,
        api_key=api_key,
        model="gpt-4o",
        temperature=0.3,
    )
    
    # Create TaskDelegator with orchestrator
    delegator = synqed.TaskDelegator(orchestrator=orchestrator)
    
    # Create and register agents
    recipe_agent = create_recipe_agent()
    shopping_agent = create_shopping_agent()
    
    # Start agent servers in the background
    recipe_server = synqed.AgentServer(recipe_agent, port=8001)
    shopping_server = synqed.AgentServer(shopping_agent, port=8002)
    
    # Start servers
    asyncio.create_task(recipe_server.start())
    asyncio.create_task(shopping_server.start())
    await asyncio.sleep(2)  # Give servers time to start
    
    # Register with delegator (which also registers with orchestrator)
    delegator.register_agent(recipe_agent)
    delegator.register_agent(shopping_agent)
    
    print(f"✓ Started {len(delegator.list_agents())} agent servers\n")
    
    # Submit tasks - orchestrator will intelligently route them
    test_task = "Find me a recipe for pasta and create a shopping list for it"
    
    print(f"Task: {test_task}\n")
    print("🤖 Orchestrator is analyzing and routing the task...\n")
    
    try:
        # The delegator will use the orchestrator to select the right agent
        result = await delegator.submit_task(test_task)
        
        print("✓ Task completed!\n")
        print("Response:")
        print(result)
    
    except Exception as e:
        print(f"❌ Error: {e}")
    
    finally:
        # Cleanup
        await recipe_server.stop()
        await shopping_server.stop()
        await delegator.close_all()
    
    print("\n" + "="*70 + "\n")


async def example_3_multiple_providers():
    """Example 3: Using different LLM providers."""
    print("\n" + "="*70)
    print("EXAMPLE 3: Multiple LLM Providers")
    print("="*70 + "\n")
    
    # Try different providers
    providers_to_test = []
    
    # OpenAI
    if os.getenv("OPENAI_API_KEY"):
        providers_to_test.append({
            "name": "OpenAI GPT-4o",
            "provider": synqed.LLMProvider.OPENAI,
            "api_key": os.getenv("OPENAI_API_KEY"),
            "model": "gpt-4o"
        })
    
    # Anthropic
    if os.getenv("ANTHROPIC_API_KEY"):
        providers_to_test.append({
            "name": "Anthropic Claude",
            "provider": synqed.LLMProvider.ANTHROPIC,
            "api_key": os.getenv("ANTHROPIC_API_KEY"),
            "model": "claude-3-5-haiku-20241022"
        })
    
    # Google
    if os.getenv("GOOGLE_API_KEY"):
        providers_to_test.append({
            "name": "Google Gemini",
            "provider": synqed.LLMProvider.GOOGLE,
            "api_key": os.getenv("GOOGLE_API_KEY"),
            "model": "gemini-pro"
        })
    
    if not providers_to_test:
        print("⚠ No API keys found. Set one of these in your .env file:")
        print("  - OPENAI_API_KEY")
        print("  - ANTHROPIC_API_KEY")
        print("  - GOOGLE_API_KEY")
        return
    
    # Create agents
    recipe_agent = create_recipe_agent()
    shopping_agent = create_shopping_agent()
    
    task = "I need a quick dinner recipe"
    
    print(f"Task: {task}\n")
    
    # Test each provider
    for config in providers_to_test:
        print(f"\n{'─'*70}")
        print(f"Testing: {config['name']}")
        print('─'*70)
        
        try:
            orchestrator = synqed.Orchestrator(
                provider=config["provider"],
                api_key=config["api_key"],
                model=config["model"],
                temperature=0.3,
            )
            
            orchestrator.register_agent(recipe_agent.card, recipe_agent.url, "RecipeAgent")
            orchestrator.register_agent(shopping_agent.card, shopping_agent.url, "ShoppingAgent")
            
            result = await orchestrator.orchestrate(task)
            
            print(f"✓ Selected: {result.selected_agents[0].agent_name}")
            print(f"  Confidence: {result.selected_agents[0].confidence:.1%}")
            print(f"  Reasoning: {result.selected_agents[0].reasoning[:100]}...")
        
        except Exception as e:
            print(f"❌ Error with {config['name']}: {e}")
    
    print("\n" + "="*70 + "\n")


# ============================================================================
# MAIN
# ============================================================================

async def main():
    """Run all examples."""
    print("\n" + "="*70)
    print("  INTELLIGENT AGENT ORCHESTRATION WITH LLMs")
    print("="*70)
    
    # Run examples
    await example_1_standalone_orchestrator()
    await example_2_orchestrator_with_delegator()
    await example_3_multiple_providers()
    
    print("\n✓ All examples completed!")
    print("\nKey Takeaways:")
    print("1. The Orchestrator uses LLMs to intelligently route tasks")
    print("2. You must specify the LLM provider and API key")
    print("3. The orchestrator provides confidence scores and reasoning")
    print("4. It integrates seamlessly with TaskDelegator")
    print("5. Multiple LLM providers are supported (OpenAI, Anthropic, Google)")


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

