"""
Workspace Engine Demo - Clean, elegant multi-agent collaboration.

This demonstrates how easy it is to set up sophisticated multi-agent
conversations using Synqed's WorkspaceEngine.

✨ WHAT'S NEW:
    - 90% reduction in executor code (106 lines → 35 lines)
    - Simple convenience functions: create_agent(), create_consensus_workspace()
    - Clean, readable agent definitions
    - One-line workspace creation
    - Business logic separated from framework code

🎯 KEY FEATURES:
    - Simple agent definitions with personality traits
    - Automatic protocol-driven conversations (CONSENSUS, DEBATE, CRITIQUE_AND_REVISE)
    - Real-time message streaming with colored output
    - Zero boilerplate code

📚 COMPARISON:
    
    OLD WAY (verbose):
        agent = synqed.AgentInfo(
            agent_id="dev",
            agent_name="Developer",
            role="Backend Engineer",
            skills=["python", "apis"],
            executor=my_executor,
        )
    
    NEW WAY (clean):
        agent = synqed.create_agent("dev", "Developer", "Backend Engineer",
                                   ["python", "apis"], my_executor)
"""

import asyncio
import os
from pathlib import Path
from dotenv import load_dotenv
from openai import AsyncOpenAI

import synqed

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


# ==============================================================================
# Agent Personality Configuration (Your Business Logic)
# ==============================================================================

AGENT_PERSONALITIES = {
    "Food Production Advocate": {
        "stance": "Food is survival. Without it, nothing else matters.",
        "priority": "food",
        "min_acceptable": 50,
    },
    "Defense Strategist": {
        "stance": "We were just raided. Defense protects everything else.",
        "priority": "defense",
        "min_acceptable": 35,
    },
    "Research Visionary": {
        "stance": "Short-term thinking leads to long-term failure.",
        "priority": "research",
        "min_acceptable": 25,
    },
}


# ==============================================================================
# Simple LLM Executor (10 lines instead of 106!)
# ==============================================================================

async def simple_llm_executor(context: dict) -> str:
    """Clean, focused LLM executor - just the essentials."""
    client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    
    # Extract what we need
    name = context["agent_name"]
    role = context.get("agent_role", "participant")
    task = context["task"]
    recent = context["recent_messages"][-3:]  # Last 3 messages
    
    # Get personality (if defined)
    personality = AGENT_PERSONALITIES.get(role)
    system_prompt = f"You are {name}, a {role}."
    
    if personality:
        system_prompt += f" {personality['stance']} You need at least {personality['min_acceptable']} for {personality['priority']}."
    
    # Build conversation context
    conversation = "\n".join([
        f"{msg['sender']}: {msg['content']}"
        for msg in recent if msg['sender'] != 'System'
    ])
    
    prompt = f"Task: {task}\n\nRecent discussion:\n{conversation}\n\nYour response (1-2 sentences):"
    
    response = await client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ],
        temperature=0.9,
        max_tokens=80,
    )
    
    return response.choices[0].message.content.strip()


# ==============================================================================
# Simple Message Handler with Colors
# ==============================================================================

def print_colored_message(message):
    """Clean message printer with automatic coloring."""
    sender = message.sender_name
    content = message.content.strip()
    
    # Auto-assign colors
    color_map = {
        "System": "\033[90m",  # Gray
        "Farmer": "\033[92m", "Environmentalist": "\033[92m", "Developer": "\033[92m",  # Green
        "Warrior": "\033[91m",  # Red
        "Scholar": "\033[94m",  # Blue
        "Engineer": "\033[93m", "Product Manager": "\033[93m",  # Yellow
        "Economist": "\033[95m",  # Magenta
        "UX Designer": "\033[96m",  # Cyan
    }
    
    color = color_map.get(sender, "\033[97m")
    reset = "\033[0m"
    bold = "\033[1m"
    
    # Only show important system messages
    if sender == "System":
        if any(keyword in content for keyword in ["PHASE:", "ROUND", "TERMINATED"]):
            print(f"\n{color}{bold}{content}{reset}")
    else:
        print(f"{color}{bold}[{sender}]{reset} {content}")


