"""Module for verifying server key signatures with constant-time comparison."""

import logging
from typing import Dict, Optional

from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ed25519

logger = logging.getLogger(__name__)

# Ed25519 constants
ED25519_SIGNATURE_LENGTH = 64
ED25519_PUBLIC_KEY_LENGTH = 32
DEFAULT_KEY_ID = "default"


class SignatureVerifier:
    """Signature verifier with constant-time comparison."""

    def __init__(self) -> None:
        """Initialize the signature verifier."""
        self.trusted_public_keys = self._load_trusted_keys()

    def _load_trusted_keys(self) -> Dict[str, Optional[ed25519.Ed25519PublicKey]]:
        """
        Load trusted public keys.

        Returns:
            Dictionary mapping key IDs to public keys.
        """
        return {DEFAULT_KEY_ID: None}

    def verify_server_key_signature(
        self,
        public_key_pem: bytes,
        signature: bytes,
        kid: Optional[str] = None,
    ) -> bool:
        """
        Verify server public key signature with constant-time comparison.

        Args:
            public_key_pem: PEM-encoded public key.
            signature: Signature bytes.
            kid: Key ID for key selection (optional).

        Returns:
            True if signature is valid, False otherwise.
        """
        try:
            if len(signature) != ED25519_SIGNATURE_LENGTH:
                logger.warning(
                    "Invalid signature length",
                    extra={
                        "expected": ED25519_SIGNATURE_LENGTH,
                        "actual": len(signature),
                        "kid": kid,
                    },
                )
                return False

            try:
                public_key = serialization.load_pem_public_key(public_key_pem)
            except Exception as e:
                logger.warning(
                    "Failed to load PEM public key",
                    extra={"error": str(e), "kid": kid},
                )
                return False

            if not isinstance(public_key, ed25519.Ed25519PublicKey):
                logger.warning(
                    "Public key is not Ed25519",
                    extra={"kid": kid},
                )
                return False

            raw_public_key = public_key.public_bytes_raw()
            if len(raw_public_key) != ED25519_PUBLIC_KEY_LENGTH:
                logger.warning(
                    "Invalid public key length",
                    extra={
                        "expected": ED25519_PUBLIC_KEY_LENGTH,
                        "actual": len(raw_public_key),
                        "kid": kid,
                    },
                )
                return False

            trusted_key = self._get_trusted_key(kid)
            if not trusted_key:
                logger.error("No trusted key found", extra={"kid": kid})
                return False

            try:
                trusted_key.verify(signature, public_key_pem)
                logger.info("Signature verification successful", extra={"kid": kid})
                return True
            except InvalidSignature:
                logger.warning("Signature verification failed", extra={"kid": kid})
                return False

        except Exception as e:
            logger.error(
                "Signature verification error",
                extra={"error": str(e), "kid": kid},
            )
            return False

    def _get_trusted_key(
        self, kid: Optional[str] = None
    ) -> Optional[ed25519.Ed25519PublicKey]:
        """
        Get trusted public key by kid.

        Args:
            kid: Key ID (optional).

        Returns:
            Trusted public key or None if not found.
        """
        key_id = kid if kid else DEFAULT_KEY_ID
        return self.trusted_public_keys.get(key_id)

    def add_trusted_key(self, kid: str, public_key_pem: bytes) -> bool:
        """
        Add trusted public key.

        Args:
            kid: Key ID.
            public_key_pem: PEM-encoded public key.

        Returns:
            True if key was added successfully.
        """
        try:
            public_key = serialization.load_pem_public_key(public_key_pem)
            if not isinstance(public_key, ed25519.Ed25519PublicKey):
                logger.error("Public key is not Ed25519", extra={"kid": kid})
                return False

            self.trusted_public_keys[kid] = public_key
            logger.info("Trusted key added", extra={"kid": kid})
            return True
        except Exception as e:
            logger.error(
                "Failed to add trusted key",
                extra={"error": str(e), "kid": kid},
            )
            return False


signature_verifier = SignatureVerifier()
