"""
03_parallel_agents.py — Fan-out / fan-in: coordinator spawns 3 agents in parallel

Tests the auto-gather feature: coordinator declares llm_input=["fetch_a","fetch_b","fetch_c"]
and the library automatically runs all three in parallel — no asyncio.gather needed.

Pipeline:
  coordinator
    ├─ fetch_a  (parallel)
    ├─ fetch_b  (parallel)
    └─ fetch_c  (parallel)
  → aggregator  (merge point)

WHAT THE DASHBOARD SHOULD SHOW:
  Live Graph:
    • coordinator node fans out to 3 sibling nodes (fetch_a, fetch_b, fetch_c)
    • All 3 siblings merge into aggregator
    • Faint "parallel" band label across the sibling column
    • coordinator → fetch_a/b/c edges + fetch_a/b/c → aggregator edges

  State View:
    • Steps show fetch_a, fetch_b, fetch_c starting roughly simultaneously
    • aggregator step shows merged input from all 3

  Time Travel:
    • Scrub back to see parallel agents still in-flight before aggregation

  Token Dashboard:
    • 5 bars: coordinator + fetch_a + fetch_b + fetch_c + aggregator
    • Aggregator bar tallest (synthesis prompt is longest)

  Run Comparison:
    • coordinator shows 0ms wall-clock (auto-gather returns immediately)
    • fetch_a/b/c show overlapping timings (true parallel)

  Failure History:
    • Empty
"""
import asyncio
import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from agentradar import radar, reset_run

load_dotenv()

llm = ChatOpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url=os.getenv("DEEPSEEK_BASE_URL"),
    model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
    temperature=0.7,
)


@radar(tag="coordinator")
async def coordinator(query: str):
    """
    Coordinator — library auto-spawns fetch_a, fetch_b, fetch_c in parallel.
    No asyncio.gather() needed: declare llm_input list and return None (pass).
    The library returns the gathered [result_a, result_b, result_c] list.
    """
    print(f"[Coordinator] Auto-spawning parallel fetchers for: {query}")


@radar(tag="fetch_a", llm_input=["coordinator"])
async def fetch_source_a(query: str) -> str:
    """Parallel worker A — technology perspective."""
    await asyncio.sleep(0.3)
    message = llm.invoke(f"From a technology perspective, answer in 2 sentences: {query}")
    return f"[Tech] {message.content}"


@radar(tag="fetch_b", llm_input=["coordinator"])
async def fetch_source_b(query: str) -> str:
    """Parallel worker B — business perspective."""
    await asyncio.sleep(0.2)
    message = llm.invoke(f"From a business perspective, answer in 2 sentences: {query}")
    return f"[Business] {message.content}"


@radar(tag="fetch_c", llm_input=["coordinator"])
async def fetch_source_c(query: str) -> str:
    """Parallel worker C — societal perspective."""
    await asyncio.sleep(0.4)
    message = llm.invoke(f"From a societal perspective, answer in 2 sentences: {query}")
    return f"[Society] {message.content}"


@radar(tag="aggregator", llm_input=["fetch_a", "fetch_b", "fetch_c"])
async def aggregator(results: list) -> str:
    """Merge point — synthesize all 3 parallel perspectives into one answer."""
    combined = "\n".join(results)
    message = llm.invoke(
        f"Synthesize these three perspectives into one balanced, cohesive answer:\n{combined}"
    )
    return message.content


async def main():
    reset_run()
    query = "What is the future of artificial intelligence?"

    print("=== Parallel Agents (Auto-Gather) ===\n")

    # coordinator auto-runs fetch_a, fetch_b, fetch_c in parallel — returns list of 3 results
    results = await coordinator(query)
    print(f"[Done] Got {len(results)} parallel results")
    for r in results:
        print(f"  • {r[:80]}...")

    final = await aggregator(results)
    print(f"\n[Aggregated]\n{final[:300]}...")

    print("\n✓ Dashboard: fan-out graph, parallel band, 3-to-1 merge edges to aggregator")


if __name__ == "__main__":
    asyncio.run(main())
    input("\nPress Enter to view dashboard (http://localhost:7111) or Ctrl+C to exit...")
