#!/usr/bin/env python3
"""11 — Parallel fan-out with streaming.

Run multiple agents concurrently on the same input and collect results.
Uses ParallelFanOut for concurrent execution and streaming events.

Requirements:
    pip install koan[llm]
    export OPENAI_API_KEY=sk-...
"""

from __future__ import annotations

import asyncio

from koan.agents import (
    AgentComplete,
    AgentConfig,
    AgentStart,
    ParallelFanOut,
    ReactAgent,
    StreamingReactAgent,
    TextDelta,
    ToolCallEnd,
    ToolCallStart,
)
from koan.llm.client import LLMClient
from koan.tools import ToolRegistry, tool


@tool
def analyze_sentiment(text: str) -> str:
    """Analyze the sentiment of text."""
    # Simulated — in production, use an NLP model
    positive = ["good", "great", "excellent", "happy", "love"]
    negative = ["bad", "terrible", "awful", "sad", "hate"]
    lower = text.lower()
    pos = sum(1 for w in positive if w in lower)
    neg = sum(1 for w in negative if w in lower)
    if pos > neg:
        return "Positive sentiment"
    if neg > pos:
        return "Negative sentiment"
    return "Neutral sentiment"


@tool
def count_words(text: str) -> int:
    """Count the number of words in text."""
    return len(text.split())


async def demo_parallel_run() -> None:
    """Run agents in parallel and collect results."""
    print("=" * 60)
    print("Parallel Fan-Out (non-streaming)")
    print("=" * 60)

    async with LLMClient() as client:
        summarizer = ReactAgent(
            "summarizer",
            client=client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="Summarize the input in one sentence.",
            ),
        )

        translator = ReactAgent(
            "translator",
            client=client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="Translate the input to French. Only output the French translation.",
            ),
        )

        sentiment_tools = ToolRegistry()
        sentiment_tools.add(analyze_sentiment)
        sentiment_tools.add(count_words)

        analyzer = ReactAgent(
            "analyzer",
            client=client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt=(
                    "Analyze the input: use the analyze_sentiment tool for sentiment and "
                    "count_words for word count. Report both results."
                ),
            ),
            tools=sentiment_tools,
        )

        fan_out = ParallelFanOut([summarizer, translator, analyzer])

        text = (
            "The new AI framework makes building agents remarkably easy. "
            "The documentation is excellent and the API is intuitive. "
            "I love how tools are automatically converted to JSON schemas."
        )

        print(f"\nInput: {text}\n")
        result = await fan_out.run(text)

        print(f"Agent: {result.agent_name}")
        print(f"Total iterations: {result.iterations}")
        print(f"Total tool calls: {result.tool_calls_made}")
        print(f"\nCombined output:\n{result.output}")


async def demo_parallel_stream() -> None:
    """Run streaming agents in parallel, printing events from all."""
    print("\n" + "=" * 60)
    print("Parallel Fan-Out (streaming)")
    print("=" * 60)

    async with LLMClient() as client:
        agent_a = StreamingReactAgent(
            "poet",
            client=client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="Write a two-line haiku about the input topic.",
            ),
        )

        agent_b = StreamingReactAgent(
            "factoid",
            client=client,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o-mini",
                system_prompt="State one surprising fact about the input topic in one sentence.",
            ),
        )

        fan_out = ParallelFanOut([agent_a, agent_b])

        print("\nInput: 'the ocean'\n")

        async for event in fan_out.run_stream("the ocean"):
            if isinstance(event, AgentStart):
                print(f"\n[{event.agent_name}] started")
            elif isinstance(event, TextDelta):
                print(event.text, end="", flush=True)
            elif isinstance(event, ToolCallStart):
                print(f"\n  -> {event.tool_name}({event.arguments})")
            elif isinstance(event, ToolCallEnd):
                print(f"  <- {event.result}")
            elif isinstance(event, AgentComplete):
                print(f"\n[Done: {event.iterations} iters]")

        print()


async def main() -> None:
    await demo_parallel_run()
    await demo_parallel_stream()


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