#!/usr/bin/env python3
"""
Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com

Comprehensive test pipeline for all MCP Proxy Adapter modes with custom commands.

This test pipeline:
1. Tests all 8 security modes (HTTP, HTTPS, mTLS with various auth combinations)
2. Verifies custom commands registration in spawn mode
3. Tests queue commands execution in child processes
4. Validates that commands with use_queue=True work correctly

This is the main test that should be run before publishing to PyPI.
"""

import json
import multiprocessing
import os
import subprocess
import sys
import time
from pathlib import Path

import requests

# Set spawn mode for multiprocessing (required for CUDA compatibility)
multiprocessing.set_start_method("spawn", force=True)

# Test configurations
TEST_CONFIGS = [
    {
        "name": "HTTP Basic",
        "config": "mcp_proxy_adapter/examples/full_application/configs/http_basic.json",
        "port": 8080,
        "use_ssl": False,
        "use_mtls": False,
        "token": None,
    },
    {
        "name": "HTTP + Token",
        "config": "mcp_proxy_adapter/examples/full_application/configs/http_token.json",
        "port": 8080,
        "use_ssl": False,
        "use_mtls": False,
        "token": "admin-secret-key",
    },
    {
        "name": "HTTP + Token + Roles",
        "config": "mcp_proxy_adapter/examples/full_application/configs/http_token_roles.json",
        "port": 8080,
        "use_ssl": False,
        "use_mtls": False,
        "token": "admin-secret-key",
    },
    {
        "name": "HTTPS Basic",
        "config": "mcp_proxy_adapter/examples/full_application/configs/https_basic.json",
        "port": 8443,
        "use_ssl": True,
        "use_mtls": False,
        "token": None,
    },
    {
        "name": "HTTPS + Token",
        "config": "mcp_proxy_adapter/examples/full_application/configs/https_token.json",
        "port": 8443,
        "use_ssl": True,
        "use_mtls": False,
        "token": "admin-secret-key-https",
    },
    {
        "name": "HTTPS + Token + Roles",
        "config": "mcp_proxy_adapter/examples/full_application/configs/https_token_roles.json",
        "port": 8443,
        "use_ssl": True,
        "use_mtls": False,
        "token": "admin-secret-key-https",
    },
    {
        "name": "mTLS Basic",
        "config": "mcp_proxy_adapter/examples/full_application/configs/mtls_no_roles_correct.json",
        "port": 8443,
        "use_ssl": True,
        "use_mtls": True,
        "token": "admin-secret-key-mtls",
    },
    {
        "name": "mTLS + Roles",
        "config": "mcp_proxy_adapter/examples/full_application/configs/mtls_with_roles_correct.json",
        "port": 8443,
        "use_ssl": True,
        "use_mtls": True,
        "token": "admin-secret-key-mtls",
    },
]

TEST_RESULTS = []


def log_result(test_name: str, status: str, details: str = ""):
    """Log test result."""
    result = {"test": test_name, "status": status, "details": details}
    TEST_RESULTS.append(result)
    status_icon = "✅" if status == "PASS" else "❌" if status == "FAIL" else "⚠️"
    print(f"{status_icon} {test_name}: {status}")
    if details:
        print(f"   {details}")


def wait_for_server(url: str, timeout: int = 30, verify: bool = True) -> bool:
    """Wait for server to be ready."""
    start_time = time.time()
    while time.time() - start_time < timeout:
        try:
            response = requests.get(f"{url}/health", timeout=2, verify=verify)
            if response.status_code == 200:
                return True
        except requests.exceptions.RequestException:
            pass
        time.sleep(0.5)
    return False


def kill_server_on_port(port: int):
    """Kill any process running on the specified port."""
    try:
        result = subprocess.run(
            ["lsof", "-ti", f":{port}"], capture_output=True, text=True
        )
        if result.stdout.strip():
            pids = result.stdout.strip().split("\n")
            for pid in pids:
                try:
                    subprocess.run(["kill", "-9", pid], check=False)
                except Exception:
                    pass
    except Exception:
        pass


