"""
multi_agent.py — Multi-Agent Orchestration with Semantic Kernel
===============================================================
Demonstrates a multi-agent pattern where two specialized agents collaborate:
  - ResearchAgent  : Finds information and provides factual answers
  - WriterAgent    : Takes research output and produces polished content

Uses Semantic Kernel's agent framework for orchestration.

Prerequisites:
    pip install semantic-kernel azure-identity python-dotenv

Usage:
    python multi_agent.py
"""

import asyncio
import os
import sys
from dotenv import load_dotenv

from semantic_kernel import Kernel
from semantic_kernel.connectors.ai.open_ai import AzureChatCompletion
from semantic_kernel.agents import ChatCompletionAgent, AgentGroupChat
from semantic_kernel.agents.strategies.termination.termination_strategy import TerminationStrategy
from semantic_kernel.contents.chat_message_content import ChatMessageContent

load_dotenv()


def get_env(key: str) -> str:
    value = os.getenv(key)
    if not value:
        print(f"ERROR: '{key}' is not set. See .env.example.")
        sys.exit(1)
    return value


# ---------------------------------------------------------------------------
# Custom termination strategy — stop after the Writer has responded
# ---------------------------------------------------------------------------
class WriterTermination(TerminationStrategy):
    """Terminate the group chat once the WriterAgent has produced output."""

    async def should_agent_terminate(self, agent, history: list[ChatMessageContent]) -> bool:
        return agent.name == "WriterAgent" and len(history) >= 2


# ---------------------------------------------------------------------------
# Build agents
# ---------------------------------------------------------------------------
def create_kernel() -> Kernel:
    """Create a Semantic Kernel instance with an Azure OpenAI chat service."""
    kernel = Kernel()

    kernel.add_service(
        AzureChatCompletion(
            service_id="chat",
            deployment_name=get_env("AZURE_OPENAI_CHAT_DEPLOYMENT"),
            endpoint=get_env("AZURE_OPENAI_ENDPOINT"),
            api_key=os.getenv("AZURE_OPENAI_API_KEY", ""),
        )
    )
    return kernel


def create_research_agent(kernel: Kernel) -> ChatCompletionAgent:
    """Agent that researches and provides factual bullet points."""
    return ChatCompletionAgent(
        kernel=kernel,
        service_id="chat",
        name="ResearchAgent",
        instructions=(
            "You are a research assistant. When given a topic, provide 5 concise, "
            "factual bullet points with key information. Be specific and cite "
            "concepts accurately. Do not write prose — only bullet points."
        ),
    )


def create_writer_agent(kernel: Kernel) -> ChatCompletionAgent:
    """Agent that writes polished content from research notes."""
    return ChatCompletionAgent(
        kernel=kernel,
        service_id="chat",
        name="WriterAgent",
        instructions=(
            "You are a professional technical writer. You receive research bullet "
            "points from the ResearchAgent. Transform them into a polished, "
            "well-structured short article (200-300 words). Use clear headings "
            "and an engaging tone suitable for a developer blog."
        ),
    )


# ---------------------------------------------------------------------------
# Orchestration
# ---------------------------------------------------------------------------
async def run_multi_agent(topic: str) -> None:
    """Run the research → writer pipeline on a given topic."""
    kernel = create_kernel()

    researcher = create_research_agent(kernel)
    writer = create_writer_agent(kernel)

    # Create a group chat with both agents
    group_chat = AgentGroupChat(
        agents=[researcher, writer],
        termination_strategy=WriterTermination(
            agents=[writer],
            maximum_iterations=4,
        ),
    )

    print(f"📝 Topic: {topic}\n")
    print("=" * 60)

    # Kick off the conversation
    await group_chat.add_chat_message(
        ChatMessageContent(role="user", content=f"Research this topic: {topic}")
    )

    # Stream agent responses
    async for message in group_chat.invoke():
        agent_name = message.name or "Unknown"
        print(f"\n🤖 [{agent_name}]:")
        print(message.content)
        print("-" * 60)

    print("\n✅ Multi-agent orchestration complete!")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> None:
    topic = "Azure AI Foundry: What it is, key features, and how it compares to using Azure OpenAI directly"
    asyncio.run(run_multi_agent(topic))


if __name__ == "__main__":
    main()