#!/usr/bin/env python3
"""
AutoGen Integration Example

Demonstrates how to monitor AutoGen multi-agent conversations with Rora.

Requirements:
    pip install rora[autogen] autogen-agentchat
"""

import os

RORA_API_KEY = os.environ.get("RORA_API_KEY", "your_rora_key")
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "your_openai_key")

os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY

from autogen import AssistantAgent, UserProxyAgent
from rora.integrations.autogen import (
    RoraAgentWrapper,
    RoraConversationTracker,
    expose,
)


def example_wrapped_agent():
    """Wrap individual agents for metrics tracking."""
    print("\n=== Wrapped Agent Example ===\n")

    assistant = AssistantAgent(
        name="assistant",
        llm_config={"model": "gpt-4"},
    )

    wrapped = RoraAgentWrapper(assistant, api_key=RORA_API_KEY)

    user_proxy = UserProxyAgent(
        name="user_proxy",
        human_input_mode="NEVER",
        max_consecutive_auto_reply=1,
    )

    user_proxy.initiate_chat(wrapped.agent, message="What is 2 + 2?")

    print(f"\nMetrics: {wrapped.metrics}")


def example_conversation_tracker():
    """Track full multi-agent conversations."""
    print("\n=== Conversation Tracker Example ===\n")

    assistant = AssistantAgent(
        name="assistant",
        llm_config={"model": "gpt-4"},
    )

    user_proxy = UserProxyAgent(
        name="user_proxy",
        human_input_mode="NEVER",
        max_consecutive_auto_reply=2,
    )

    tracker = RoraConversationTracker(api_key=RORA_API_KEY)
    tracker.track(assistant, user_proxy)

    user_proxy.initiate_chat(assistant, message="Write a haiku about AI")

    print(f"\nConversation metrics: {tracker.metrics}")


def example_expose():
    """Expose an AutoGen agent for outside-in monitoring."""
    print("\n=== Expose Example ===\n")

    assistant = AssistantAgent(
        name="assistant",
        llm_config={"model": "gpt-4"},
        system_message="You are a helpful coding assistant.",
    )

    print("Exposing agent to Rora's distributed network...")
    print("Press Ctrl+C to stop\n")

    expose(
        assistant,
        api_key=RORA_API_KEY,
        agent_name="AutoGen Coding Assistant",
    )


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1:
        if sys.argv[1] == "wrapped":
            example_wrapped_agent()
        elif sys.argv[1] == "tracker":
            example_conversation_tracker()
        elif sys.argv[1] == "expose":
            example_expose()
        else:
            print(f"Unknown example: {sys.argv[1]}")
            print("Usage: python autogen_integration.py [wrapped|tracker|expose]")
    else:
        example_wrapped_agent()
