#!/usr/bin/env python3
"""
Test client that connects to an MCP server over stdio and runs a search.
Usage:
  python test_client.py google       # tests google-mcp container
  python test_client.py duckduckgo   # tests duckduckgo-mcp container
"""

import asyncio
import subprocess
import sys


async def test_mcp_server(service_name: str):
    container_name = f"{service_name}-mcp"
    query = "Python programming language"
    print(f"\n{'='*60}")
    print(f"Testing {container_name} with query: '{query}'")
    print(f"{'='*60}")

    proc = subprocess.Popen(
        [
            "docker", "exec", "-i", container_name,
            "python", "-c",
            f"""
import sys
sys.path.insert(0, '/app')
import asyncio
from mcp.server.stdio import stdio_server
from {service_name.replace('-', '_')}_mcp.server import app

async def test():
    async with stdio_server() as (read, write):
        await asyncio.sleep(0.5)
        init_req = json.loads(await read.__anext__())
        # Send initialize response
        init_result = {{
            'jsonrpc': '2.0', 'id': init_req['id'],
            'result': {{
                'protocolVersion': '2024-11-05', 'capabilities': {{}},
                'serverInfo': {{'name': 'test', 'version': '1.0'}},
            }},
        }}
        await write.send(json.dumps(init_result))
        # Client sends initialized notification
        await read.__anext__()
        # List tools
        list_req = {{'jsonrpc': '2.0', 'id': 1, 'method': 'tools/list', 'params': {{}}}}
        await write.send(json.dumps(list_req))
        list_resp = json.loads(await read.__anext__())
        print("TOOLS:", json.dumps(list_resp.get('result', {{}}).get('tools', []), indent=2))
""",
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        stdin=subprocess.PIPE,
    )

    try:
        stdout, stderr = proc.communicate(timeout=30)
        print("STDOUT:", stdout.decode())
        if stderr:
            print("STDERR:", stderr.decode())
    except subprocess.TimeoutExpired:
        proc.kill()
        print("TIMEOUT: Test timed out after 30s")
        return False

    return proc.returncode == 0


async def test_with_curl(service_name: str):
    """Test using docker exec to run a Python test directly in the container."""
    container_name = f"{service_name}-mcp"

    test_script = f"""
import sys, asyncio, os
sys.path.insert(0, '/app')
os.environ.setdefault('RATE_LIMIT_SECONDS', '1')

from lib.playwright_manager import PlaywrightManager
from lib.google_search import GoogleSearch
from lib.duckduckgo_search import DuckDuckGoSearch

async def test():
    await PlaywrightManager.initialize()
    try:
        query = "Python programming language"
        if "{service_name}" == "google":
            engine = GoogleSearch(min_interval=1)
            results = await engine.search(query, num_results=5)
        else:
            engine = DuckDuckGoSearch(min_interval=1)
            results = await engine.search(query, num_results=5)

        print(f"Query: {{query}}")
        print(f"Results found: {{len(results)}}")
        for i, r in enumerate(results, 1):
            title = r.get('title', 'N/A')
            url = r.get('url', 'N/A')
            snippet = r.get('snippet', '')[:100]
            print(f"  {{i}}. {{title}}")
            print(f"     URL: {{url}}")
            if snippet:
                print(f"     {{snippet}}")

        if results:
            print("\\nSUCCESS: Got {{len(results)}} results")
            sys.exit(0)
        else:
            print("\\nFAILURE: No results returned")
            sys.exit(1)
    except Exception as e:
        print(f"ERROR: {{e}}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
    finally:
        await PlaywrightManager.shutdown()

asyncio.run(test())
"""

    print(f"\n{'='*60}")
    print(f"Testing {container_name}")
    print(f"{'='*60}")

    try:
        result = subprocess.run(
            [
                "docker", "exec", "-i", container_name,
                "python", "-c", test_script,
            ],
            capture_output=True,
            text=True,
            timeout=60,
        )
        print(result.stdout)
        if result.stderr:
            print("STDERR:", result.stderr)
        return result.returncode == 0

    except subprocess.TimeoutExpired:
        print("TIMEOUT: Test timed out after 60s")
        return False
    except FileNotFoundError:
        print("ERROR: docker not found")
        return False


async def main():
    if len(sys.argv) < 2:
        print("Usage: python test_client.py [google|duckduckgo|both]")
        sys.exit(1)

    target = sys.argv[1]
    services = ["google", "duckduckgo"] if target == "both" else [target]

    results = {}
    for service in services:
        success = await test_with_curl(service)
        results[service] = success

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

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


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