#!/usr/bin/env python3
"""Test that agents communicate with each other through swarm event bus."""

import asyncio
import os
import sys

sys.path.insert(0, os.path.dirname(__file__))

from logos_adk.event_bus import InMemoryEventBus
from logos_adk.swarm import Swarm
from logos_adk.agent import Agent
from logos_adk.providers.openai import OpenAICompatibleProvider


async def get_provider():
    return OpenAICompatibleProvider(
        model=os.getenv("OPENAI_COMPAT_MODEL", "gpt-4o-mini"),
        base_url=os.getenv("OPENAI_COMPAT_BASE_URL", "https://api.example.com/v1"),
        api_key=os.getenv("OPENAI_COMPAT_API_KEY", "test-api-key"),
    )


async def test_agent_to_agent_message():
    """Test 1: Agent sends message to another agent via event bus."""
    print("\n" + "=" * 60)
    print("TEST 1: Agent → Agent message via event bus")
    print("=" * 60)

    from logos_adk.runtime import _default_runtime_holder

    _default_runtime_holder.reset()

    bus = InMemoryEventBus()
    swarm = Swarm("message-swarm", event_bus=bus)

    provider = await get_provider()
    sender = Agent("sender", model=provider)
    receiver = Agent("receiver", model=provider)

    swarm.subscribe("send_message", sender, publish_to="deliver_message")
    swarm.subscribe("deliver_message", receiver, publish_to="message_received")

    bus.publish(
        "send_message",
        {"prompt": "Tell the receiver: The secret code is 42"},
        correlation_id="msg-1",
        run_id="run-1",
    )

    await swarm.drain()
    await swarm.drain()

    events = bus.list_events(topic="message_received", correlation_id="msg-1")
    assert len(events) == 1

    content = events[0].payload.get("content", "")
    print(f"✓ Sender sent: 'Tell the receiver: The secret code is 42'")
    print(f"✓ Receiver responded: {content[:200]}...")

    assert "42" in content or "secret" in content.lower(), "Receiver should acknowledge the message"
    print(f"✓ Inter-agent communication works!")
    return True


async def test_multi_agent_workflow():
    """Test 2: Three agents in a pipeline workflow."""
    print("\n" + "=" * 60)
    print("TEST 2: Three agents pipeline (Researcher → Writer → Editor)")
    print("=" * 60)

    from logos_adk.runtime import _default_runtime_holder

    _default_runtime_holder.reset()

    bus = InMemoryEventBus()
    swarm = Swarm("pipeline-swarm", event_bus=bus)

    provider = await get_provider()
    researcher = Agent("researcher", model=provider)
    writer = Agent("writer", model=provider)
    editor = Agent("editor", model=provider)

    swarm.subscribe("start_research", researcher, publish_to="research_done")
    swarm.subscribe("research_done", writer, publish_to="writing_done")
    swarm.subscribe("writing_done", editor, publish_to="final_review")

    bus.publish(
        "start_research",
        {"prompt": "Research: What is quantum computing? Give me 3 key facts."},
        correlation_id="pipeline-1",
        run_id="run-pipeline",
    )

    for _ in range(4):
        await swarm.drain()

    events = bus.list_events(topic="final_review", correlation_id="pipeline-1")
    assert len(events) == 1

    content = events[0].payload.get("content", "")
    print(f"✓ Pipeline completed!")
    print(f"✓ Final output: {content[:300]}...")
    return True


async def test_agent_shares_context():
    """Test 3: Agents share context through event payload."""
    print("\n" + "=" * 60)
    print("TEST 3: Agents share context through event payload")
    print("=" * 60)

    from logos_adk.runtime import _default_runtime_holder

    _default_runtime_holder.reset()

    bus = InMemoryEventBus()
    swarm = Swarm("context-swarm", event_bus=bus)

    provider = await get_provider()
    analyzer = Agent("analyzer", model=provider)
    summarizer = Agent("summarizer", model=provider)

    swarm.subscribe("raw_data", analyzer, publish_to="analysis_done")
    swarm.subscribe("analysis_done", summarizer, publish_to="summary_done")

    bus.publish(
        "raw_data",
        {
            "prompt": "Analyze this data: The company had 100 users in January, 150 in February, 200 in March",
            "context": {"company": "Acme Corp", "year": 2024},
        },
        correlation_id="context-1",
        run_id="run-context",
    )

    await swarm.drain()

    analysis_events = bus.list_events(topic="analysis_done", correlation_id="context-1")
    assert len(analysis_events) == 1

    analysis_payload = analysis_events[0].payload
    print(f"✓ Analyzer received context: {analysis_payload.get('context')}")
    print(f"✓ Analyzer output: {analysis_payload.get('content', '')[:150]}...")

    await swarm.drain()

    summary_events = bus.list_events(topic="summary_done", correlation_id="context-1")
    assert len(summary_events) == 1

    summary_content = summary_events[0].payload.get("content", "")
    print(f"✓ Summarizer received analysis")
    print(f"✓ Summary: {summary_content[:200]}...")

    print(f"✓ Context sharing works!")
    return True


async def test_agent_responses_to_agent():
    """Test 4: Agent responds directly to another agent's output."""
    print("\n" + "=" * 60)
    print("TEST 4: Agent responds to another agent (Q&A chain)")
    print("=" * 60)

    from logos_adk.runtime import _default_runtime_holder

    _default_runtime_holder.reset()

    bus = InMemoryEventBus()
    swarm = Swarm("qa-swarm", event_bus=bus)

    provider = await get_provider()
    asker = Agent("asker", model=provider)
    answerer = Agent("answerer", model=provider)

    swarm.subscribe("question", asker, publish_to="question_ready")
    swarm.subscribe("question_ready", answerer, publish_to="answer_ready")
    swarm.subscribe("answer_ready", asker, publish_to="follow_up_ready")

    bus.publish(
        "question",
        {"prompt": "What is machine learning?"},
        correlation_id="qa-1",
        run_id="run-qa",
    )

    for _ in range(4):
        await swarm.drain()

    events = bus.list_events(topic="follow_up_ready", correlation_id="qa-1")
    assert len(events) == 1

    content = events[0].payload.get("content", "")
    print(f"✓ Q&A chain completed!")
    print(f"✓ Final response: {content[:300]}...")
    return True


async def main():
    """Run all inter-agent communication tests."""
    print("=" * 60)
    print("INTER-AGENT COMMUNICATION TESTS")
    print("=" * 60)

    results = []

    tests = [
        ("Agent → Agent message", test_agent_to_agent_message),
        ("Three-agent pipeline", test_multi_agent_workflow),
        ("Context sharing", test_agent_shares_context),
        ("Q&A chain", test_agent_responses_to_agent),
    ]

    for name, test_func in tests:
        try:
            await test_func()
            results.append((name, True, None))
        except Exception as e:
            print(f"✗ FAILED: {e}")
            import traceback

            traceback.print_exc()
            results.append((name, False, str(e)))

    print("\n" + "=" * 60)
    print("RESULTS")
    print("=" * 60)

    passed = sum(1 for _, ok, _ in results if ok)
    total = len(results)

    for name, ok, err in results:
        status = "✓ PASS" if ok else "✗ FAIL"
        print(f"{status}: {name}")
        if err:
            print(f"    Error: {err}")

    print(f"\n{passed}/{total} tests passed")
    return passed == total


if __name__ == "__main__":
    success = asyncio.run(main())
    sys.exit(0 if success else 1)
