"""
Colored Agent Output Demo

This example demonstrates the colored display feature for multi-agent workspaces.
Each agent automatically gets a unique color, making conversations easy to follow.
"""

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()


async def main():
    """Run colored output demo."""
    if not os.getenv("OPENAI_API_KEY"):
        print("❌ Set OPENAI_API_KEY in your .env file")
        return
    
    print("\n" + "=" * 70)
    print("  🎨 COLORED AGENT OUTPUT DEMO")
    print("  Each agent gets a unique color automatically!")
    print("=" * 70 + "\n")
    
    # Create colored printer - this is all you need!
    printer = synqed.create_colored_printer()
    
    # Show the agents with their colors
    print("🤖 Agents participating:\n")
    printer.print_agent_list([
        ("farmer", "Farmer"),
        ("warrior", "Warrior"),
        ("scholar", "Scholar"),
    ])
    print("\n" + "-" * 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
        ),
    ]
    
    print(f"📋 Task: {task}\n")
    print("🎬 Starting discussion...\n")
    print("=" * 70 + "\n")
    
    # Create workspace with colored output
    workspace = synqed.create_workspace(
        name="Resource Council",
        task=task,
        agents=agents,
        protocol="consensus",
        max_rounds=3,
        convergence_threshold=0.8,
        timeout_seconds=300,
        on_message=printer.print_message,  # 🎨 Colored output enabled!
    )
    
    result = await workspace.run()
    
    print("\n" + "=" * 70)
    print("✓ Discussion Complete!")
    print("=" * 70 + "\n")
    
    print(f"Result: {result[:300]}...\n")
    
    # Show statistics
    stats = workspace.stats.to_dict()
    print("\n📊 Statistics:")
    print(f"  • Total rounds: {stats['total_rounds']}")
    print(f"  • Total messages: {stats['total_messages']}")
    print(f"  • Messages by agent:")
    for agent_id, count in stats['messages_by_agent'].items():
        if agent_id != "system":
            agent_name = next((a.agent_name for a in agents if a.agent_id == agent_id), agent_id)
            formatted_name = printer.format_agent_name(agent_id, agent_name)
            print(f"    - {formatted_name}: {count} messages")
    print()


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

