"""Test Microsoft AutoGen integration with Arzule observability.

This test demonstrates the AutoGen instrumentation capabilities.

Requirements:
    pip install pyautogen openai

Usage:
    # With mock (no API key needed)
    python test_autogen_integration.py

    # With real LLM (requires OPENAI_API_KEY)
    OPENAI_API_KEY=sk-... python test_autogen_integration.py --real
"""

from __future__ import annotations

import json
import os
import sys


def test_autogen_mock():
    """Test AutoGen integration with mock responses (no LLM calls)."""
    print("\n=== Testing AutoGen Integration (Mock Mode) ===\n")
    
    # Import and instrument before creating agents
    from arzule_ingest.autogen import instrument_autogen
    from arzule_ingest import ArzuleRun
    from arzule_ingest.sinks.file_jsonl import JsonlFileSink

    # Instrument AutoGen
    instrument_autogen()

    # Create output sink
    output_file = "out/autogen_test_mock.jsonl"
    os.makedirs("out", exist_ok=True)
    sink = JsonlFileSink(output_file)

    try:
        from autogen import ConversableAgent

        # Create agents with mock LLM (human_input_mode="NEVER" and no llm_config)
        # This allows testing the instrumentation without actual LLM calls
        
        agent_a = ConversableAgent(
            name="AgentA",
            llm_config=False,  # Disable LLM
            human_input_mode="NEVER",
        )
        
        agent_b = ConversableAgent(
            name="AgentB", 
            llm_config=False,
            human_input_mode="NEVER",
        )

        # Register auto-reply functions to simulate conversation
        message_count = {"a": 0, "b": 0}
        
        def agent_a_reply(recipient, messages, sender, config):
            message_count["a"] += 1
            if message_count["a"] >= 2:
                return True, "Goodbye from AgentA!"
            return True, f"AgentA response #{message_count['a']}"
        
        def agent_b_reply(recipient, messages, sender, config):
            message_count["b"] += 1
            if message_count["b"] >= 2:
                return True, "Goodbye from AgentB!"
            return True, f"AgentB response #{message_count['b']}"

        agent_a.register_reply([ConversableAgent], agent_a_reply)
        agent_b.register_reply([ConversableAgent], agent_b_reply)

        # Run with ArzuleRun context
        with ArzuleRun(
            tenant_id="test-tenant",
            project_id="autogen-test",
            sink=sink,
        ) as run:
            print(f"Run ID: {run.run_id}")
            print(f"Trace ID: {run.trace_id}")
            print()
            
            # Start a conversation
            result = agent_a.initiate_chat(
                agent_b,
                message="Hello from AgentA!",
                max_turns=3,
                summary_method="last_msg",
            )
            
            print("\nConversation completed!")
            if hasattr(result, "chat_history"):
                print(f"Messages exchanged: {len(result.chat_history)}")

    finally:
        sink.close()

    # Print captured events
    print(f"\n=== Captured Events ({output_file}) ===\n")
    with open(output_file) as f:
        for line in f:
            event = json.loads(line)
            event_type = event.get("event_type", "unknown")
            summary = event.get("summary", "")
            print(f"  {event_type}: {summary}")

    print(f"\nFull trace written to: {output_file}")


