"""
Cryptographic signature verification utilities for EDS authentication.

Handles X.509 certificate validation and nonce-based challenge-response 
authentication using raw signature verification (RSA, ECDSA, DSA).

Security Note: This module works exclusively with public certificates and keys.
NO private keys are stored or processed on the server side.
"""

import binascii
import hashlib
from datetime import datetime
from datetime import timezone as python_timezone
from typing import Optional, Tuple

from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import dsa, ec, padding, rsa
from cryptography.x509.oid import ExtensionOID


class CertificateValidationError(Exception):
    """Raised when certificate validation fails."""

    pass


class SignatureVerificationError(Exception):
    """Raised when signature verification fails."""

    pass


class EDSCertificateValidator:
    """
    Validates X.509 certificates for EDS authentication.

    Performs checks:
    - Certificate format and PEM encoding
    - Validity period
    - Certificate chain (when available)
    - Key usage constraints
    """

    def __init__(self, ca_certs_pem: Optional[str] = None):
        """
        Initialize validator.

        Args:
            ca_certs_pem: PEM-encoded CA certificates for chain validation (optional)
        """
        self.ca_certs_pem = ca_certs_pem
        self.ca_certs: list = []

        if ca_certs_pem:
            self._load_ca_certificates()

    def _load_ca_certificates(self) -> None:
        """Load and parse CA certificates."""
        if not self.ca_certs_pem:
            return

        # Split multiple PEM blocks
        pem_blocks = self.ca_certs_pem.split("-----END CERTIFICATE-----")

        for block in pem_blocks:
            if "-----BEGIN CERTIFICATE-----" not in block:
                continue

            pem_block = block.strip() + "\n-----END CERTIFICATE-----"
            try:
                cert_data = pem_block.encode("utf-8")
                cert = x509.load_pem_x509_certificate(cert_data)
                self.ca_certs.append(cert)
            except Exception as e:
                raise CertificateValidationError(f"Failed to load CA certificate: {e}")

    def load_certificate(self, pem_data: str) -> x509.Certificate:
        """
        Load certificate from PEM format.

        Args:
            pem_data: PEM-encoded certificate

        Returns:
            Loaded x509.Certificate object

        Raises:
            CertificateValidationError: If certificate cannot be loaded
        """
        try:
            cert_bytes = pem_data.encode("utf-8")
            cert = x509.load_pem_x509_certificate(cert_bytes)
            return cert
        except Exception as e:
            raise CertificateValidationError(f"Invalid certificate format: {e}")

    def validate_certificate(self, cert: x509.Certificate) -> Tuple[bool, str]:
        """
        Validate certificate.

        Args:
            cert: x509.Certificate to validate

        Returns:
            Tuple of (is_valid: bool, reason: str)
        """
        # Check validity period
        now = datetime.now(tz=python_timezone.utc)

        if now < cert.not_valid_before_utc:
            return False, "Certificate not yet valid"

        if now > cert.not_valid_after_utc:
            return False, "Certificate expired"

        # Check key usage extension
        try:
            key_usage = cert.extensions.get_extension_for_oid(
                ExtensionOID.KEY_USAGE
            ).value

            if not key_usage.digital_signature:
                return False, "Certificate not authorized for digital signatures"
        except x509.ExtensionNotFound:
            # Key usage not present, continue
            pass

        # Validate certificate chain if CA certs provided
        if self.ca_certs:
            chain_valid, chain_reason = self._validate_chain(cert)
            if not chain_valid:
                return False, chain_reason

        return True, "Certificate is valid"

    def _validate_chain(self, cert: x509.Certificate) -> Tuple[bool, str]:
        """Validate certificate chain against CA certificates."""
        # Check if certificate is self-signed
        if cert.issuer == cert.subject:
            # For self-signed, check if it's in our CA list
            for ca_cert in self.ca_certs:
                if ca_cert.public_key().public_bytes(
                    encoding=serialization.Encoding.DER,
                    format=serialization.PublicFormat.SubjectPublicKeyInfo,
                ) == cert.public_key().public_bytes(
                    encoding=serialization.Encoding.DER,
                    format=serialization.PublicFormat.SubjectPublicKeyInfo,
                ):
                    return True, "Self-signed certificate is trusted"
            return False, "Self-signed certificate is not trusted"

        # Find issuer certificate
        issuer_cert = None
        for ca_cert in self.ca_certs:
            if ca_cert.subject == cert.issuer:
                issuer_cert = ca_cert
                break

        if not issuer_cert:
            return False, "Issuer certificate not found in trusted pool"

        # Verify signature - handle different key types
        try:
            issuer_public_key = issuer_cert.public_key()

            if isinstance(issuer_public_key, rsa.RSAPublicKey):
                issuer_public_key.verify(
                    cert.signature,
                    cert.tbs_certificate_bytes,
                    padding.PKCS1v15(),
                    cert.signature_hash_algorithm,
                )
            elif isinstance(issuer_public_key, ec.EllipticCurvePublicKey):
                issuer_public_key.verify(
                    cert.signature,
                    cert.tbs_certificate_bytes,
                    ec.ECDSA(cert.signature_hash_algorithm),
                )
            elif isinstance(issuer_public_key, dsa.DSAPublicKey):
                issuer_public_key.verify(
                    cert.signature,
                    cert.tbs_certificate_bytes,
                    cert.signature_hash_algorithm,
                )
            else:
                return False, f"Unsupported issuer key type: {type(issuer_public_key)}"

            return True, "Certificate chain is valid"
        except Exception as e:
            return False, f"Certificate signature verification failed: {e}"

    def extract_certificate_info(self, cert: x509.Certificate) -> dict:
        """
        Extract certificate metadata.

        Args:
            cert: x509.Certificate to process

        Returns:
            Dictionary with certificate information
        """
        return {
            "subject": cert.subject.rfc4514_string(),
            "issuer": cert.issuer.rfc4514_string(),
            "serial_number": format(cert.serial_number, "x"),
            "valid_from": cert.not_valid_before_utc,
            "valid_to": cert.not_valid_after_utc,
            "thumbprint": hashlib.sha256(
                cert.public_bytes(serialization.Encoding.DER)
            ).hexdigest(),
        }


