"""
Universal Authentication Validator

This module provides a universal authentication validator that supports
certificate, token, mTLS, and SSL validation with JSON-RPC error codes.

Author: MCP Proxy Adapter Team
Version: 1.0.0
"""

import json
import logging
import os
from datetime import datetime
from pathlib import Path
from typing import Optional, Dict, Any, List

from cryptography import x509


# Standard JSON-RPC error codes
JSON_RPC_ERRORS = {
    # General errors
    -32600: "Invalid Request",  # Invalid request
    -32601: "Method not found",  # Method not found
    -32602: "Invalid params",  # Invalid parameters
    -32603: "Internal error",  # Internal error
    -32700: "Parse error",  # Parse error
    # Custom codes for authentication
    -32001: "Authentication disabled",  # Authentication disabled
    -32002: "Invalid configuration",  # Invalid configuration
    -32003: "Certificate validation failed",  # Certificate validation failed
    -32004: "Token validation failed",  # Token validation failed
    -32005: "MTLS validation failed",  # MTLS validation failed
    -32006: "SSL validation failed",  # SSL validation failed
    -32007: "Role validation failed",  # Role validation failed
    -32008: "Certificate expired",  # Certificate expired
    -32009: "Certificate not found",  # Certificate not found
    -32010: "Token expired",  # Token expired
    -32011: "Token not found",  # Token not found
}


class AuthValidationResult:
    """
    Authentication validation result.

    Contains validation status, error information, and extracted roles.
    """

    def __init__(
        self,
        is_valid: bool,
        error_code: Optional[int] = None,
        error_message: Optional[str] = None,
        roles: Optional[List[str]] = None,
    ):
        """
        Initialize authentication validation result.

        Args:
            is_valid: Whether authentication is valid
            error_code: JSON-RPC error code if validation failed
            error_message: Error message if validation failed
            roles: List of roles extracted from authentication data
        """
        self.is_valid = is_valid
        self.error_code = error_code
        self.error_message = error_message
        self.roles = roles or []




