#!/usr/bin/env python3
"""Functional tests for Swarm orchestration with real APIs."""

import asyncio
import os
import sys

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

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


async def get_test_agent():
    """Create a test agent with real API."""
    provider = 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"),
    )
    return Agent("test_agent", model=provider)


async def test_swarm_in_memory():
    """Test 1: Basic swarm with InMemoryEventBus."""
    print("\n" + "=" * 60)
    print("TEST 1: Swarm with InMemoryEventBus (simple handler)")
    print("=" * 60)

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

    async def uppercase_handler(event):
        return {"result": event.payload.get("prompt", "").upper()}

    swarm.subscribe("uppercase", uppercase_handler, publish_to="uppercase_done")

    request = bus.publish(
        "uppercase",
        {"prompt": "hello world"},
        correlation_id="corr-1",
        run_id="run-1",
    )

    processed = await swarm.drain()

    assert processed == 1, f"Expected 1, got {processed}"
    responses = bus.list_events(topic="uppercase_done", correlation_id="corr-1")
    assert len(responses) == 1
    assert responses[0].payload["result"] == "HELLO WORLD"
    print(f"✓ Processed {processed} event")
    print(f"✓ Response: {responses[0].payload}")
    return True


async def test_swarm_dispatch():
    """Test 2: Swarm dispatch method."""
    print("\n" + "=" * 60)
    print("TEST 2: Swarm dispatch() method")
    print("=" * 60)

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

    async def echo_handler(event):
        return {"echo": f"got: {event.payload.get('prompt')}"}

    swarm.subscribe("echo_topic", echo_handler, publish_to="echo_done")

    response = await swarm.dispatch(
        prompt="test message",
        request_topic="echo_topic",
        response_topic="echo_done",
        session_id="session-1",
    )

    assert "got: test message" in response.content, f"Unexpected: {response.content}"
    print(f"✓ Response content: {response.content}")
    return True


async def test_swarm_with_runnable():
    """Test 3: Swarm with Runnable (Agent) handler."""
    print("\n" + "=" * 60)
    print("TEST 3: Swarm with Agent as Runnable handler")
    print("=" * 60)

    from logos_adk.runtime import _default_runtime_holder

    _default_runtime_holder.reset()

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

    agent = await get_test_agent()
    swarm.subscribe("ask_agent", agent, publish_to="agent_done")

    request = bus.publish(
        "ask_agent",
        {"prompt": "What is 2+2?"},
        correlation_id="corr-agent",
        run_id="run-agent",
    )

    processed = await swarm.drain()

    assert processed == 1, f"Expected 1, got {processed}"
    responses = bus.list_events(topic="agent_done", correlation_id="corr-agent")
    assert len(responses) == 1
    content = responses[0].payload.get("content", "")
    print(f"✓ Processed {processed} event")
    print(f"✓ Agent response: {content[:100]}...")
    assert "4" in content, "Expected '4' in response"
    return True


async def test_swarm_chaining():
    """Test 4: Swarm with chained handlers (publish_to -> subscribe)."""
    print("\n" + "=" * 60)
    print("TEST 4: Swarm with chained handlers")
    print("=" * 60)

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

    async def step1(event):
        return {"step1": "done", "value": event.payload.get("value", 0) * 2}

    async def step2(event):
        return {"step2": "done", "value": event.payload.get("value", 0) + 10}

    swarm.subscribe("start", step1, publish_to="step1_done")
    swarm.subscribe("step1_done", step2, publish_to="final")

    request = bus.publish(
        "start",
        {"value": 5, "prompt": "start"},
        correlation_id="corr-chain",
        run_id="run-chain",
    )

    processed = await swarm.drain()

    assert processed == 2, f"First drain processed: {processed} (chained events processed together)"

    processed2 = await swarm.drain()
    assert processed2 == 0, f"Second drain processed: {processed2} (should be 0, all done)"

    responses = bus.list_events(topic="final", correlation_id="corr-chain")
    assert len(responses) == 1
    assert responses[0].payload["value"] == 20, f"Expected 20, got {responses[0].payload['value']}"
    print(f"✓ Step 1: value * 2 = 10")
    print(f"✓ Step 2: value + 10 = 20")
    print(f"✓ Chained execution works!")
    return True


async def test_swarm_with_sqlite():
    """Test 5: Swarm with SQLite event bus."""
    print("\n" + "=" * 60)
    print("TEST 5: Swarm with SQLiteEventBus")
    print("=" * 60)

    import tempfile

    with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as f:
        db_path = f.name

    bus = SQLiteEventBus(db_path)
    swarm = Swarm("sqlite-swarm", event_bus=bus)

    async def handler(event):
        return {"result": "sqlite works!"}

    swarm.subscribe("sqlite_test", handler, publish_to="sqlite_done")

    request = bus.publish(
        "sqlite_test",
        {"prompt": "test"},
        correlation_id="corr-sqlite",
        run_id="run-sqlite",
    )

    processed = await swarm.drain()

    assert processed == 1
    responses = bus.list_events(topic="sqlite_done", correlation_id="corr-sqlite")
    assert len(responses) == 1
    print(f"✓ SQLiteEventBus works!")

    os.unlink(db_path)
    return True


