"""
Universal Client for MCP Proxy Adapter Framework

This module provides a universal client that can connect to MCP Proxy Adapter servers
using various authentication methods and protocols. It's designed to be used for
proxy registration and general API communication.

Author: Vasiliy Zdanovskiy
email: vasilyvz@gmail.com
"""

import asyncio
import json
import os
import ssl
import time
from pathlib import Path
from typing import Any, Dict, Optional
from urllib.parse import urljoin

import aiohttp
import requests

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

    SECURITY_FRAMEWORK_AVAILABLE = True
except ImportError:
    SECURITY_FRAMEWORK_AVAILABLE = False
    # Define stubs for missing imports
    generate_api_key = None
    create_jwt_token = None
    validate_jwt_token = None
    extract_roles_from_cert = None
    validate_certificate_chain = None
    create_ssl_context = None
    validate_server_certificate = None


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[ssl.SSLContext]:
        """Create SSL context for secure connections."""
        ssl_config = self.security_config.get("ssl", {})
        if not ssl_config.get("enabled", False):
            return None

        try:
            context: Optional[ssl.SSLContext] = None

            # Try security framework first
            if self.security_manager:
                try:
                    context = self.security_manager.create_client_ssl_context()
                except Exception:
                    context = None

            # Fallback SSL context creation
            if context is None:
                context = ssl.create_default_context()

            # Always honor explicit client certificate config for mTLS
            cert_config = self.security_config.get("certificate", {})
            if cert_config.get("enabled", False):
                cert_file = cert_config.get("cert_file")
                key_file = cert_config.get("key_file")
                if cert_file and key_file:
                    context.load_cert_chain(certfile=cert_file, keyfile=key_file)

            # Add CA certificate if provided
            ca_cert_file = ssl_config.get("ca_cert_file") or ssl_config.get("ca_cert")
            if ca_cert_file and os.path.exists(ca_cert_file):
                context.load_verify_locations(cafile=ca_cert_file)

            # Configure verification
            if ssl_config.get("check_hostname", True):
                context.check_hostname = True
                context.verify_mode = ssl.CERT_REQUIRED
            else:
                context.check_hostname = False
                context.verify_mode = ssl.CERT_NONE

            return context
        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

        return results

    async def register_proxy(self, proxy_config: Dict[str, Any]) -> Dict[str, Any]:
        """
        Register with proxy server.

        Args:
            proxy_config: Proxy registration configuration

        Returns:
            Registration result
        """
        try:
            result = await self.post(
                "/api/jsonrpc",
                {
                    "jsonrpc": "2.0",
                    "method": "proxy_registration",
                    "params": proxy_config,
                    "id": 1,
                },
            )
            return result
        except Exception as e:
            print(f"Proxy registration failed: {e}")
            return {"error": str(e)}

    async def execute_command(
        self, command: str, params: Dict[str, Any] = None
    ) -> Dict[str, Any]:
        """
        Execute a command on the server.

        Args:
            command: Command name
            params: Command parameters

        Returns:
            Command result
        """
        try:
            result = await self.post(
                "/api/jsonrpc",
                {"jsonrpc": "2.0", "method": command, "params": params or {}, "id": 1},
            )
            return result
        except Exception as e:
            print(f"Command execution failed: {e}")
            return {"error": str(e)}


def create_client_from_config(config_file: str) -> UniversalClient:
    """
    Create a UniversalClient instance from a configuration file.

    Args:
        config_file: Path to configuration file

    Returns:
        UniversalClient instance
    """
    try:
        with open(config_file, "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}"

        # 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

        return UniversalClient(client_config)

    except Exception as e:
        raise ValueError(f"Failed to create client from config: {e}")


# CLI interface for standalone usage
async def main():
    """Main function for CLI usage."""
    import argparse

    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:
            client = create_client_from_config(args.config)

            print(f"🚀 Testing --config connection")
            print("=" * 40)
            print(f"Universal client initialized with auth method: --config")
            print(f"Connecting to {client.base_url} with --config authentication...")

            async with 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

                        result = await client.execute_command(args.method, params)
                        print(f"✅ Method '{args.method}' executed successfully:")
                        print(json.dumps(result, indent=2))
                    else:
                        # Default to help command
                        result = await client.execute_command("help")
                        print("✅ Help command executed successfully:")
                        print(json.dumps(result, indent=2))
                else:
                    print("❌ Connection failed")
                print("Connection closed")

        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}")
    else:
        print("❌ Configuration file required. Use --config option.")


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