"""
Module for PKI/CA certificate handling for server key signature verification.
"""

import logging
import ssl
from datetime import datetime, timezone
from typing import Dict, List, Optional

import requests
from cryptography import x509
from cryptography.hazmat.primitives.asymmetric import ed25519

logger = logging.getLogger(__name__)

class PKICertificateVerifier:
    """PKI/CA certificate verifier for server key signature verification."""

    def __init__(self, ca_cert_path: str, crl_urls: Optional[List[str]] = None) -> None:
        """Initialize PKI certificate verifier.

        Args:
            ca_cert_path: Path to CA certificate file.
            crl_urls: List of CRL URLs (optional).
        """
        self.ca_cert_path = ca_cert_path
        self.crl_urls = crl_urls or []
        self.ca_cert = self._load_ca_certificate()
        self.cert_store = self._build_cert_store()

    def _load_ca_certificate(self) -> x509.Certificate:
        """Load CA certificate.

        Returns:
            Loaded CA certificate.

        Raises:
            Exception: If certificate cannot be loaded.
        """
        try:
            with open(self.ca_cert_path, "rb") as f:
                ca_cert_data = f.read()
            return x509.load_pem_x509_certificate(ca_cert_data)
        except Exception as e:
            logger.error("Failed to load CA certificate", extra={"error": str(e)})
            raise

    def _build_cert_store(self) -> x509.CertificateStore:
        """Build certificate store.

        Returns:
            Certificate store with CA certificate.
        """
        store = x509.CertificateStore()
        store.add_cert(self.ca_cert)
        return store
    
    def verify_server_certificate(self, server_cert_pem: bytes) -> bool:
        """
        Verify server certificate through PKI/CA.

        Args:
            server_cert_pem: PEM-encoded server certificate.

        Returns:
            True if certificate is valid, False otherwise.
        """
        try:
            server_cert = x509.load_pem_x509_certificate(server_cert_pem)

            if not self._verify_certificate_chain(server_cert):
                logger.warning("Certificate chain verification failed")
                return False

            if not self._check_certificate_revocation(server_cert):
                logger.warning("Certificate is revoked")
                return False

            if not self._check_certificate_validity(server_cert):
                logger.warning("Certificate is expired or not yet valid")
                return False

            logger.info("Server certificate verified successfully")
            return True

        except Exception as e:
            logger.error("Certificate verification error", extra={"error": str(e)})
            return False
    
    def _verify_certificate_chain(self, server_cert: x509.Certificate) -> bool:
        """Verify certificate chain.

        Args:
            server_cert: Server certificate to verify.

        Returns:
            True if chain is valid, False otherwise.
        """
        try:
            ca_public_key = self.ca_cert.public_key()
            ca_public_key.verify(
                server_cert.signature,
                server_cert.tbs_certificate_bytes,
                server_cert.signature_algorithm_oid,
            )
            return True
        except Exception as e:
            logger.error(
                "Certificate chain verification failed", extra={"error": str(e)}
            )
            return False
    
    def _check_certificate_revocation(self, server_cert: x509.Certificate) -> bool:
        """Check certificate revocation via CRL.

        Args:
            server_cert: Server certificate to check.

        Returns:
            True if certificate is not revoked, False otherwise.
        """
        try:
            crl_dps = server_cert.extensions.get_extension_for_oid(
                x509.ExtensionOID.CRL_DISTRIBUTION_POINTS
            ).value

            for crl_dp in crl_dps:
                for crl_url in crl_dp.full_name:
                    if self._check_crl(server_cert, crl_url.value):
                        return False

            return True

        except Exception as e:
            logger.warning("CRL check failed", extra={"error": str(e)})
            return True
    
    def _check_crl(self, server_cert: x509.Certificate, crl_url: str) -> bool:
        """Check specific CRL for certificate revocation.

        Args:
            server_cert: Server certificate to check.
            crl_url: CRL URL.

        Returns:
            True if certificate is revoked, False otherwise.
        """
        try:
            response = requests.get(crl_url, timeout=10)
            if response.status_code == 200:
                crl_data = response.content
                crl = x509.load_der_x509_crl(crl_data)
                return server_cert.serial_number in [revoked.serial_number for revoked in crl]
            return False
        except Exception as e:
            logger.warning("Failed to check CRL", extra={"crl_url": crl_url, "error": str(e)})
            return False
    
    def _check_certificate_validity(self, server_cert: x509.Certificate) -> bool:
        """Check certificate validity period.

        Args:
            server_cert: Server certificate to check.

        Returns:
            True if certificate is valid, False otherwise.
        """
        now = datetime.now(timezone.utc)
        return server_cert.not_valid_before <= now <= server_cert.not_valid_after

    def extract_public_key_from_certificate(
        self, server_cert_pem: bytes
    ) -> ed25519.Ed25519PublicKey:
        """Extract public key from certificate.

        Args:
            server_cert_pem: PEM-encoded server certificate.

        Returns:
            Extracted Ed25519 public key.
        """
        server_cert = x509.load_pem_x509_certificate(server_cert_pem)
        public_key = server_cert.public_key()
        if not isinstance(public_key, ed25519.Ed25519PublicKey):
            raise ValueError("Certificate does not contain Ed25519 public key")
        return public_key

class ServerCertificateManager:
    """Server certificate manager."""

    def __init__(self, pki_verifier: PKICertificateVerifier) -> None:
        """Initialize server certificate manager.

        Args:
            pki_verifier: PKI certificate verifier instance.
        """
        self.pki_verifier = pki_verifier
        self.server_certificates: Dict[str, bytes] = {}

    def get_server_certificate(self, server_hostname: str) -> Optional[bytes]:
        """Get server certificate via TLS handshake.

        Args:
            server_hostname: Server hostname.

        Returns:
            PEM-encoded server certificate or None if failed.
        """
        try:
            context = ssl.create_default_context()
            context.check_hostname = True
            context.verify_mode = ssl.CERT_REQUIRED

            with ssl.create_connection((server_hostname, 443)) as sock:
                with context.wrap_socket(sock, server_hostname=server_hostname) as ssock:
                    cert_der = ssock.getpeercert_chain()[0]
                    cert_pem = ssl.DER_cert_to_PEM_cert(cert_der)
                    cert_bytes = cert_pem.encode()

                    if self.pki_verifier.verify_server_certificate(cert_bytes):
                        self.server_certificates[server_hostname] = cert_bytes
                        return cert_bytes
                    else:
                        logger.error(
                            "Server certificate verification failed",
                            extra={"server_hostname": server_hostname},
                        )
                        return None

        except Exception as e:
            logger.error(
                "Failed to get server certificate",
                extra={"server_hostname": server_hostname, "error": str(e)},
            )
            return None

    def get_trusted_public_key(
        self, server_hostname: str
    ) -> Optional[ed25519.Ed25519PublicKey]:
        """Get trusted public key from server certificate.

        Args:
            server_hostname: Server hostname.

        Returns:
            Ed25519 public key or None if failed.
        """
        cert_pem = self.get_server_certificate(server_hostname)
        if cert_pem:
            return self.pki_verifier.extract_public_key_from_certificate(cert_pem)
        return None


pki_verifier = PKICertificateVerifier(
    ca_cert_path="/etc/ssl/certs/ca-certificates.crt",
    crl_urls=["http://crl.example.com/crl.pem"],
)

certificate_manager = ServerCertificateManager(pki_verifier)
