#!/usr/bin/env python3
"""Integration test: connect to MCP servers via SSE and test search."""
import asyncio
import json
import sys
import aiohttp


class SSEClient:
    def __init__(self):
        self.events = []
        self.lock = asyncio.Lock()

    async def add_event(self, data):
        async with self.lock:
            self.events.append(data)

    async def wait_for_id(self, expected_id, timeout=30):
        start = asyncio.get_event_loop().time()
        while asyncio.get_event_loop().time() - start < timeout:
            async with self.lock:
                for evt in self.events:
                    if isinstance(evt, dict) and evt.get("id") == expected_id:
                        return evt
            await asyncio.sleep(0.1)
        return None


async def test_server(name, host, port, tool_name, query):
    print(f"\n{'='*60}")
    print(f"Testing {name} at http://{host}:{port}")
    print(f"{'='*60}")

    base = f"http://{host}:{port}"
    collector = SSEClient()
    messages_endpoint = None

    async with aiohttp.ClientSession() as session:
        async with session.get(f"{base}/sse") as sse_resp:
            if sse_resp.status != 200:
                print(f"FAIL: SSE status {sse_resp.status}")
                return False

            async def read_sse():
                try:
                    async for line in sse_resp.content:
                        line_str = line.decode().strip()
                        if not line_str or not line_str.startswith("data:"):
                            continue
                        data_str = line_str[5:].strip()
                        try:
                            data = json.loads(data_str)
                            await collector.add_event(data)
                        except json.JSONDecodeError:
                            nonlocal messages_endpoint
                            messages_endpoint = base.rstrip("/") + data_str
                            print(f"    Endpoint: {messages_endpoint}")
                            await collector.add_event(data_str)
                except Exception as e:
                    print(f"    SSE error: {e}")

            reader = asyncio.create_task(read_sse())

            try:
                for _ in range(50):
                    if messages_endpoint:
                        break
                    await asyncio.sleep(0.1)

                if not messages_endpoint:
                    print("FAIL: No endpoint")
                    return False

                # Initialize
                async with session.post(messages_endpoint, json={
                    "jsonrpc": "2.0", "id": 1, "method": "initialize",
                    "params": {"protocolVersion": "2024-11-05", "capabilities": {},
                               "clientInfo": {"name": "test", "version": "1.0"}},
                }) as resp:
                    pass

                init = await collector.wait_for_id(1, timeout=10)
                if not init:
                    print("FAIL: No init response")
                    return False
                print(f"    Protocol: {init['result']['protocolVersion']}")

                # Initialized notification
                async with session.post(messages_endpoint, json={
                    "jsonrpc": "2.0", "method": "notifications/initialized",
                }) as resp:
                    pass

                # List tools
                async with session.post(messages_endpoint, json={
                    "jsonrpc": "2.0", "id": 2, "method": "tools/list", "params": {},
                }) as resp:
                    pass

                tools = await collector.wait_for_id(2, timeout=10)
                if tools:
                    print(f"    Tools: {[t['name'] for t in tools.get('result', {}).get('tools', [])]}")

                # Call search
                print(f"    Searching: '{query}'")
                async with session.post(messages_endpoint, json={
                    "jsonrpc": "2.0", "id": 3, "method": "tools/call",
                    "params": {"name": tool_name, "arguments": {"query": query, "num_results": 5}},
                }) as resp:
                    pass

                result = await collector.wait_for_id(3, timeout=30)
                if not result:
                    print("FAIL: No search response")
                    return False

                content = result.get("result", {}).get("content", [])
                if content:
                    text = content[0].get("text", "")
                    print(f"\n    --- RESULTS ---\n{text}\n    --- END ---")
                    if "Search failed" in text:
                        print("    PARTIAL: Error in response")
                        return False
                    if "No results" in text:
                        print("    PARTIAL: No results")
                        return False
                    return True

                print("FAIL: Empty response")
                return False

            finally:
                reader.cancel()
                try:
                    await reader
                except asyncio.CancelledError:
                    pass


async def main():
    host = "localhost"
    results = {}

    # Wait for services to be ready
    print("Waiting for services to start...")
    await asyncio.sleep(5)

    results["google"] = await test_server("Google", host, 3001, "google_search", "Python programming language")

    await asyncio.sleep(2)

    results["duckduckgo"] = await test_server("DuckDuckGo", host, 3002, "duckduckgo_search", "Python programming language")

    print(f"\n{'='*60}")
    print("SUMMARY")
    print(f"{'='*60}")
    for name, ok in results.items():
        print(f"  {name}: {'PASS' if ok else 'FAIL'}")

    sys.exit(0 if all(results.values()) else 1)


if __name__ == "__main__":
    asyncio.run(main())