"""
Clean Workspace API Demo - The new way to create workspaces!

This example shows the new unified create_workspace() API.
ONE function with all configuration as parameters.

Features demonstrated:
- Colored agent output (each agent gets a unique color)
- Simple workspace creation
- Protocol configuration
- Custom callbacks
"""

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

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


async def simple_llm_executor(context: dict) -> str:
    """Simple LLM executor for agent responses."""
    client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))
    
    name = context["agent_name"]
    role = context.get("agent_role", "participant")
    task = context["task"]
    recent = context["recent_messages"][-3:]
    
    system_prompt = f"You are {name}, a {role}."
    
    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()


def print_message(message):
    """Print non-system messages (basic version without colors)."""
    sender = message.sender_name
    content = message.content.strip()
    if sender != "System":
        print(f"[{sender}] {content}")


# Create colored printer for better visualization
colored_printer = synqed.create_colored_printer()


async def simple_example():
    """The simplest possible workspace - just 3 required parameters!"""
    print("\n" + "=" * 70)
    print("EXAMPLE 1: Simple Consensus (3 parameters)")
    print("=" * 70 + "\n")
    
    workspace = synqed.create_workspace(
        name="Simple Team",
        task="Decide what to have for lunch",
        agents=[
            synqed.AgentInfo("alice", "Alice", "Chef", executor=simple_llm_executor),
            synqed.AgentInfo("bob", "Bob", "Foodie", executor=simple_llm_executor),
        ],
    )
    
    result = await workspace.run()
    print(f"\n✓ Result: {result[:200]}...\n")


async def consensus_example():
    """Consensus with custom configuration."""
    print("\n" + "=" * 70)
    print("EXAMPLE 2: Consensus with Configuration")
    print("=" * 70 + "\n")
    
    task = "Allocate 100 resources between Food, Defense, Research, and Trade."
    
    agents = [
        synqed.AgentInfo(
            agent_id="farmer",
            agent_name="Farmer",
            role="Food Production Advocate",
            skills=["agriculture"],
            executor=simple_llm_executor
        ),
        synqed.AgentInfo(
            agent_id="warrior",
            agent_name="Warrior",
            role="Defense Strategist",
            skills=["combat"],
            executor=simple_llm_executor
        ),
        synqed.AgentInfo(
            agent_id="scholar",
            agent_name="Scholar",
            role="Research Visionary",
            skills=["science"],
            executor=simple_llm_executor
        ),
    ]
    
    # ONE function call with all configuration!
    workspace = synqed.create_workspace(
        name="Resource Council",
        task=task,
        agents=agents,
        protocol="consensus",  # Can use string or ProtocolType.CONSENSUS
        max_rounds=4,
        convergence_threshold=0.8,
        timeout_seconds=300,
        on_message=colored_printer.print_message,  # Use colored output!
    )
    
    result = await workspace.run()
    print(f"\n✓ Result: {result[:200]}...\n")


async def debate_example():
    """Debate protocol with different settings."""
    print("\n" + "=" * 70)
    print("EXAMPLE 3: Debate Protocol")
    print("=" * 70 + "\n")
    
    workspace = synqed.create_workspace(
        name="Policy Debate",
        task="Should we invest in renewable energy?",
        agents=[
            synqed.AgentInfo("advocate", "Advocate", "Pro-renewable", executor=simple_llm_executor),
            synqed.AgentInfo("critic", "Critic", "Traditional energy", executor=simple_llm_executor),
        ],
        protocol="debate",  # Different protocol!
        max_rounds=3,
        on_message=colored_printer.print_message,
    )
    
    result = await workspace.run()
    print(f"\n✓ Result: {result[:200]}...\n")


async def advanced_example():
    """Advanced configuration with all the bells and whistles."""
    print("\n" + "=" * 70)
    print("EXAMPLE 4: Advanced Configuration")
    print("=" * 70 + "\n")
    
    workspace = synqed.create_workspace(
        name="Research Team",
        task="Analyze quarterly sales data and provide recommendations",
        agents=[
            synqed.AgentInfo("analyst", "DataAnalyst", "Data Expert", executor=simple_llm_executor),
            synqed.AgentInfo("reviewer", "Reviewer", "Quality Control", executor=simple_llm_executor),
        ],
        # Protocol settings
        protocol=synqed.ProtocolType.CRITIQUE_AND_REVISE,
        max_rounds=3,
        max_turns_per_agent_per_round=2,
        # Execution limits
        max_messages=50,
        max_total_tokens=10000,
        timeout_seconds=600,
        # Scheduling
        scheduling_strategy=synqed.SchedulingStrategy.PROTOCOL_DRIVEN,
        # Callbacks
        on_message=colored_printer.print_message,
        on_round_end=lambda round_num, stats: print(f"\n--- Round {round_num} complete ---\n"),
        # Metadata
        description="Collaborative data analysis workspace",
    )
    
    result = await workspace.run()
    print(f"\n✓ Result: {result[:200]}...\n")


async def manual_init_example():
    """Create workspace first, initialize later."""
    print("\n" + "=" * 70)
    print("EXAMPLE 5: Manual Initialization")
    print("=" * 70 + "\n")
    
    # Create empty workspace
    workspace = synqed.create_workspace(
        name="Flexible Workspace",
        protocol="consensus",
        max_rounds=3,
    )
    
    print("✓ Workspace created (not initialized yet)\n")
    
    # Initialize later with task and agents
    await workspace.initialize(
        task="Plan a team event",
        agents=[
            synqed.AgentInfo("planner", "Planner", "Event Organizer", executor=simple_llm_executor),
            synqed.AgentInfo("budgeter", "Budgeter", "Finance", executor=simple_llm_executor),
        ],
    )
    
    print("✓ Workspace initialized\n")
    
    result = await workspace.run()
    print(f"\n✓ Result: {result[:200]}...\n")


async def main():
    """Run all examples."""
    if not os.getenv("OPENAI_API_KEY"):
        print("❌ Set OPENAI_API_KEY in your .env file")
        return
    
    print("\n" + "=" * 70)
    print("  NEW WORKSPACE API DEMO")
    print("  One function. All parameters. Clean and simple.")
    print("=" * 70)
    
    # Run examples
    # await simple_example()
    await consensus_example()
    # await debate_example()
    # await advanced_example()
    # await manual_init_example()
    
    print("\n" + "=" * 70)
    print("  ✓ All examples complete!")
    print("=" * 70 + "\n")


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