async def test_swarm_with_postgresql():
    """Test 6: Swarm with PostgreSQL event bus."""
    print("\n" + "=" * 60)
    print("TEST 6: Swarm with PostgreSQLEventBus")
    print("=" * 60)

    import uuid

    table_name = f"test_swarm_{uuid.uuid4().hex[:8]}"
    dsn = "postgresql://crm:crm@localhost:5432/crm"

    bus = PostgreSQLEventBus(dsn, table_name=table_name)

    swarm = Swarm("postgres-swarm", event_bus=bus)

    async def handler(event):
        return {"result": "postgresql works!", "prompt": event.payload.get("prompt")}

    swarm.subscribe("postgres_test", handler, publish_to="postgres_done")

    request = bus.publish(
        "postgres_test",
        {"prompt": "hello postgres"},
        correlation_id="corr-pg",
        run_id="run-pg",
    )

    processed = await swarm.drain()

    assert processed == 1, f"Expected 1, got {processed}"
    responses = bus.list_events(topic="postgres_done", correlation_id="corr-pg")
    assert len(responses) == 1
    print(f"✓ PostgreSQLEventBus works!")
    print(f"✓ Response: {responses[0].payload}")

    return True


async def test_swarm_drain_error_handling():
    """Test 7: Swarm drain error handling."""
    print("\n" + "=" * 60)
    print("TEST 7: Swarm drain error handling")
    print("=" * 60)

    bus = InMemoryEventBus()
    swarm = Swarm("error-swarm", event_bus=bus, retry_delay_seconds=0.0)

    async def fail_handler(event):
        raise RuntimeError("intentional failure")

    async def success_handler(event):
        return {"result": "success!"}

    swarm.subscribe("fail_topic", fail_handler)
    swarm.subscribe("success_topic", success_handler, publish_to="success_done")

    bus.publish("fail_topic", {"prompt": "fail"}, correlation_id="corr-fail")
    bus.publish("success_topic", {"prompt": "success"}, correlation_id="corr-success")

    processed = await swarm.drain(max_events=10)

    assert processed == 2, f"Expected 2, got {processed}"
    print(f"✓ Processed both events despite error: {processed}")

    responses = bus.list_events(topic="success_done", correlation_id="corr-success")
    assert len(responses) == 1
    print(f"✓ Success handler worked: {responses[0].payload}")
    return True


async def test_swarm_real_agent_chain():
    """Test 8: Full swarm with real agents in a chain."""
    print("\n" + "=" * 60)
    print("TEST 8: Real agent chain with swarm")
    print("=" * 60)

    from logos_adk.runtime import _default_runtime_holder

    _default_runtime_holder.reset()

    bus = InMemoryEventBus()
    swarm = Swarm("agent-chain", event_bus=bus)

    provider1 = 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"),
    )
    provider2 = 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"),
    )
    agent1 = Agent("analyzer_agent", model=provider1)
    agent2 = Agent("writer_agent", model=provider2)

    swarm.subscribe("analyze", agent1, publish_to="analyzed")
    swarm.subscribe("analyzed", agent2, publish_to="final")

    request = bus.publish(
        "analyze",
        {"prompt": "What is the capital of France?"},
        correlation_id="corr-real",
        run_id="run-real",
    )

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

    responses = bus.list_events(topic="final", correlation_id="corr-real")
    assert len(responses) == 1
    content = responses[0].payload.get("content", "")
    print(f"✓ First agent asked: What is the capital of France?")
    print(f"✓ Second agent responded: {content[:150]}...")
    assert "Paris" in content, "Expected Paris in response"
    print(f"✓ Real agent chain works!")
    return True


async def main():
    """Run all swarm tests."""
    print("=" * 60)
    print("SWARM FUNCTIONAL TESTS")
    print("=" * 60)

    results = []

    tests = [
        ("InMemoryEventBus", test_swarm_in_memory),
        ("dispatch()", test_swarm_dispatch),
        ("Runnable handler", test_swarm_with_runnable),
        ("Chained handlers", test_swarm_chaining),
        ("SQLiteEventBus", test_swarm_with_sqlite),
        ("PostgreSQLEventBus", test_swarm_with_postgresql),
        ("Error handling", test_swarm_drain_error_handling),
        ("Real agent chain", test_swarm_real_agent_chain),
    ]

    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)