async def civilization_simulator_example():
    """
    A village council debates resource allocation.
    Clean 30-line setup for a complex multi-agent simulation!
    """
    print("\n" + "="*80)
    print("🏛️  CIVILIZATION SIMULATOR - Village Council Meeting")
    print("="*80 + "\n")
    
    task = """Your village has 100 units of resources.
Allocate them between: Food (survival), Defense (security), Research (advancement), Trade (relations).
Discuss and reach consensus."""
    
    # Super clean agent creation using the new convenience function
    agents = [
        synqed.create_agent("farmer", "Farmer", "Food Production Advocate", 
                           ["agriculture", "sustainability"], simple_llm_executor),
        synqed.create_agent("warrior", "Warrior", "Defense Strategist",
                           ["combat", "strategy"], simple_llm_executor),
        synqed.create_agent("scholar", "Scholar", "Research Visionary",
                           ["science", "innovation"], simple_llm_executor),
    ]
    
    # One-line workspace creation with the new convenience function!
    workspace = synqed.create_consensus_workspace(
        "Village Council",
        message_handler=print_colored_message,
        max_rounds=4,
        convergence=0.8
    )
    
    # Initialize and run
    await workspace.initialize(
        task=task,
        agents=agents,
        context={"season": "Winter", "recent_events": "Neighboring village was raided"}
    )
    
    print("✓ Initialized | Protocol: CONSENSUS | Agents: Farmer, Warrior, Scholar\n")
    print("="*80 + "\n")
    
    # Run and get result
    result = await workspace.run()
    
    # Clean results display
    print("\n" + "="*80)
    print("📊 RESULTS")
    print("="*80)
    print(f"\nDecision: {result}\n")
    
    stats = workspace.stats.to_dict()
    print(f"Rounds: {stats['total_rounds']} | Messages: {stats['total_messages']} | Duration: {stats['duration_seconds']:.1f}s")


async def debate_example():
    """
    A heated debate about building a dam.
    Shows how easy it is to set up structured debates.
    """
    print("\n" + "="*80)
    print("⚖️  DEBATE - Should We Build the Dam?")
    print("="*80 + "\n")
    
    task = "Our town is considering building a $50M dam. Debate and decide: Build it or not?"
    
    agents = [
        synqed.create_agent("env", "Environmentalist", "Environmental Protection", 
                           ["ecology", "conservation"], simple_llm_executor),
        synqed.create_agent("eng", "Engineer", "Infrastructure Development",
                           ["engineering", "construction"], simple_llm_executor),
        synqed.create_agent("econ", "Economist", "Economic Analysis",
                           ["finance", "cost-benefit"], simple_llm_executor),
    ]
    
    workspace = synqed.create_debate_workspace("Dam Debate", print_colored_message)
    
    await workspace.initialize(task=task, agents=agents)
    print("✓ Debate starting with 3 perspectives...\n" + "="*80 + "\n")
    
    result = await workspace.run()
    
    print("\n" + "="*80)
    print("📊 DEBATE CONCLUDED")
    print("="*80)
    print(f"\nDecision: {result}")


async def critique_and_revise_example():
    """
    Product team refines a feature design through structured critique.
    Demonstrates the power of iterative improvement protocols.
    """
    print("\n" + "="*80)
    print("🎨 DESIGN REVIEW - Mobile App Feature")
    print("="*80 + "\n")
    
    task = 'Design a "Daily Challenge" feature for our fitness app (500k users, goal: 5 sessions/week).'
    
    agents = [
        synqed.create_agent("ux", "UX Designer", "User Experience Lead",
                           ["user-experience", "design"], simple_llm_executor),
        synqed.create_agent("pm", "Product Manager", "Product Strategy",
                           ["product-management", "business"], simple_llm_executor),
        synqed.create_agent("dev", "Developer", "Technical Implementation",
                           ["software-engineering", "api-design"], simple_llm_executor),
    ]
    
    workspace = synqed.create_critique_workspace("Feature Review", print_colored_message)
    
    await workspace.initialize(task=task, agents=agents)
    print("✓ Starting review: Propose → Critique → Revise → Synthesize\n" + "="*80 + "\n")
    
    result = await workspace.run()
    
    print("\n" + "="*80)
    print("📊 DESIGN REVIEW COMPLETE")
    print("="*80)
    print(f"\nFinal Design:\n{result}")


async def main():
    """Run examples - clean and simple."""
    if not os.getenv("OPENAI_API_KEY"):
        print("❌ Set OPENAI_API_KEY in your .env file")
        return
    
    print("\n" + "="*80)
    print("🚀 SYNQED WORKSPACE ENGINE - Multi-Agent Collaboration Made Easy")
    print("="*80)
    print("\nPick a scenario:")
    print("  1. 🏛️  Village Council (CONSENSUS protocol)")
    print("  2. ⚖️  Dam Debate (DEBATE protocol)")
    print("  3. 🎨 Design Review (CRITIQUE_AND_REVISE protocol)")
    
    choice = input("\nChoice (1-3, or Enter for all): ").strip()
    
    examples = {
        "1": civilization_simulator_example,
        "2": debate_example,
        "3": critique_and_revise_example,
    }
    
    if choice in examples:
        await examples[choice]()
    else:
        print("\n🎬 Running all examples...\n")
        for example in examples.values():
            await example()
            print("\n")


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