class AuthValidator:
    """
    Universal authentication validator.

    Provides methods to validate different types of authentication:
    - Certificate validation
    - Token validation
    - mTLS validation
    - SSL validation
    """

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        """
        Initialize authentication validator.

        Args:
            config: Configuration dictionary for validation settings
        """
        self.config = config or {}
        self.logger = logging.getLogger(__name__)

        # Custom OID for roles
        self.role_oid = "1.3.6.1.4.1.99999.1"


    def validate_certificate(
        self, cert_path: Optional[str], cert_type: str = "server"
    ) -> AuthValidationResult:
        """
        Validate certificate.

        Args:
            cert_path: Path to certificate file
            cert_type: Type of certificate (server/client/ca)

        Returns:
            Certificate validation result
        """
        try:
            if not cert_path:
                return AuthValidationResult(
                    is_valid=False,
                    error_code=-32009,
                    error_message="Certificate path not provided",
                )

            if not os.path.exists(cert_path):
                return AuthValidationResult(
                    is_valid=False,
                    error_code=-32009,
                    error_message=f"Certificate file not found: {cert_path}",
                )

            # Load and validate certificate
            with open(cert_path, "rb") as f:
                cert_data = f.read()

            cert = x509.load_pem_x509_certificate(cert_data)

            # Check if certificate is not expired
            now = datetime.utcnow()
            if cert.not_valid_before > now or cert.not_valid_after < now:
                return AuthValidationResult(
                    is_valid=False,
                    error_code=-32008,
                    error_message="Certificate has expired",
                )

            # Extract roles from certificate
            roles = self._extract_roles_from_certificate(cert)

            # Validate certificate type
            if cert_type == "server":
                # Check for server-specific extensions
                if not self._validate_server_certificate(cert):
                    return AuthValidationResult(
                        is_valid=False,
                        error_code=-32003,
                        error_message="Invalid server certificate",
                    )
            elif cert_type == "client":
                # Check for client-specific extensions
                if not self._validate_client_certificate(cert):
                    return AuthValidationResult(
                        is_valid=False,
                        error_code=-32003,
                        error_message="Invalid client certificate",
                    )

            return AuthValidationResult(is_valid=True, roles=roles)

        except Exception as e:
            self.get_global_logger().error(f"Certificate validation error: {e}")
            return AuthValidationResult(
                is_valid=False,
                error_code=-32003,
                error_message=f"Certificate validation failed: {str(e)}",
            )

    def validate_token(
        self, token: Optional[str], token_type: str = "jwt"
    ) -> AuthValidationResult:
        """
        Validate token.

        Args:
            token: Token string to validate
            token_type: Type of token (jwt/api)

        Returns:
            Token validation result
        """
        try:
            if not token:
                return AuthValidationResult(
                    is_valid=False,
                    error_code=-32011,
                    error_message="Token not provided",
                )

            if token_type == "jwt":
                return self._validate_jwt_token(token)
            elif token_type == "api":
                return self._validate_api_token(token)
            else:
                return AuthValidationResult(
                    is_valid=False,
                    error_code=-32602,
                    error_message=f"Unsupported token type: {token_type}",
                )

        except Exception as e:
            self.get_global_logger().error(f"Token validation error: {e}")
            return AuthValidationResult(
                is_valid=False,
                error_code=-32004,
                error_message=f"Token validation failed: {str(e)}",
            )

    def validate_mtls(
        self, client_cert: Optional[str], ca_cert: Optional[str]
    ) -> AuthValidationResult:
        """
        Validate mTLS connection.

        Args:
            client_cert: Path to client certificate
            ca_cert: Path to CA certificate

        Returns:
            mTLS validation result
        """
        try:
            if not client_cert or not ca_cert:
                return AuthValidationResult(
                    is_valid=False,
                    error_code=-32005,
                    error_message="Client certificate and CA certificate required for mTLS",
                )

            # Validate client certificate
            client_result = self.validate_certificate(client_cert, "client")
            if not client_result.is_valid:
                return client_result

            # Validate CA certificate
            ca_result = self.validate_certificate(ca_cert, "ca")
            if not ca_result.is_valid:
                return ca_result

            # Verify client certificate is signed by CA
            if not self._verify_certificate_chain(client_cert, ca_cert):
                return AuthValidationResult(
                    is_valid=False,
                    error_code=-32005,
                    error_message="Client certificate not signed by provided CA",
                )

            return AuthValidationResult(is_valid=True, roles=client_result.roles)

        except Exception as e:
            self.get_global_logger().error(f"mTLS validation error: {e}")
            return AuthValidationResult(
                is_valid=False,
                error_code=-32005,
                error_message=f"mTLS validation failed: {str(e)}",
            )

    def validate_ssl(self, server_cert: Optional[str]) -> AuthValidationResult:
        """
        Validate SSL connection.

        Args:
            server_cert: Path to server certificate

        Returns:
            SSL validation result
        """
        try:
            if not server_cert:
                return AuthValidationResult(
                    is_valid=False,
                    error_code=-32006,
                    error_message="Server certificate required for SSL validation",
                )

            # Validate server certificate
            return self.validate_certificate(server_cert, "server")

        except Exception as e:
            self.get_global_logger().error(f"SSL validation error: {e}")
            return AuthValidationResult(
                is_valid=False,
                error_code=-32006,
                error_message=f"SSL validation failed: {str(e)}",
            )

    def _check_config(self, auth_type: str) -> bool:
        """
        Check if authentication is enabled in configuration.

        Args:
            auth_type: Type of authentication

        Returns:
            True if authentication is enabled, False otherwise
        """
        ssl_config = self.config.get("ssl", {})

        if not ssl_config.get("enabled", False):
            return False

        # Check specific authentication type
        if auth_type == "token":
            return ssl_config.get("token_auth", {}).get("enabled", False)
        elif auth_type == "mtls":
            return ssl_config.get("mtls", {}).get("enabled", False)
        elif auth_type in ["ssl", "certificate"]:
            return True

        # For unsupported types, return False to indicate authentication is not enabled
        return False

    def _get_validation_mode(self) -> str:
        """
        Get validation mode from configuration.

        Returns:
            Validation mode string
        """
        ssl_config = self.config.get("ssl", {})

        if ssl_config.get("token_auth", {}).get("enabled", False):
            return "token"
        elif ssl_config.get("mtls", {}).get("enabled", False):
            return "mtls"
        elif ssl_config.get("enabled", False):
            return "ssl"

        return "none"

    def _extract_roles_from_certificate(self, cert: x509.Certificate) -> List[str]:
        """
        Extract roles from certificate.

        Args:
            cert: Certificate object

        Returns:
            List of roles extracted from certificate
        """
        try:
            for extension in cert.extensions:
                if extension.oid.dotted_string == self.role_oid:
                    roles_data = extension.value.value.decode("utf-8")
                    return [role.strip() for role in roles_data.split(",")]

            return []

        except Exception as e:
            self.get_global_logger().error(f"Failed to extract roles from certificate: {e}")
            return []

    def _validate_server_certificate(self, cert: x509.Certificate) -> bool:
        """
        Validate server certificate extensions.

        Args:
            cert: Certificate object

        Returns:
            True if server certificate is valid, False otherwise
        """
        try:
            # Check for server authentication key usage
            for extension in cert.extensions:
                if extension.oid == x509.oid.ExtensionOID.KEY_USAGE:
                    key_usage = extension.value
                    return key_usage.digital_signature and key_usage.key_encipherment

            return True

        except Exception as e:
            self.get_global_logger().error(f"Server certificate validation error: {e}")
            return False

    def _validate_client_certificate(self, cert: x509.Certificate) -> bool:
        """
        Validate client certificate extensions.

        Args:
            cert: Certificate object

        Returns:
            True if client certificate is valid, False otherwise
        """
        try:
            # Check for client authentication extended key usage
            for extension in cert.extensions:
                if extension.oid == x509.oid.ExtensionOID.EXTENDED_KEY_USAGE:
                    extended_key_usage = extension.value
                    return (
                        x509.oid.ExtendedKeyUsageOID.CLIENT_AUTH in extended_key_usage
                    )

            return True

        except Exception as e:
            self.get_global_logger().error(f"Client certificate validation error: {e}")
            return False

    def _validate_jwt_token(self, token: str) -> AuthValidationResult:
        """
        Validate JWT token.

        Args:
            token: JWT token string

        Returns:
            JWT validation result
        """
        try:
            # This is a placeholder for JWT validation
            # In a real implementation, you would use a JWT library
            # like PyJWT to validate the token

            # For now, just check if token exists and has basic format
            if not token or len(token.split(".")) != 3:
                return AuthValidationResult(
                    is_valid=False,
                    error_code=-32004,
                    error_message="Invalid JWT token format",
                )

            # Extract roles from JWT payload (placeholder)
            roles = []

            return AuthValidationResult(is_valid=True, roles=roles)

        except Exception as e:
            return AuthValidationResult(
                is_valid=False,
                error_code=-32004,
                error_message=f"JWT validation failed: {str(e)}",
            )

    def _validate_api_token(self, token: str) -> AuthValidationResult:
        """
        Validate API token.

        Args:
            token: API token string

        Returns:
            API token validation result
        """
        try:
            # This is a placeholder for API token validation
            # In a real implementation, you would validate against
            # a token store or database

            if not token:
                return AuthValidationResult(
                    is_valid=False,
                    error_code=-32011,
                    error_message="API token not found",
                )

            # Extract roles from API token (placeholder)
            roles = []

            return AuthValidationResult(is_valid=True, roles=roles)

        except Exception as e:
            return AuthValidationResult(
                is_valid=False,
                error_code=-32004,
                error_message=f"API token validation failed: {str(e)}",
            )

    def _verify_certificate_chain(
        self, client_cert_path: str, ca_cert_path: str
    ) -> bool:
        """
        Verify that client certificate is signed by CA.

        Args:
            client_cert_path: Path to client certificate
            ca_cert_path: Path to CA certificate

        Returns:
            True if certificate chain is valid, False otherwise
        """
        try:
            # Load client certificate
            with open(client_cert_path, "rb") as f:
                client_cert_data = f.read()
            client_cert = x509.load_pem_x509_certificate(client_cert_data)

            # Load CA certificate
            with open(ca_cert_path, "rb") as f:
                ca_cert_data = f.read()
            ca_cert = x509.load_pem_x509_certificate(ca_cert_data)

            # Verify client certificate is signed by CA
            ca_public_key = ca_cert.public_key()

            # This is a simplified verification
            # In a real implementation, you would use proper signature verification
            return True

        except Exception as e:
            self.get_global_logger().error(f"Certificate chain verification error: {e}")
            return False
