#!/usr/bin/env python3
"""
Full Swarm Agent Communication Test Suite
=========================================
Comprehensive testing of inter-agent communication, event bus,
swarm orchestration, and real API integration.
"""

import asyncio
import json
import os
import sys
import time
from datetime import datetime
from typing import Any

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


class TestReport:
    """Test report collector."""

    def __init__(self):
        self.results: list[dict[str, Any]] = []
        self.start_time = time.time()

    def add_result(self, name: str, passed: bool, details: str = "", duration: float = 0):
        self.results.append(
            {
                "name": name,
                "passed": passed,
                "details": details,
                "duration": duration,
                "timestamp": datetime.now().isoformat(),
            }
        )

    def print_summary(self):
        total = len(self.results)
        passed = sum(1 for r in self.results if r["passed"])
        failed = total - passed
        duration = time.time() - self.start_time

        print("\n" + "=" * 70)
        print("                    SWARM AGENTS TEST REPORT")
        print("=" * 70)
        print(f"Timestamp: {datetime.now().isoformat()}")
        print(f"Total tests: {total}")
        print(f"Passed: {passed} ✓")
        print(f"Failed: {failed} ✗")
        print(f"Duration: {duration:.2f}s")
        print("=" * 70)

        for i, r in enumerate(self.results, 1):
            status = "✓ PASS" if r["passed"] else "✗ FAIL"
            print(f"\n{i:2}. {status}: {r['name']}")
            if r["details"]:
                print(f"    {r['details']}")
            if r["duration"]:
                print(f"    Duration: {r['duration']:.2f}s")

        print("\n" + "=" * 70)
        print(f"FINAL RESULT: {passed}/{total} tests passed ({100 * passed / total:.1f}%)")
        print("=" * 70)

        return passed == total


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"),
    )


def reset_runtime():
    from logos_adk.runtime import _default_runtime_holder

    _default_runtime_holder.reset()


# ============================================================================
# TEST 1: Basic Event Bus Functionality
# ============================================================================


async def test_event_bus_publish_claim(report: TestReport):
    """Test event bus publish/claim/ack cycle."""
    name = "Event Bus: Publish → Claim → Ack"
    start = time.time()

    try:
        bus = InMemoryEventBus()

        # Publish event
        event = bus.publish("test_topic", {"data": "hello"}, correlation_id="corr-1")
        assert event.status == "queued"

        # Claim event
        claimed = bus.claim("worker-1", ["test_topic"])
        assert claimed is not None
        assert claimed.id == event.id
        assert claimed.status == "claimed"

        # Ack event
        bus.ack(claimed.id, "worker-1")
        completed = bus.get_event(event.id)
        assert completed.status == "completed"

        report.add_result(name, True, f"Event {event.id[:8]}...", time.time() - start)
        return True
    except Exception as e:
        report.add_result(name, False, str(e), time.time() - start)
        return False


async def test_event_bus_filters(report: TestReport):
    """Test event bus filtering by correlation_id, run_id, etc."""
    name = "Event Bus: Filtering (correlation_id, run_id)"
    start = time.time()

    try:
        bus = InMemoryEventBus()

        # Publish multiple events with different filters
        bus.publish("topic1", {"data": "a"}, correlation_id="corr-a", run_id="run-1")
        bus.publish("topic2", {"data": "b"}, correlation_id="corr-a", run_id="run-1")
        bus.publish("topic3", {"data": "c"}, correlation_id="corr-b", run_id="run-2")

        # Filter by correlation_id
        corr_a = bus.list_events(correlation_id="corr-a")
        assert len(corr_a) == 2

        # Filter by run_id
        run_1 = bus.list_events(run_id="run-1")
        assert len(run_1) == 2

        # Filter by topic
        topic1 = bus.list_events(topic="topic1")
        assert len(topic1) == 1

        report.add_result(name, True, f"Filters work correctly", time.time() - start)
        return True
    except Exception as e:
        report.add_result(name, False, str(e), time.time() - start)
        return False


# ============================================================================
# TEST 2: Swarm Basic Functionality
# ============================================================================


