import os
from dotenv import load_dotenv
from IPython.core.error import (
    StdinNotImplementedError,
)
import asyncio

import agentops

from autogen_agentchat.agents import AssistantAgent, UserProxyAgent
from autogen_ext.models.openai import OpenAIChatCompletionClient

from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.conditions import MaxMessageTermination
from autogen_agentchat.ui import Console

load_dotenv()
os.environ["AGENTOPS_API_KEY"] = os.getenv("AGENTOPS_API_KEY", "your_api_key_here")
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "your_openai_api_key_here")

# When initializing AgentOps, you can pass in optional tags to help filter sessions
agentops.init(auto_start_session=False)
tracer = agentops.start_trace(
    trace_name="Microsoft Agent Chat Example", tags=["autogen-chat", "microsoft-autogen", "agentops-example"]
)

# Define model and API key
model_name = "gpt-4-turbo"  # Or "gpt-4o" / "gpt-4o-mini" as per migration guide examples
api_key = os.getenv("OPENAI_API_KEY")

# Create the model client
model_client = OpenAIChatCompletionClient(model=model_name, api_key=api_key)

# Create the agent that uses the LLM.
assistant = AssistantAgent(
    name="assistant",
    system_message="You are a helpful assistant.",  # Added system message for clarity
    model_client=model_client,
)

user_proxy_initiator = UserProxyAgent("user_initiator")


async def main():
    termination = MaxMessageTermination(max_messages=2)

    group_chat = RoundRobinGroupChat(
        [user_proxy_initiator, assistant],  # Corrected: agents as positional argument
        termination_condition=termination,
    )

    chat_task = "How can I help you today?"
    print(f"User Initiator: {chat_task}")

    try:
        stream = group_chat.run_stream(task=chat_task)
        await Console().run(stream)
        agentops.end_trace(tracer, end_state="Success")

    except StdinNotImplementedError:
        print("StdinNotImplementedError: This typically happens in non-interactive environments.")
        print("Skipping interactive part of chat for automation.")
        agentops.end_trace(tracer, end_state="Indeterminate")
    except Exception as e:
        print(f"An error occurred: {e}")
        agentops.end_trace(tracer, end_state="Error")
    finally:
        await model_client.close()


if __name__ == "__main__":
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        loop = None

    if loop and loop.is_running():
        import nest_asyncio

        nest_asyncio.apply()
        asyncio.run(main())
    else:
        asyncio.run(main())