def test_autogen_real_llm():
    """Test AutoGen integration with real LLM calls."""
    print("\n=== Testing AutoGen Integration (Real LLM Mode) ===\n")
    
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        print("ERROR: OPENAI_API_KEY environment variable required for real LLM test")
        print("Run with: OPENAI_API_KEY=sk-... python test_autogen_integration.py --real")
        sys.exit(1)

    # Import and instrument
    from arzule_ingest.autogen import instrument_autogen
    from arzule_ingest import ArzuleRun
    from arzule_ingest.sinks.file_jsonl import JsonlFileSink

    instrument_autogen()

    output_file = "out/autogen_test_real.jsonl"
    os.makedirs("out", exist_ok=True)
    sink = JsonlFileSink(output_file)

    try:
        from autogen import AssistantAgent, UserProxyAgent

        # LLM config
        llm_config = {
            "model": "gpt-4o-mini",
            "api_key": api_key,
        }

        # Create assistant agent
        assistant = AssistantAgent(
            name="Assistant",
            llm_config=llm_config,
            system_message="You are a helpful AI assistant. Keep responses brief.",
        )

        # Create user proxy (no human input, auto-terminates)
        user_proxy = UserProxyAgent(
            name="UserProxy",
            human_input_mode="NEVER",
            max_consecutive_auto_reply=2,
            is_termination_msg=lambda x: x.get("content", "").strip().endswith("TERMINATE"),
            code_execution_config=False,  # Disable code execution for this test
        )

        # Run conversation
        with ArzuleRun(
            tenant_id="test-tenant",
            project_id="autogen-real-test",
            sink=sink,
        ) as run:
            print(f"Run ID: {run.run_id}")
            print(f"Trace ID: {run.trace_id}")
            print()

            result = user_proxy.initiate_chat(
                assistant,
                message="What is 2 + 2? Answer briefly then say TERMINATE.",
                max_turns=3,
            )

            print("\nConversation completed!")
            if hasattr(result, "chat_history"):
                print(f"Messages exchanged: {len(result.chat_history)}")
                for msg in result.chat_history:
                    role = msg.get("role", "unknown")
                    content = msg.get("content", "")[:100]
                    print(f"  [{role}]: {content}")

    finally:
        sink.close()

    # Print captured events
    print(f"\n=== Captured Events ({output_file}) ===\n")
    with open(output_file) as f:
        for line in f:
            event = json.loads(line)
            event_type = event.get("event_type", "unknown")
            summary = event.get("summary", "")[:80]
            print(f"  {event_type}: {summary}")

    print(f"\nFull trace written to: {output_file}")


def test_autogen_group_chat():
    """Test AutoGen group chat with multiple agents."""
    print("\n=== Testing AutoGen Group Chat ===\n")

    from arzule_ingest.autogen import instrument_autogen
    from arzule_ingest import ArzuleRun
    from arzule_ingest.sinks.file_jsonl import JsonlFileSink

    instrument_autogen()

    output_file = "out/autogen_group_chat.jsonl"
    os.makedirs("out", exist_ok=True)
    sink = JsonlFileSink(output_file)

    try:
        from autogen import ConversableAgent

        # Create multiple agents
        agents = []
        for name in ["Alice", "Bob", "Charlie"]:
            agent = ConversableAgent(
                name=name,
                llm_config=False,
                human_input_mode="NEVER",
            )
            agents.append(agent)

        alice, bob, charlie = agents

        # Simple reply handlers
        reply_counts = {a.name: 0 for a in agents}
        
        def make_reply(agent_name):
            def reply(recipient, messages, sender, config):
                reply_counts[agent_name] += 1
                if reply_counts[agent_name] >= 2:
                    return True, f"Bye from {agent_name}!"
                return True, f"{agent_name} says hello back!"
            return reply

        for agent in agents:
            agent.register_reply([ConversableAgent], make_reply(agent.name))

        with ArzuleRun(
            tenant_id="test-tenant",
            project_id="autogen-group-test",
            sink=sink,
        ) as run:
            print(f"Run ID: {run.run_id}")
            
            # Alice talks to Bob
            alice.initiate_chat(bob, message="Hi Bob!", max_turns=2, summary_method="last_msg")
            
            # Bob talks to Charlie
            bob.initiate_chat(charlie, message="Hi Charlie!", max_turns=2, summary_method="last_msg")

    finally:
        sink.close()

    print(f"\nEvents written to: {output_file}")
    
    # Count events by type
    event_counts = {}
    with open(output_file) as f:
        for line in f:
            event = json.loads(line)
            event_type = event.get("event_type", "unknown")
            event_counts[event_type] = event_counts.get(event_type, 0) + 1

    print("\nEvent counts:")
    for event_type, count in sorted(event_counts.items()):
        print(f"  {event_type}: {count}")


if __name__ == "__main__":
    if "--real" in sys.argv:
        test_autogen_real_llm()
    elif "--group" in sys.argv:
        test_autogen_group_chat()
    else:
        test_autogen_mock()
        print("\n" + "=" * 60)
        test_autogen_group_chat()