def test_health_endpoint(config: dict) -> bool:
    """Test health endpoint."""
    protocol = "https" if config["use_ssl"] else "http"
    url = f"{protocol}://localhost:{config['port']}"
    verify = not config["use_ssl"]  # Don't verify SSL for test certificates
    
    try:
        response = requests.get(f"{url}/health", timeout=5, verify=verify)
        if response.status_code == 200:
            data = response.json()
            if data.get("status") == "ok":
                return True
        return False
    except Exception as e:
        print(f"   Health check error: {e}")
        return False


def test_echo_command(config: dict) -> bool:
    """Test basic echo command."""
    protocol = "https" if config["use_ssl"] else "http"
    url = f"{protocol}://localhost:{config['port']}"
    verify = not config["use_ssl"]
    
    headers = {"Content-Type": "application/json"}
    if config["token"]:
        headers["X-API-Key"] = config["token"]
    
    payload = {
        "jsonrpc": "2.0",
        "method": "echo",
        "params": {"message": "Hello from comprehensive test"},
        "id": 1,
    }
    
    try:
        response = requests.post(
            f"{url}/api/jsonrpc", json=payload, headers=headers, timeout=10, verify=verify
        )
        if response.status_code == 200:
            result = response.json()
            if "result" in result and result["result"].get("success"):
                return True
        return False
    except Exception as e:
        print(f"   Echo command error: {e}")
        return False


def test_queue_command(config: dict) -> bool:
    """Test queue command (long_running_task) execution."""
    protocol = "https" if config["use_ssl"] else "http"
    url = f"{protocol}://localhost:{config['port']}"
    verify = not config["use_ssl"]
    
    headers = {"Content-Type": "application/json"}
    if config["token"]:
        headers["X-API-Key"] = config["token"]
    
    # Execute long_running_task (has use_queue=True)
    payload = {
        "jsonrpc": "2.0",
        "method": "long_running_task",
        "params": {"task_name": "comprehensive_test", "duration": 2, "steps": 4},
        "id": 1,
    }
    
    try:
        response = requests.post(
            f"{url}/api/jsonrpc", json=payload, headers=headers, timeout=10, verify=verify
        )
        
        if response.status_code != 200:
            return False
        
        result = response.json()
        
        # Check for errors
        if "error" in result:
            error = result["error"]
            error_code = error.get("code", 0)
            error_message = error.get("message", "")
            
            if error_code == -32601:  # Method not found
                print(f"   ❌ Command not found: {error_message}")
                print(f"   ❌ This indicates registration issue in spawn mode!")
                return False
            else:
                print(f"   ⚠️  Command error: {error_message}")
                return False
        
        if "result" not in result:
            return False
        
        command_result = result["result"]
        
        # Check if job was queued
        if isinstance(command_result, dict) and "job_id" in command_result:
            job_id = command_result["job_id"]
            
            # Wait for job to complete
            max_wait = 30
            waited = 0
            while waited < max_wait:
                status_payload = {
                    "jsonrpc": "2.0",
                    "method": "job_status",
                    "params": {"job_id": job_id},
                    "id": 2,
                }
                status_response = requests.post(
                    f"{url}/api/jsonrpc",
                    json=status_payload,
                    headers=headers,
                    timeout=5,
                    verify=verify,
                )
                
                if status_response.status_code == 200:
                    status_result = status_response.json()
                    if "result" in status_result:
                        job_status = status_result["result"]
                        status = job_status.get("status", "unknown")
                        
                        if status == "completed":
                            return True
                        elif status == "failed":
                            error = job_status.get("error", "Unknown error")
                            print(f"   ❌ Job failed: {error}")
                            return False
                
                time.sleep(1)
                waited += 1
            
            return False
        else:
            # Command executed synchronously (unexpected but OK)
            return True
            
    except Exception as e:
        print(f"   Queue command error: {e}")
        return False