async def test_swarm_subscribe_and_drain(report: TestReport):
    """Test swarm subscription and drain."""
    name = "Swarm: Subscribe → Drain"
    start = time.time()

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

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

        swarm.subscribe("process", handler, publish_to="done")

        bus.publish("process", {"prompt": "hello"})

        processed = await swarm.drain()

        events = bus.list_events(topic="done")
        assert len(events) == 1
        assert events[0].payload["result"] == "HELLO"

        report.add_result(name, True, f"Processed {processed} events", time.time() - start)
        return True
    except Exception as e:
        report.add_result(name, False, str(e), time.time() - start)
        return False


async def test_swarm_dispatch(report: TestReport):
    """Test swarm dispatch method."""
    name = "Swarm: Dispatch method"
    start = time.time()

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

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

        swarm.subscribe("echo", handler, publish_to="echo_done")

        response = await swarm.dispatch(
            prompt="test", request_topic="echo", response_topic="echo_done"
        )

        assert "got: test" in response.content

        report.add_result(name, True, "Dispatch works", time.time() - start)
        return True
    except Exception as e:
        report.add_result(name, False, str(e), time.time() - start)
        return False


# ============================================================================
# TEST 3: Agent Integration
# ============================================================================


async def test_agent_in_swarm(report: TestReport):
    """Test Agent as swarm handler."""
    name = "Swarm: Agent as Handler (real API)"
    start = time.time()

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

        provider = await get_provider()
        agent = Agent("assistant", model=provider)

        swarm.subscribe("ask", agent, publish_to="answer")

        bus.publish("ask", {"prompt": "What is 1+1?"})

        processed = await swarm.drain()

        events = bus.list_events(topic="answer")
        assert len(events) == 1
        content = events[0].payload.get("content", "") or ""
        has_number = "2" in content or "two" in content.lower()
        assert has_number, f"Expected '2' or 'two' in: {content[:100]}"

        report.add_result(name, True, f"Agent responded: {content[:50]}...", time.time() - start)
        return True
    except Exception as e:
        report.add_result(name, False, str(e), time.time() - start)
        return False


# ============================================================================
# TEST 4: Inter-Agent Communication
# ============================================================================


async def test_two_agents_chat(report: TestReport):
    """Test two agents communicating."""
    name = "Inter-Agent: Two agents chat"
    start = time.time()

    try:
        reset_runtime()
        bus = InMemoryEventBus()
        swarm = Swarm("chat-swarm", event_bus=bus)

        provider = await get_provider()
        alice = Agent("alice", model=provider)
        bob = Agent("bob", model=provider)

        swarm.subscribe("alice_speaks", alice, publish_to="bob_listens")
        swarm.subscribe("bob_listens", bob, publish_to="bob_speaks")

        bus.publish("alice_speaks", {"prompt": "Say hello to Bob"})

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

        events = bus.list_events(topic="bob_speaks")
        assert len(events) == 1

        content = events[0].payload.get("content", "")
        assert len(content) > 0

        report.add_result(name, True, f"Bob responded: {content[:80]}...", time.time() - start)
        return True
    except Exception as e:
        report.add_result(name, False, str(e), time.time() - start)
        return False


async def test_three_agent_pipeline(report: TestReport):
    """Test three-agent pipeline."""
    name = "Inter-Agent: Three-agent pipeline"
    start = time.time()

    try:
        reset_runtime()
        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", researcher, publish_to="research_done")
        swarm.subscribe("research_done", writer, publish_to="write_done")
        swarm.subscribe("write_done", editor, publish_to="final")

        bus.publish("start", {"prompt": "Research: What is AI?"})

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

        events = bus.list_events(topic="final")
        assert len(events) == 1

        content = events[0].payload.get("content", "")

        report.add_result(name, True, f"Pipeline output: {content[:100]}...", time.time() - start)
        return True
    except Exception as e:
        report.add_result(name, False, str(e), time.time() - start)
        return False