class EDSSignatureVerifier:
    """
    Verifies digital signatures over nonce values.

    Supports multiple signature formats:
    - Raw signatures (RSA, EC, DSA) in hex-encoded format
    """

    def __init__(self, certificate_pem: str):
        """
        Initialize verifier with a certificate.

        Args:
            certificate_pem: PEM-encoded certificate for verification
        """
        self.cert_pem = certificate_pem
        self.cert: Optional[x509.Certificate] = None
        self.public_key = None

        self._load_certificate()

    def _load_certificate(self) -> None:
        """Load and extract public key from certificate."""
        try:
            cert_bytes = self.cert_pem.encode("utf-8")
            self.cert = x509.load_pem_x509_certificate(cert_bytes)
            self.public_key = self.cert.public_key()
        except Exception as e:
            raise SignatureVerificationError(f"Failed to load certificate: {e}")

    def verify_signature(
        self, nonce: str, signature_hex: str, algorithm: str = "sha256"
    ) -> Tuple[bool, str]:
        """
        Verify signature over nonce.

        Args:
            nonce: Original nonce that was signed
            signature_hex: Hex-encoded signature value
            algorithm: Hash algorithm used ('sha256', 'sha384', 'sha512')

        Returns:
            Tuple of (is_valid: bool, reason: str)
        """
        try:
            nonce_bytes = nonce.encode("utf-8")
            signature_bytes = binascii.unhexlify(signature_hex)
        except (ValueError, binascii.Error) as e:
            return False, f"Invalid signature format: {e}"

        # Select hash algorithm
        hash_algo = self._get_hash_algorithm(algorithm)
        if not hash_algo:
            return False, f"Unsupported hash algorithm: {algorithm}"

        # Determine key type and verify accordingly
        try:
            if isinstance(self.public_key, rsa.RSAPublicKey):
                return self._verify_rsa_signature(
                    nonce_bytes, signature_bytes, hash_algo
                )
            elif isinstance(self.public_key, ec.EllipticCurvePublicKey):
                return self._verify_ec_signature(
                    nonce_bytes, signature_bytes, hash_algo
                )
            elif isinstance(self.public_key, dsa.DSAPublicKey):
                return self._verify_dsa_signature(
                    nonce_bytes, signature_bytes, hash_algo
                )
            else:
                return False, f"Unsupported key type: {type(self.public_key)}"
        except Exception as e:
            return False, f"Signature verification error: {e}"

    def _get_hash_algorithm(self, algorithm: str):
        """Get cryptography hash algorithm object."""
        algorithms = {
            "sha256": hashes.SHA256(),
            "sha384": hashes.SHA384(),
            "sha512": hashes.SHA512(),
        }
        return algorithms.get(algorithm)

    def _verify_rsa_signature(
        self, data: bytes, signature: bytes, hash_algo
    ) -> Tuple[bool, str]:
        """Verify RSA signature."""
        try:
            self.public_key.verify(signature, data, padding.PKCS1v15(), hash_algo)
            return True, "RSA signature verified"
        except Exception:
            return False, "RSA signature verification failed"

    def _verify_ec_signature(
        self, data: bytes, signature: bytes, hash_algo
    ) -> Tuple[bool, str]:
        """Verify ECDSA signature."""
        try:
            self.public_key.verify(signature, data, ec.ECDSA(hash_algo))
            return True, "ECDSA signature verified"
        except Exception:
            return False, "ECDSA signature verification failed"

    def _verify_dsa_signature(
        self, data: bytes, signature: bytes, hash_algo
    ) -> Tuple[bool, str]:
        """Verify DSA signature."""
        try:
            self.public_key.verify(signature, data, hash_algo)
            return True, "DSA signature verified"
        except Exception as e:
            return False, f"DSA signature verification failed: {e}"
