"""
Universal Client Example
This module demonstrates all possible secure connection methods to the server
using the mcp_security_framework. The client supports all authentication methods
and connection types supported by the security framework.
Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com
"""

import asyncio
import json
import os
import time
from pathlib import Path
from typing import TYPE_CHECKING, Optional, Dict, Any
import aiohttp
import requests

# Add project root to path
import sys
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from mcp_proxy_adapter.core.ssl_utils import SSLUtils

if TYPE_CHECKING:
    from ssl import SSLContext

# Import security framework components
try:
    from mcp_security_framework import (
        SecurityManager,
        AuthManager,
        CertificateManager,
        PermissionManager,
    )
    from mcp_security_framework.utils.auth_utils import (
        generate_api_key,
        create_jwt_token,
        validate_jwt_token,
    )
    from mcp_security_framework.utils.cert_utils import (
        extract_roles_from_cert,
        validate_certificate_chain,
    )
    from mcp_security_framework.utils.ssl_utils import (
        create_ssl_context,
        validate_server_certificate,
    )

    SECURITY_FRAMEWORK_AVAILABLE = True
except ImportError:
    SECURITY_FRAMEWORK_AVAILABLE = False
    print("Warning: mcp_security_framework not available. Using basic HTTP client.")


class UniversalClient:
    """
    Universal client that demonstrates all possible secure connection methods.
    Supports:
    - HTTP/HTTPS connections
    - API Key authentication
    - JWT token authentication
    - Certificate-based authentication
    - SSL/TLS with custom certificates
    - Role-based access control
    - Rate limiting awareness
    """

    def __init__(self, config: Dict[str, Any]):
        """
        Initialize universal client with configuration.
        Args:
            config: Client configuration with security settings
        """
        self.config = config
        self.base_url = config.get("server_url", "http://localhost:8000")
        self.timeout = config.get("timeout", 30)
        self.retry_attempts = config.get("retry_attempts", 3)
        self.retry_delay = config.get("retry_delay", 1)
        # Security configuration
        self.security_config = config.get("security", {})
        self.auth_method = self.security_config.get("auth_method", "none")
        # Initialize security managers if framework is available
        self.security_manager = None
        self.auth_manager = None
        self.cert_manager = None
        if SECURITY_FRAMEWORK_AVAILABLE:
            self._initialize_security_managers()
        # Session management
        self.session: Optional[aiohttp.ClientSession] = None
        self.current_token: Optional[str] = None
        self.token_expiry: Optional[float] = None
        print(f"Universal client initialized with auth method: {self.auth_method}")

    def _initialize_security_managers(self) -> None:
        """Initialize security framework managers."""
        try:
            # Initialize security manager
            self.security_manager = SecurityManager(self.security_config)
            # Initialize permission manager first
            permissions_config = self.security_config.get("permissions", {})
            self.permission_manager = PermissionManager(permissions_config)
            # Initialize auth manager with permission_manager
            auth_config = self.security_config.get("auth", {})
            self.auth_manager = AuthManager(auth_config, self.permission_manager)
            # Initialize certificate manager
            cert_config = self.security_config.get("certificates", {})
            self.cert_manager = CertificateManager(cert_config)
            print("Security framework managers initialized successfully")
        except Exception as e:
            print(f"Warning: Failed to initialize security managers: {e}")

    async def __aenter__(self):
        """Async context manager entry."""
        await self.connect()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Async context manager exit."""
        await self.disconnect()

    async def connect(self) -> None:
        """Establish connection with authentication."""
        print(
            f"Connecting to {self.base_url} with {self.auth_method} authentication..."
        )
        # Create SSL context
        ssl_context = self._create_ssl_context()
        # Create connector with SSL context
        connector = None
        if ssl_context:
            connector = aiohttp.TCPConnector(ssl=ssl_context)
        # Create session
        self.session = aiohttp.ClientSession(connector=connector)
        # Perform authentication based on method
        if self.auth_method == "api_key":
            await self._authenticate_api_key()
        elif self.auth_method == "jwt":
            await self._authenticate_jwt()
        elif self.auth_method == "certificate":
            await self._authenticate_certificate()
        elif self.auth_method == "basic":
            await self._authenticate_basic()
        else:
            print("No authentication required")
        print("Connection established successfully")

    async def disconnect(self) -> None:
        """Close connection and cleanup."""
        if self.session:
            await self.session.close()
            self.session = None
        print("Connection closed")

    async def _authenticate_api_key(self) -> None:
        """Authenticate using API key."""
        api_key_config = self.security_config.get("api_key", {})
        api_key = api_key_config.get("key")
        if not api_key:
            raise ValueError("API key not provided in configuration")
        # Store API key for requests
        self.current_token = api_key
        print(f"Authenticated with API key: {api_key[:8]}...")

    async def _authenticate_jwt(self) -> None:
        """Authenticate using JWT token."""
        jwt_config = self.security_config.get("jwt", {})
        # Check if we have a stored token that's still valid
        if self.current_token and self.token_expiry and time.time() < self.token_expiry:
            print("Using existing JWT token")
            return
        # Get credentials for JWT
        username = jwt_config.get("username")
        password = jwt_config.get("password")
        secret = jwt_config.get("secret")
        if not all([username, password, secret]):
            raise ValueError("JWT credentials not provided in configuration")
        # Create JWT token
        if SECURITY_FRAMEWORK_AVAILABLE:
            self.current_token = create_jwt_token(
                username, secret, expiry_hours=jwt_config.get("expiry_hours", 24)
            )
        else:
            # Simple JWT creation (for demonstration)
            import jwt

            payload = {
                "username": username,
                "exp": time.time() + (jwt_config.get("expiry_hours", 24) * 3600),
            }
            self.current_token = jwt.encode(payload, secret, algorithm="HS256")
        self.token_expiry = time.time() + (jwt_config.get("expiry_hours", 24) * 3600)
        print(f"Authenticated with JWT token: {self.current_token[:20]}...")

    async def _authenticate_certificate(self) -> None:
        """Authenticate using client certificate."""
        cert_config = self.security_config.get("certificate", {})
        cert_file = cert_config.get("cert_file")
        key_file = cert_config.get("key_file")
        if not cert_file or not key_file:
            raise ValueError("Certificate files not provided in configuration")
        # Validate certificate
        if SECURITY_FRAMEWORK_AVAILABLE and self.cert_manager:
            try:
                cert_info = self.cert_manager.validate_certificate(cert_file, key_file)
                print(f"Certificate validated: {cert_info.get('subject', 'Unknown')}")
                # Extract roles from certificate
                roles = extract_roles_from_cert(cert_file)
                if roles:
                    print(f"Certificate roles: {roles}")
            except Exception as e:
                print(f"Warning: Certificate validation failed: {e}")
        print("Certificate authentication prepared")

    async def _authenticate_basic(self) -> None:
        """Authenticate using basic authentication."""
        basic_config = self.security_config.get("basic", {})
        username = basic_config.get("username")
        password = basic_config.get("password")
        if not username or not password:
            raise ValueError("Basic auth credentials not provided in configuration")
        import base64

        credentials = base64.b64encode(f"{username}:{password}".encode()).decode()
        self.current_token = f"Basic {credentials}"
        print(f"Authenticated with basic auth: {username}")

    def _get_auth_headers(self) -> Dict[str, str]:
        """Get authentication headers for requests."""
        headers = {"Content-Type": "application/json"}
        if not self.current_token:
            return headers
        if self.auth_method == "api_key":
            api_key_config = self.security_config.get("api_key", {})
            header_name = api_key_config.get("header", "X-API-Key")
            headers[header_name] = self.current_token
        elif self.auth_method == "jwt":
            headers["Authorization"] = f"Bearer {self.current_token}"
        elif self.auth_method == "basic":
            headers["Authorization"] = self.current_token
        return headers

    def _create_ssl_context(self) -> Optional["SSLContext"]:
        """Create SSL context for secure connections using security framework."""
        ssl_config = self.security_config.get("ssl", {})
        if not ssl_config.get("enabled", False):
            return None
        try:
            # Get certificate configuration
            cert_config = self.security_config.get("certificate", {})
            cert_file = cert_config.get("cert_file") if cert_config.get("enabled", False) else None
            key_file = cert_config.get("key_file") if cert_config.get("enabled", False) else None
            
            # Get CA certificate
            ca_cert_file = ssl_config.get("ca_cert_file") or ssl_config.get("ca_cert")
            
            # Get verification settings
            verify = ssl_config.get("verify", True)
            check_hostname = ssl_config.get("check_hostname", True)
            
            # Create SSL context using security framework
            # Try security manager first if available
            if self.security_manager and hasattr(self.security_manager, 'create_client_ssl_context'):
                try:
                    return self.security_manager.create_client_ssl_context()
                except Exception:
                    pass  # Fall through to SSLUtils
            
            # Use SSLUtils from mcp_proxy_adapter (always uses security framework)
            return SSLUtils.create_client_ssl_context(
                ca_cert=ca_cert_file,
                client_cert=cert_file,
                client_key=key_file,
                verify=verify,
                check_hostname=check_hostname,
            )
        except Exception as e:
            print(f"Warning: Failed to create SSL context: {e}")
            return None

    async def request(
        self,
        method: str,
        endpoint: str,
        data: Optional[Dict[str, Any]] = None,
        headers: Optional[Dict[str, str]] = None,
    ) -> Dict[str, Any]:
        """
        Make authenticated request to server.
        Args:
            method: HTTP method (GET, POST, etc.)
            endpoint: API endpoint
            data: Request data
            headers: Additional headers
        Returns:
            Response data
        """
        url = urljoin(self.base_url, endpoint)
        # Prepare headers
        request_headers = self._get_auth_headers()
        if headers:
            request_headers.update(headers)
        try:
            for attempt in range(self.retry_attempts):
                try:
                    async with self.session.request(
                        method,
                        url,
                        json=data,
                        headers=request_headers,
                        timeout=aiohttp.ClientTimeout(total=self.timeout),
                    ) as response:
                        result = await response.json()
                        # Validate response if security framework available
                        if SECURITY_FRAMEWORK_AVAILABLE and self.security_manager:
                            self.security_manager.validate_server_response(
                                dict(response.headers)
                            )
                        if response.status >= 400:
                            print(
                                f"Request failed with status {response.status}: {result}"
                            )
                            return {"error": result, "status": response.status}
                        return result
                except Exception as e:
                    print(f"Request attempt {attempt + 1} failed: {e}")
                    if attempt < self.retry_attempts - 1:
                        await asyncio.sleep(self.retry_delay)
                    else:
                        raise
        except Exception as e:
            print(f"Request failed: {e}")
            raise

    async def get(self, endpoint: str, **kwargs) -> Dict[str, Any]:
        """Make GET request."""
        return await self.request("GET", endpoint, **kwargs)

    async def post(
        self, endpoint: str, data: Dict[str, Any], **kwargs
    ) -> Dict[str, Any]:
        """Make POST request."""
        return await self.request("POST", endpoint, data=data, **kwargs)

    async def put(
        self, endpoint: str, data: Dict[str, Any], **kwargs
    ) -> Dict[str, Any]:
        """Make PUT request."""
        return await self.request("PUT", endpoint, data=data, **kwargs)

    async def delete(self, endpoint: str, **kwargs) -> Dict[str, Any]:
        """Make DELETE request."""
        return await self.request("DELETE", endpoint, **kwargs)

    async def test_connection(self) -> bool:
        """Test connection to server."""
        try:
            result = await self.get("/health")
            if "error" not in result:
                print("✅ Connection test successful")
                return True
            else:
                print(f"❌ Connection test failed: {result}")
                return False
        except Exception as e:
            print(f"❌ Connection test failed: {e}")
            return False

    async def test_security_features(self) -> Dict[str, bool]:
        """Test various security features."""
        results = {}
        # Test basic connectivity
        results["connectivity"] = await self.test_connection()
        # Test authentication
        if self.auth_method != "none":
            try:
                result = await self.get("/api/auth/status")
                results["authentication"] = "error" not in result
            except:
                results["authentication"] = False
        # Test SSL/TLS
        if self.base_url.startswith("https"):
            results["ssl_tls"] = True
        else:
            results["ssl_tls"] = False
        # Test certificate validation
        if self.auth_method == "certificate" and SECURITY_FRAMEWORK_AVAILABLE:
            results["certificate_validation"] = True
        else:
            results["certificate_validation"] = False
        return results


def create_client_config(
    server_url: str, auth_method: str = "none", **kwargs
) -> Dict[str, Any]:
    """
    Create client configuration for different authentication methods.
    Args:
        server_url: Server URL
        auth_method: Authentication method (none, api_key, jwt, certificate, basic)
        **kwargs: Additional configuration parameters
    Returns:
        Client configuration dictionary
    """
    config = {
        "server_url": server_url,
        "timeout": 30,
        "retry_attempts": 3,
        "retry_delay": 1,
        "security": {"auth_method": auth_method},
    }
    if auth_method == "api_key":
        config["security"]["api_key"] = {
            "key": kwargs.get("api_key", "your_api_key_here"),
            "header": kwargs.get("header", "X-API-Key"),
        }
    elif auth_method == "jwt":
        config["security"]["jwt"] = {
            "username": kwargs.get("username", "user"),
            "password": kwargs.get("password", "password"),
            "secret": kwargs.get("secret", "your_jwt_secret"),
            "expiry_hours": kwargs.get("expiry_hours", 24),
        }
    elif auth_method == "certificate":
        config["security"]["certificate"] = {
            "enabled": True,
            "cert_file": kwargs.get("cert_file", "./certs/client.crt"),
            "key_file": kwargs.get("key_file", "./keys/client.key"),
            "ca_cert_file": kwargs.get("ca_cert_file", "./certs/ca.crt"),
        }
        config["security"]["ssl"] = {
            "enabled": True,
            "check_hostname": kwargs.get("check_hostname", True),
            "ca_cert_file": kwargs.get("ca_cert_file", "./certs/ca.crt"),
        }
    elif auth_method == "basic":
        config["security"]["basic"] = {
            "username": kwargs.get("username", "user"),
            "password": kwargs.get("password", "password"),
        }
    return config


async def demo_all_connection_methods():
    """Demonstrate all possible connection methods."""
    print("🚀 Universal Client Demo - All Connection Methods")
    print("=" * 60)
    # Test configurations for different auth methods
    test_configs = [
        {
            "name": "No Authentication",
            "config": create_client_config("http://localhost:8000", "none"),
        },
        {
            "name": "API Key Authentication",
            "config": create_client_config(
                "http://localhost:8000", "api_key", api_key="demo_api_key_123"
            ),
        },
        {
            "name": "JWT Authentication",
            "config": create_client_config(
                "http://localhost:8000",
                "jwt",
                username="demo_user",
                password="demo_password",
                secret="demo_jwt_secret",
            ),
        },
        {
            "name": "Basic Authentication",
            "config": create_client_config(
                "http://localhost:8000",
                "basic",
                username="demo_user",
                password="demo_password",
            ),
        },
        {
            "name": "Certificate Authentication (HTTPS)",
            "config": create_client_config(
                "https://localhost:8443",
                "certificate",
                cert_file="./certs/client.crt",
                key_file="./keys/client.key",
                ca_cert_file="./certs/ca.crt",
            ),
        },
    ]
    for test_config in test_configs:
        print(f"\n📋 Testing: {test_config['name']}")
        print("-" * 40)
        try:
            async with UniversalClient(test_config["config"]) as client:
                # Test connection
                success = await client.test_connection()
                if success:
                    # Test security features
                    security_results = await client.test_security_features()
                    print("Security Features:")
                    for feature, status in security_results.items():
                        status_icon = "✅" if status else "❌"
                        print(f"  {status_icon} {feature}: {status}")
                    # Make a test API call
                    try:
                        result = await client.get("/api/status")
                        print(f"API Status: {result}")
                    except Exception as e:
                        print(f"API call failed: {e}")
                else:
                    print("❌ Connection failed")
        except Exception as e:
            print(f"❌ Test failed: {e}")
    print("\n🎉 Demo completed!")


async def demo_specific_connection(auth_method: str, **kwargs):
    """
    Demo specific connection method.
    Args:
        auth_method: Authentication method to test
        **kwargs: Configuration parameters
    """
    print(f"🚀 Testing {auth_method} connection")
    print("=" * 40)
    config = create_client_config("http://localhost:8000", auth_method, **kwargs)
    async with UniversalClient(config) as client:
        # Test connection
        success = await client.test_connection()
        if success:
            print("✅ Connection successful!")
            # Make some API calls
            try:
                # Get server status
                status = await client.get("/api/status")
                print(f"Server Status: {status}")
                # Test command execution
                command_data = {
                    "jsonrpc": "2.0",
                    "method": "test_command",
                    "params": {"message": "Hello from universal client!"},
                    "id": 1,
                }
                result = await client.post("/api/jsonrpc", command_data)
                print(f"Command Result: {result}")
            except Exception as e:
                print(f"API calls failed: {e}")
        else:
            print("❌ Connection failed")


if __name__ == "__main__":
    import sys
    import argparse
    import json

    parser = argparse.ArgumentParser(
        description="Universal Client for MCP Proxy Adapter"
    )
    parser.add_argument("--config", help="Path to configuration file")
    parser.add_argument("--method", help="JSON-RPC method to call")
    parser.add_argument("--params", help="JSON-RPC parameters (JSON string)")
    parser.add_argument("--auth-method", help="Authentication method")
    parser.add_argument("--server-url", help="Server URL")
    args = parser.parse_args()
    if args.config:
        # Load configuration from file
        try:
            with open(args.config, "r") as f:
                config_data = json.load(f)
            # Extract server configuration
            server_config = config_data.get("server", {})
            host = server_config.get("host", "127.0.0.1")
            port = server_config.get("port", 8000)
            # Determine protocol
            ssl_config = config_data.get("ssl", {})
            ssl_enabled = ssl_config.get("enabled", False)
            protocol = "https" if ssl_enabled else "http"
            server_url = f"{protocol}://{host}:{port}"
            print(f"🚀 Testing --config connection")
            print("=" * 40)
            print(f"Universal client initialized with auth method: --config")
            print(f"Connecting to {server_url} with --config authentication...")
            # Create client configuration
            client_config = {
                "server_url": server_url,
                "timeout": 30,
                "retry_attempts": 3,
                "retry_delay": 1,
                "security": {"auth_method": "none"},
            }
            # Add SSL configuration if needed
            if ssl_enabled:
                client_config["security"]["ssl"] = {
                    "enabled": True,
                    "check_hostname": False,
                    "verify": False,
                }
                # Add CA certificate if available
                ca_cert = ssl_config.get("ca_cert")
                if ca_cert and os.path.exists(ca_cert):
                    client_config["security"]["ssl"]["ca_cert_file"] = ca_cert

            async def test_config_connection():
                async with UniversalClient(client_config) as client:
                    # Test connection
                    success = await client.test_connection()
                    if success:
                        print("No authentication required")
                        print("Connection established successfully")
                        if args.method:
                            # Execute JSON-RPC method
                            params = {}
                            if args.params:
                                try:
                                    params = json.loads(args.params)
                                except json.JSONDecodeError:
                                    print("❌ Invalid JSON parameters")
                                    return
                            command_data = {
                                "jsonrpc": "2.0",
                                "method": args.method,
                                "params": params,
                                "id": 1,
                            }
                            try:
                                result = await client.post("/api/jsonrpc", command_data)
                                print(
                                    f"✅ Method '{args.method}' executed successfully:"
                                )
                                print(json.dumps(result, indent=2))
                            except Exception as e:
                                print(f"❌ Method execution failed: {e}")
                        else:
                            # Default to help command
                            command_data = {
                                "jsonrpc": "2.0",
                                "method": "help",
                                "params": {},
                                "id": 1,
                            }
                            try:
                                result = await client.post("/api/jsonrpc", command_data)
                                print("✅ Help command executed successfully:")
                                print(json.dumps(result, indent=2))
                            except Exception as e:
                                print(f"❌ Help command failed: {e}")
                    else:
                        print("❌ Connection failed")
                    print("Connection closed")

            asyncio.run(test_config_connection())
        except FileNotFoundError:
            print(f"❌ Configuration file not found: {args.config}")
        except json.JSONDecodeError:
            print(f"❌ Invalid JSON in configuration file: {args.config}")
        except Exception as e:
            print(f"❌ Error loading configuration: {e}")
    elif len(sys.argv) > 1:
        # Test specific auth method (legacy mode)
        auth_method = sys.argv[1]
        kwargs = {}
        if auth_method == "api_key":
            kwargs["api_key"] = "demo_key_123"
        elif auth_method == "jwt":
            kwargs.update(
                {
                    "username": "demo_user",
                    "password": "demo_password",
                    "secret": "demo_secret",
                }
            )
        elif auth_method == "certificate":
            kwargs.update(
                {
                    "cert_file": "./certs/client.crt",
                    "key_file": "./keys/client.key",
                    "ca_cert_file": "./certs/ca.crt",
                }
            )
        elif auth_method == "basic":
            kwargs.update({"username": "demo_user", "password": "demo_password"})
        asyncio.run(demo_specific_connection(auth_method, **kwargs))
    else:
        # Demo all connection methods
        asyncio.run(demo_all_connection_methods())