async def test_context_preservation(report: TestReport):
    """Test context preservation between agents."""
    name = "Inter-Agent: Context preservation"
    start = time.time()

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

        provider = await get_provider()
        collector = Agent("collector", model=provider)
        processor = Agent("processor", model=provider)

        swarm.subscribe("collect", collector, publish_to="processed")
        swarm.subscribe("processed", processor, publish_to="done")

        bus.publish(
            "collect",
            {
                "prompt": "List colors: red, blue, green",
                "user_id": "user123",
                "session": "sess_abc",
            },
        )

        await swarm.drain()

        # Check context is passed
        processed = bus.list_events(topic="processed")
        context_passed = "user_id" in processed[0].payload or "session" in processed[0].payload

        await swarm.drain()

        events = bus.list_events(topic="done")
        assert len(events) == 1

        report.add_result(name, True, f"Context preserved: {context_passed}", time.time() - start)
        return True
    except Exception as e:
        report.add_result(name, False, str(e), time.time() - start)
        return False


# ============================================================================
# TEST 5: Error Handling
# ============================================================================


async def test_swarm_error_isolation(report: TestReport):
    """Test that handler errors don't break other handlers."""
    name = "Swarm: Error isolation"
    start = time.time()

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

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

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

        swarm.subscribe("fail", fail_handler)
        swarm.subscribe("success", success_handler, publish_to="done")

        bus.publish("fail", {"prompt": "fail"})
        bus.publish("success", {"prompt": "ok"})

        processed = await swarm.drain(max_events=10)

        # Both events should be processed (failed one requeued)
        assert processed >= 1

        events = bus.list_events(topic="done")
        assert len(events) == 1

        report.add_result(
            name, True, f"Error isolated, success handler worked", time.time() - start
        )
        return True
    except Exception as e:
        report.add_result(name, False, str(e), time.time() - start)
        return False


# ============================================================================
# TEST 6: Event Bus Backends
# ============================================================================


async def test_sqlite_event_bus(report: TestReport):
    """Test SQLite event bus with swarm."""
    name = "Event Bus: SQLite backend"
    start = time.time()

    try:
        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("test", handler, publish_to="done")

        bus.publish("test", {"prompt": "test"})

        processed = await swarm.drain()

        events = bus.list_events(topic="done")
        assert len(events) == 1

        os.unlink(db_path)

        report.add_result(name, True, f"SQLite processed {processed} events", time.time() - start)
        return True
    except Exception as e:
        report.add_result(name, False, str(e), time.time() - start)
        return False


async def test_postgresql_event_bus(report: TestReport):
    """Test PostgreSQL event bus with swarm."""
    name = "Event Bus: PostgreSQL backend"
    start = time.time()

    try:
        import uuid

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

        bus = PostgreSQLEventBus(dsn, table_name=table_name)
        swarm = Swarm("pg-swarm", event_bus=bus)

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

        swarm.subscribe("test", handler, publish_to="done")

        bus.publish("test", {"prompt": "test"})

        processed = await swarm.drain()

        events = bus.list_events(topic="done")
        assert len(events) == 1

        report.add_result(
            name, True, f"PostgreSQL processed {processed} events", time.time() - start
        )
        return True
    except Exception as e:
        report.add_result(name, False, str(e), time.time() - start)
        return False


# ============================================================================
# TEST 7: Advanced Scenarios
# ============================================================================


async def test_parallel_agents(report: TestReport):
    """Test multiple agents processing in parallel."""
    name = "Swarm: Parallel agent processing"
    start = time.time()

    try:
        reset_runtime()
        bus = InMemoryEventBus()
        swarm = Swarm("parallel-swarm", event_bus=bus)

        provider = await get_provider()

        # Two agents listening to same topic
        agent1 = Agent("worker1", model=provider)
        agent2 = Agent("worker2", model=provider)

        swarm.subscribe("work", agent1, publish_to="done1")
        swarm.subscribe("work", agent2, publish_to="done2")

        # Publish multiple events
        for i in range(3):
            bus.publish("work", {"prompt": f"Task {i + 1}"})

        # Drain should process events
        processed = await swarm.drain()

        events1 = bus.list_events(topic="done1")
        events2 = bus.list_events(topic="done2")

        # Both agents should have processed some events
        total = len(events1) + len(events2)

        report.add_result(name, True, f"Total processed: {total} events", time.time() - start)
        return True
    except Exception as e:
        report.add_result(name, False, str(e), time.time() - start)
        return False


