"""Test Microsoft AutoGen integration with real LLM (Anthropic/OpenAI).

Usage:
    # With Anthropic (via litellm)
    ANTHROPIC_API_KEY=... python test_autogen_real.py

    # With OpenAI  
    OPENAI_API_KEY=... python test_autogen_real.py --openai
"""

from __future__ import annotations

import json
import os
import sys


def test_autogen_with_anthropic():
    """Test AutoGen with Anthropic Claude via the OpenAI-compatible endpoint."""
    print("\n=== Testing AutoGen with Anthropic Claude ===\n")

    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        print("ERROR: ANTHROPIC_API_KEY environment variable required")
        print("Run with: ANTHROPIC_API_KEY=... python test_autogen_real.py")
        sys.exit(1)

    # 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_anthropic_test.jsonl"
    os.makedirs("out", exist_ok=True)
    sink = JsonlFileSink(output_file)

    try:
        from autogen import AssistantAgent, UserProxyAgent

        # Configure for Anthropic Claude
        # AutoGen 0.2.x supports various model providers through config
        llm_config = {
            "config_list": [
                {
                    "model": "claude-sonnet-4-20250514",
                    "api_key": api_key,
                    "api_type": "anthropic",
                }
            ],
            "temperature": 0.7,
        }

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

        # Create user proxy (no human input, terminates on specific message)
        user_proxy = UserProxyAgent(
            name="UserProxy",
            human_input_mode="NEVER",
            max_consecutive_auto_reply=2,
            is_termination_msg=lambda x: "TERMINATE" in x.get("content", "").upper(),
            code_execution_config=False,
        )

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

            # Start conversation
            result = user_proxy.initiate_chat(
                assistant,
                message="What is 2 + 2? Give a one-word answer then say TERMINATE.",
                max_turns=3,
                summary_method="last_msg",
            )

            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", msg.get("name", "unknown"))
                    content = str(msg.get("content", ""))[:100]
                    print(f"  [{role}]: {content}")

    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
    finally:
        sink.close()

    # Print captured events
    print(f"\n=== Captured Events ({output_file}) ===\n")
    try:
        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}")
    except FileNotFoundError:
        print("  (no events captured)")

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


def test_autogen_with_openai():
    """Test AutoGen with OpenAI GPT."""
    print("\n=== Testing AutoGen with OpenAI GPT ===\n")

    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        print("ERROR: OPENAI_API_KEY environment variable required")
        print("Run with: OPENAI_API_KEY=... python test_autogen_real.py --openai")
        sys.exit(1)

    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_openai_test.jsonl"
    os.makedirs("out", exist_ok=True)
    sink = JsonlFileSink(output_file)

    try:
        from autogen import AssistantAgent, UserProxyAgent

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

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

        user_proxy = UserProxyAgent(
            name="UserProxy",
            human_input_mode="NEVER",
            max_consecutive_auto_reply=2,
            is_termination_msg=lambda x: "TERMINATE" in x.get("content", "").upper(),
            code_execution_config=False,
        )

        with ArzuleRun(
            tenant_id="test-tenant",
            project_id="autogen-openai-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? Give a one-word answer then say TERMINATE.",
                max_turns=3,
                summary_method="last_msg",
            )

            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", msg.get("name", "unknown"))
                    content = str(msg.get("content", ""))[:100]
                    print(f"  [{role}]: {content}")

    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
    finally:
        sink.close()

    print(f"\n=== Captured Events ({output_file}) ===\n")
    try:
        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}")
    except FileNotFoundError:
        print("  (no events captured)")

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


if __name__ == "__main__":
    if "--openai" in sys.argv:
        test_autogen_with_openai()
    else:
        test_autogen_with_anthropic()