def test_mode(config: dict) -> bool:
    """Test a single mode configuration."""
    print(f"\n{'='*70}")
    print(f"Testing: {config['name']}")
    print(f"{'='*70}")
    
    # Kill any existing server on this port
    kill_server_on_port(config["port"])
    time.sleep(1)
    
    # Start server
    config_path = Path(config["config"])
    if not config_path.exists():
        log_result(f"{config['name']} - Server start", "FAIL", f"Config file not found: {config_path}")
        return False
    
    protocol = "https" if config["use_ssl"] else "http"
    url = f"{protocol}://localhost:{config['port']}"
    verify = not config["use_ssl"]
    
    # Build command
    cmd = [
        sys.executable,
        "mcp_proxy_adapter/examples/full_application/main.py",
        "--config",
        str(config_path),
        "--port",
        str(config["port"]),
    ]
    
    print(f"Starting server: {' '.join(cmd)}")
    server_process = subprocess.Popen(
        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )
    
    try:
        # Wait for server to start
        print("Waiting for server to start...")
        if not wait_for_server(url, timeout=30, verify=verify):
            log_result(f"{config['name']} - Server start", "FAIL", "Server did not start")
            return False
        
        log_result(f"{config['name']} - Server start", "PASS", "Server started successfully")
        
        # Test 1: Health endpoint
        if test_health_endpoint(config):
            log_result(f"{config['name']} - Health endpoint", "PASS")
        else:
            log_result(f"{config['name']} - Health endpoint", "FAIL")
            return False
        
        # Test 2: Echo command
        if test_echo_command(config):
            log_result(f"{config['name']} - Echo command", "PASS")
        else:
            log_result(f"{config['name']} - Echo command", "FAIL")
            return False
        
        # Test 3: Queue command (CRITICAL - tests spawn mode registration)
        if test_queue_command(config):
            log_result(f"{config['name']} - Queue command", "PASS", "Command executed successfully in child process")
        else:
            log_result(f"{config['name']} - Queue command", "FAIL", "Command failed in child process - registration issue!")
            return False
        
        return True
        
    finally:
        # Stop server
        try:
            server_process.terminate()
            server_process.wait(timeout=5)
        except Exception:
            server_process.kill()
        kill_server_on_port(config["port"])
        time.sleep(1)


def main():
    """Run comprehensive test pipeline."""
    print("=" * 70)
    print("MCP Proxy Adapter - Comprehensive Test Pipeline")
    print("Testing all modes with custom commands in spawn mode")
    print("=" * 70)
    
    # Check if configs exist
    missing_configs = []
    for config in TEST_CONFIGS:
        if not Path(config["config"]).exists():
            missing_configs.append(config["config"])
    
    if missing_configs:
        print("❌ Missing configuration files:")
        for cfg in missing_configs:
            print(f"   - {cfg}")
        return 1
    
    # Run tests for each mode
    results = []
    for config in TEST_CONFIGS:
        try:
            success = test_mode(config)
            results.append((config["name"], success))
        except Exception as e:
            print(f"❌ Error testing {config['name']}: {e}")
            import traceback
            traceback.print_exc()
            results.append((config["name"], False))
    
    # Summary
    print("\n" + "=" * 70)
    print("Test Pipeline Summary")
    print("=" * 70)
    
    passed = sum(1 for _, success in results if success)
    failed = sum(1 for _, success in results if not success)
    
    for name, success in results:
        status = "✅ PASS" if success else "❌ FAIL"
        print(f"{status}: {name}")
    
    print(f"\nTotal: {passed} passed, {failed} failed out of {len(results)} modes")
    
    if failed > 0:
        print("\n⚠️  Some modes failed. Check the output above for details.")
        print("\n🔍 Key things to check:")
        print("   1. Are custom commands registered correctly?")
        print("   2. Do commands with use_queue=True work in spawn mode?")
        print("   3. Are modules auto-imported in child processes?")
        return 1
    
    print("\n🎉 All modes passed! Custom commands work correctly in spawn mode.")
    return 0


if __name__ == "__main__":
    sys.exit(main())