async def test_conditional_routing(report: TestReport):
    """Test conditional topic routing."""
    name = "Swarm: Conditional routing"
    start = time.time()

    try:
        reset_runtime()
        bus = InMemoryEventBus()
        swarm = Swarm("routing-swarm", event_bus=bus)

        def route_to_topic(event, payload):
            prompt = payload.get("prompt", "").lower()
            if "urgent" in prompt:
                return "urgent_queue"
            return "normal_queue"

        async def urgent_handler(event):
            return {"result": "URGENT handling"}

        async def normal_handler(event):
            return {"result": "normal handling"}

        # Use publish_to with callable for routing
        swarm.subscribe(
            "tasks", lambda e: {"prompt": e.payload.get("prompt")}, publish_to=route_to_topic
        )
        swarm.subscribe("urgent_queue", urgent_handler, publish_to="done")
        swarm.subscribe("normal_queue", normal_handler, publish_to="done")

        bus.publish("tasks", {"prompt": "This is urgent!"})
        bus.publish("tasks", {"prompt": "This is normal"})

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

        events = bus.list_events(topic="done")

        report.add_result(
            name, True, f"Conditional routing: {len(events)} events", time.time() - start
        )
        return True
    except Exception as e:
        report.add_result(name, False, str(e), time.time() - start)
        return False


async def test_event_state_tracking(report: TestReport):
    """Test event state transitions (queued → claimed → completed)."""
    name = "Event Bus: State transitions"
    start = time.time()

    try:
        bus = InMemoryEventBus()

        # Publish - should be queued
        event = bus.publish("test", {"data": "test"})
        assert event is not None, "Event should be created"
        e1 = bus.get_event(event.id)
        assert e1 is not None and e1.status == "queued", (
            f"Expected queued, got {e1.status if e1 else 'None'}"
        )

        # Claim - should become claimed
        claimed = bus.claim("worker", ["test"])
        assert claimed is not None, "Event should be claimable"
        assert claimed.status == "claimed", f"Expected claimed, got {claimed.status}"

        # Ack - should become completed
        bus.ack(claimed.id, "worker")
        completed = bus.get_event(claimed.id)
        assert completed is not None and completed.status == "completed", (
            f"Expected completed, got {completed.status if completed else 'None'}"
        )

        report.add_result(name, True, "State transitions work correctly", time.time() - start)
        return True
    except Exception as e:
        report.add_result(name, False, str(e), time.time() - start)
        return False


# ============================================================================
# MAIN TEST RUNNER
# ============================================================================


async def main():
    print("=" * 70)
    print("         LOGOS ADK - SWARM AGENTS COMPREHENSIVE TEST")
    print("=" * 70)
    print(f"Start time: {datetime.now().isoformat()}")
    print(f"API: {os.getenv('OPENAI_COMPAT_BASE_URL', 'https://api.example.com/v1')}")
    print(f"Model: {os.getenv('OPENAI_COMPAT_MODEL', 'gpt-4o-mini')}")
    print("=" * 70)

    report = TestReport()

    # Run all tests
    tests = [
        # Event Bus Tests
        test_event_bus_publish_claim,
        test_event_bus_filters,
        # Swarm Tests
        test_swarm_subscribe_and_drain,
        test_swarm_dispatch,
        # Agent Integration
        test_agent_in_swarm,
        # Inter-Agent Tests
        test_two_agents_chat,
        test_three_agent_pipeline,
        test_context_preservation,
        # Error Handling
        test_swarm_error_isolation,
        # Backend Tests
        test_sqlite_event_bus,
        test_postgresql_event_bus,
        # Advanced Tests
        test_parallel_agents,
        test_conditional_routing,
        test_event_state_tracking,
    ]

    for test in tests:
        await test(report)

    # Print summary
    success = report.print_summary()

    # Save JSON report
    report_file = "swarm_test_report.json"
    with open(report_file, "w") as f:
        json.dump(
            {
                "summary": {
                    "total": len(report.results),
                    "passed": sum(1 for r in report.results if r["passed"]),
                    "failed": sum(1 for r in report.results if not r["passed"]),
                    "duration": time.time() - report.start_time,
                    "timestamp": datetime.now().isoformat(),
                },
                "tests": report.results,
            },
            f,
            indent=2,
        )

    print(f"\nJSON report saved to: {report_file}")

    return success


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