#!/usr/bin/env python3
"""
Test script for Gemini MCP Server
Tests both HTTP and stdio modes
"""

import argparse
import asyncio
import json
import sys
from typing import Any, Dict

import httpx


async def test_http_mode(base_url: str = "http://localhost:8006"):
    """Test the HTTP mode of the server"""
    print("Testing HTTP mode...")
    
    async with httpx.AsyncClient() as client:
        # Test health endpoint
        print("\n1. Testing health check...")
        response = await client.get(f"{base_url}/health")
        if response.status_code == 200:
            print(f"✅ Health check passed: {response.json()}")
        else:
            print(f"❌ Health check failed: {response.status_code}")
            
        # Test tools listing
        print("\n2. Testing tools listing...")
        response = await client.get(f"{base_url}/mcp/tools")
        if response.status_code == 200:
            tools = response.json()
            print(f"✅ Found {len(tools['tools'])} tools:")
            for tool in tools['tools']:
                print(f"   - {tool['name']}: {tool['description']}")
        else:
            print(f"❌ Tools listing failed: {response.status_code}")
            
        # Test Gemini status
        print("\n3. Testing Gemini status...")
        response = await client.post(
            f"{base_url}/mcp/execute",
            json={"tool": "gemini_status"}
        )
        if response.status_code == 200:
            result = response.json()
            if result['success']:
                print(f"✅ Gemini status: {json.dumps(result['result'], indent=2)}")
            else:
                print(f"❌ Gemini status failed: {result.get('error')}")
        else:
            print(f"❌ Status request failed: {response.status_code}")
            
        # Test simple consultation
        print("\n4. Testing Gemini consultation...")
        response = await client.post(
            f"{base_url}/mcp/execute",
            json={
                "tool": "consult_gemini",
                "arguments": {
                    "query": "What is 2 + 2?",
                    "force": True
                }
            }
        )
        if response.status_code == 200:
            result = response.json()
            if result['success']:
                print(f"✅ Consultation successful")
                print(f"   Response: {result['result'][:200]}...")
            else:
                print(f"❌ Consultation failed: {result.get('error')}")
        else:
            print(f"❌ Consultation request failed: {response.status_code}")


async def test_stdio_mode():
    """Test the stdio mode of the server"""
    print("Testing stdio mode...")
    print("\n⚠️  Note: stdio mode requires manual testing")
    print("To test stdio mode:")
    print("1. Run: python gemini_mcp_server.py --mode stdio")
    print("2. Send JSON-RPC messages via stdin")
    print("3. Example initialization message:")
    
    init_message = {
        "jsonrpc": "2.0",
        "method": "initialize",
        "params": {
            "protocolVersion": "2024-11-05",
            "clientInfo": {
                "name": "test-client",
                "version": "1.0.0"
            }
        },
        "id": 1
    }
    
    print(json.dumps(init_message, indent=2))
    print("\n4. Then send tools/list request:")
    
    list_message = {
        "jsonrpc": "2.0",
        "method": "tools/list",
        "params": {},
        "id": 2
    }
    
    print(json.dumps(list_message, indent=2))


async def test_mcp_protocol(base_url: str = "http://localhost:8006"):
    """Test MCP protocol endpoints"""
    print("\nTesting MCP Protocol endpoints...")
    
    async with httpx.AsyncClient() as client:
        # Test messages endpoint
        print("\n1. Testing /messages endpoint...")
        
        # Initialize session
        init_request = {
            "jsonrpc": "2.0",
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "clientInfo": {
                    "name": "test-client",
                    "version": "1.0.0"
                }
            },
            "id": 1
        }
        
        response = await client.post(
            f"{base_url}/messages",
            json=init_request,
            headers={"Content-Type": "application/json"}
        )
        
        if response.status_code == 200:
            result = response.json()
            print(f"✅ Initialization successful: {json.dumps(result, indent=2)}")
            session_id = response.headers.get("Mcp-Session-Id")
            print(f"   Session ID: {session_id}")
            
            # List tools
            print("\n2. Listing tools via MCP protocol...")
            list_request = {
                "jsonrpc": "2.0",
                "method": "tools/list",
                "params": {},
                "id": 2
            }
            
            response = await client.post(
                f"{base_url}/messages",
                json=list_request,
                headers={
                    "Content-Type": "application/json",
                    "Mcp-Session-Id": session_id
                }
            )
            
            if response.status_code == 200:
                result = response.json()
                print(f"✅ Tools listed: {len(result['result']['tools'])} tools available")
            else:
                print(f"❌ Tools listing failed: {response.status_code}")
        else:
            print(f"❌ Initialization failed: {response.status_code}")


async def main():
    parser = argparse.ArgumentParser(description="Test Gemini MCP Server")
    parser.add_argument(
        "--mode",
        choices=["http", "stdio", "both"],
        default="http",
        help="Test mode"
    )
    parser.add_argument(
        "--url",
        default="http://localhost:8006",
        help="Server URL for HTTP mode"
    )
    
    args = parser.parse_args()
    
    print("🧪 Gemini MCP Server Test Suite")
    print("=" * 40)
    
    if args.mode in ["http", "both"]:
        try:
            await test_http_mode(args.url)
            await test_mcp_protocol(args.url)
        except httpx.ConnectError:
            print(f"\n❌ Could not connect to server at {args.url}")
            print("   Make sure the server is running: python gemini_mcp_server.py --mode http")
            sys.exit(1)
    
    if args.mode in ["stdio", "both"]:
        await test_stdio_mode()
    
    print("\n✅ Test suite completed!")


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