"""
Nexus Demo — Multi-turn Conversation + Streaming Output

Demonstrates three conversation modes:
  1. chat()        — multi-turn, non-streaming
  2. stream()      — single-turn, streaming
  3. chat_stream() — multi-turn + streaming (recommended)

Usage:
    python examples/demo.py                  # run all demos
    python examples/demo.py chat_stream      # run one demo
"""

import asyncio
import sys

from _config import get_embedding_config, get_model_config

from nexus import AgentCapability, Team
from nexus.orchestration import TopologyMode


async def demo_chat():
    """Multi-turn conversation (non-streaming)."""
    print("=" * 60)
    print("Demo 1: Multi-turn Conversation (non-streaming)")
    print("=" * 60)

    team = Team(
        **get_model_config(),
        mode=TopologyMode.ADAPTIVE,
        **get_embedding_config(),
        verbose=False,
    )
    team.add_agent(
        "coder",
        instructions="You are a programmer.",
        capabilities=AgentCapability(coding=1.0, reasoning=0.8),
    )
    team.add_agent(
        "writer",
        instructions="You are a writer.",
        capabilities=AgentCapability(writing=1.0, creativity=0.9),
    )

    # Turn 1
    print("\nUser: Write a short poem about spring")
    answer1 = await team.chat("Write a short poem about spring")
    print(f"Assistant: {answer1}")

    # Turn 2 — references previous context
    print("\nUser: Translate this poem to Chinese")
    answer2 = await team.chat("Translate this poem to Chinese")
    print(f"Assistant: {answer2}")

    # Turn 3
    print("\nUser: Add a title to this poem")
    answer3 = await team.chat("Add a title to this poem")
    print(f"Assistant: {answer3}")

    await team.close()


async def demo_stream():
    """Single-turn streaming output."""
    print("\n" + "=" * 60)
    print("Demo 2: Streaming Output (single-turn)")
    print("=" * 60)

    team = Team(
        **get_model_config(),
        mode=TopologyMode.ADAPTIVE,
        **get_embedding_config(),
        verbose=False,
    )
    team.add_agent(
        "coder",
        instructions="You are a programmer.",
        capabilities=AgentCapability(coding=1.0, reasoning=0.8),
    )
    team.add_agent(
        "writer",
        instructions="You are a writer.",
        capabilities=AgentCapability(writing=1.0, creativity=0.9),
    )

    print("\nUser: Write a quicksort in Python")
    print("Assistant: ", end="", flush=True)
    async for chunk in team.stream("Write a quicksort in Python"):
        print(chunk, end="", flush=True)
    print()

    await team.close()


async def demo_chat_stream():
    """Multi-turn conversation with streaming."""
    print("\n" + "=" * 60)
    print("Demo 3: Multi-turn + Streaming (chat_stream)")
    print("=" * 60)

    team = Team(
        **get_model_config(),
        mode=TopologyMode.ADAPTIVE,
        **get_embedding_config(),
        verbose=False,
    )
    team.add_agent(
        "coder",
        instructions="You are a programmer. Keep replies concise.",
        capabilities=AgentCapability(coding=1.0, reasoning=0.8),
    )
    team.add_agent(
        "writer",
        instructions="You are a writer. Keep replies concise.",
        capabilities=AgentCapability(writing=1.0, creativity=0.9),
    )

    prompts = [
        "Write a bubble sort function in Python",
        "Add type annotations and a docstring to this function",
        "Now write a unit test for it",
    ]
    for prompt in prompts:
        print(f"\nUser: {prompt}")
        print("Assistant: ", end="", flush=True)
        async for chunk in team.chat_stream(prompt):
            print(chunk, end="", flush=True)
        print()

    team.reset_conversation()
    print("\n--- Conversation reset! ---")

    print("\nUser: Hello, what did we talk about before?")
    print("Assistant: ", end="", flush=True)
    async for chunk in team.chat_stream("Hello, what did we talk about before?"):
        print(chunk, end="", flush=True)
    print()

    await team.close()


async def main():
    demos = {
        "chat": demo_chat,
        "stream": demo_stream,
        "chat_stream": demo_chat_stream,
    }
    choice = sys.argv[1] if len(sys.argv) > 1 else "all"
    if choice == "all":
        for fn in demos.values():
            await fn()
    elif choice in demos:
        await demos[choice]()
    else:
        print(f"Usage: python demo.py [{'|'.join(demos.keys())}|all]")


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