"""
Nexus Demo — Three Orchestration Modes

  1. GRAPH mode    — static graph execution, like LangGraph
  2. ROUTER mode   — LLM-based dynamic routing
  3. ADAPTIVE mode — embedding-based capability matching

Usage:
    python examples/demo_modes.py              # run all
    python examples/demo_modes.py graph        # run one
"""

import asyncio
import sys

from pydantic import BaseModel

from _config import get_embedding_config, get_model_config

from nexus import AgentCapability, SharedState, Team
from nexus.orchestration import TopologyMode
from nexus.orchestration.graph import END, START


# ================================================================
# Demo 1: GRAPH mode
# ================================================================


class CodeReviewState(BaseModel):
    """Shared state flowing through the graph."""

    requirement: str = ""
    plan: str = ""
    code: str = ""
    review: str = ""
    approved: bool = False


async def demo_graph():
    """Graph mode: planner -> coder -> reviewer pipeline with conditional retry."""
    print("=" * 60)
    print("Demo 1: GRAPH Mode (static pipeline)")
    print("=" * 60)

    team = Team(**get_model_config(), mode=TopologyMode.GRAPH, verbose=False)

    planner = team.add_agent(
        "planner",
        instructions=(
            "You are a software architect. Given a requirement, "
            "produce a brief implementation plan (3-5 bullet points). "
            "Be concise, reply in plain text only."
        ),
    )
    coder = team.add_agent(
        "coder",
        instructions=(
            "You are a Python programmer. Given a plan, write clean Python code. "
            "Output ONLY the code, no explanation."
        ),
    )
    reviewer = team.add_agent(
        "reviewer",
        instructions=(
            "You are a code reviewer. Review the code for bugs and style. "
            "If acceptable, start your reply with 'APPROVED'. "
            "Otherwise start with 'NEEDS_WORK' and explain why. Be concise."
        ),
    )

    async def plan_node(state: dict) -> dict:
        from nexus.core.message import Message
        from nexus.core.task import Task

        task = Task(name="plan", description=state["requirement"])
        task.conversation.append(
            Message.user(f"Requirement: {state['requirement']}\n\nWrite a brief plan.")
        )
        result = await planner.run(task)
        print(f"\n[planner] Plan:\n{result.content[:300]}")
        return {"plan": result.content}

    async def code_node(state: dict) -> dict:
        from nexus.core.message import Message
        from nexus.core.task import Task

        task = Task(name="code", description="Write code")
        task.conversation.append(
            Message.user(f"Plan:\n{state['plan']}\n\nWrite the Python code.")
        )
        result = await coder.run(task)
        print(f"\n[coder] Code:\n{result.content[:500]}")
        return {"code": result.content}

    async def review_node(state: dict) -> dict:
        from nexus.core.message import Message
        from nexus.core.task import Task

        task = Task(name="review", description="Review code")
        task.conversation.append(Message.user(f"Review this code:\n{state['code']}"))
        result = await reviewer.run(task)
        approved = "APPROVED" in result.content.upper()
        print(
            f"\n[reviewer] {'APPROVED' if approved else 'NEEDS_WORK'}: "
            f"{result.content[:200]}"
        )
        return {"review": result.content, "approved": approved}

    def should_approve(state: dict) -> str:
        return "end" if state.get("approved", False) else "retry"

    graph = team.topology.graph
    graph.add_node("planner", plan_node)
    graph.add_node("coder", code_node)
    graph.add_node("reviewer", review_node)

    graph.add_edge(START, "planner")
    graph.add_edge("planner", "coder")
    graph.add_edge("coder", "reviewer")
    graph.add_conditional_edge(
        "reviewer", should_approve, {"end": END, "retry": "planner"}
    )

    print(f"\nGraph topology (Mermaid):\n{graph.to_mermaid()}\n")

    state = SharedState(
        CodeReviewState(
            requirement="Write a Python function that checks if a string is a palindrome"
        )
    )
    final = await team.topology.run_graph(state)

    print(f"\n{'=' * 40}")
    print(f"Final state: approved={final.data.approved}")
    print(f"{'=' * 40}")
    await team.close()


# ================================================================
# Demo 2: ROUTER mode
# ================================================================


async def demo_router():
    """Router mode: LLM triage agent picks the specialist for each task."""
    print("\n" + "=" * 60)
    print("Demo 2: ROUTER Mode (LLM-based routing)")
    print("=" * 60)

    team = Team(**get_model_config(), mode=TopologyMode.ROUTER, verbose=False)

    team.add_agent("coder", instructions="You are a Python programmer. Be concise.")
    team.add_agent("writer", instructions="You are a creative writer. Be concise.")
    team.add_agent("analyst", instructions="You are a data analyst. Be concise.")

    prompts = [
        ("Write a Python decorator for caching", "coder"),
        ("Write a haiku about the ocean", "writer"),
        ("What are the key trends in AI adoption in 2025?", "analyst"),
    ]
    for prompt, expected in prompts:
        print(f"\nUser: {prompt}")
        print(f"  (expected: {expected})")
        result = await team.run(prompt)
        print(f"  Assistant: {result[:200]}")

    await team.close()


# ================================================================
# Demo 3: ADAPTIVE mode
# ================================================================


async def demo_adaptive():
    """Adaptive mode: embedding similarity matches tasks to agent capabilities."""
    print("\n" + "=" * 60)
    print("Demo 3: ADAPTIVE Mode (embedding capability matching)")
    print("=" * 60)

    team = Team(
        **get_model_config(),
        mode=TopologyMode.ADAPTIVE,
        **get_embedding_config(),
        verbose=False,
    )

    team.add_agent(
        "coder",
        instructions="You are a Python programmer. Be concise.",
        capabilities=AgentCapability(coding=1.0, reasoning=0.8, math=0.7),
    )
    team.add_agent(
        "writer",
        instructions="You are a creative writer. Be concise.",
        capabilities=AgentCapability(writing=1.0, creativity=0.9, communication=0.8),
    )
    team.add_agent(
        "analyst",
        instructions="You are a data analyst. Be concise.",
        capabilities=AgentCapability(data_analysis=1.0, math=0.9, reasoning=0.8),
    )

    prompts = [
        ("Implement binary search in Python", "coder"),
        ("Write a short story about a robot learning to paint", "writer"),
        ("Compare bubble sort vs quicksort with Big-O analysis", "analyst"),
    ]
    for prompt, expected in prompts:
        print(f"\nUser: {prompt}")
        print(f"  (expected: {expected})")
        print("  Assistant: ", end="", flush=True)
        async for chunk in team.chat_stream(prompt):
            print(chunk, end="", flush=True)
        print()
        team.reset_conversation()

    await team.close()


# ================================================================
# Main
# ================================================================


async def main():
    demos = {"graph": demo_graph, "router": demo_router, "adaptive": demo_adaptive}
    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_modes.py [{'|'.join(demos.keys())}|all]")


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