"""Cryptographic operations for command signature verification."""

from __future__ import annotations

import base64
import contextlib
import hashlib
import hmac
from dataclasses import dataclass
from typing import Any

import structlog
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec, padding, rsa

from curadise_agent.errors import ValidationError

log = structlog.get_logger(__name__)


@dataclass
class SignatureInfo:
    """Information about a signature."""

    algorithm: str
    key_id: str | None
    signature: bytes
    timestamp: int | None = None


class SignatureVerifier:
    """
    Verifies cryptographic signatures on commands.

    Supports HMAC, RSA, and ECDSA signatures for command verification.
    """

    def __init__(self) -> None:
        """Initialize signature verifier."""
        self._public_keys: dict[str, Any] = {}
        self._hmac_keys: dict[str, bytes] = {}

    def add_public_key(self, key_id: str, key_pem: str | bytes) -> None:
        """
        Add a public key for verification.

        Args:
            key_id: Key identifier
            key_pem: PEM-encoded public key
        """
        if isinstance(key_pem, str):
            key_pem = key_pem.encode("utf-8")

        try:
            public_key = serialization.load_pem_public_key(key_pem)
            self._public_keys[key_id] = public_key
            log.info("public_key_added", key_id=key_id)
        except Exception as e:
            raise ValidationError(f"Invalid public key: {e}", cause=e) from e

    def add_hmac_key(self, key_id: str, key: bytes) -> None:
        """
        Add an HMAC key for verification.

        Args:
            key_id: Key identifier
            key: HMAC key bytes
        """
        self._hmac_keys[key_id] = key
        log.info("hmac_key_added", key_id=key_id)

    def remove_key(self, key_id: str) -> None:
        """Remove a key by ID."""
        self._public_keys.pop(key_id, None)
        self._hmac_keys.pop(key_id, None)

    def verify_hmac(
        self,
        data: bytes,
        signature: bytes,
        key_id: str | None = None,
        algorithm: str = "sha256",
    ) -> bool:
        """
        Verify HMAC signature.

        Args:
            data: Data that was signed
            signature: HMAC signature
            key_id: Key identifier (uses first key if None)
            algorithm: Hash algorithm (sha256, sha384, sha512)

        Returns:
            True if signature is valid
        """
        key = self._hmac_keys.get(key_id) if key_id else next(iter(self._hmac_keys.values()), None)

        if not key:
            log.error("hmac_key_not_found", key_id=key_id)
            return False

        hash_func = getattr(hashlib, algorithm, hashlib.sha256)
        expected = hmac.new(key, data, hash_func).digest()

        return hmac.compare_digest(signature, expected)

    def verify_rsa(
        self,
        data: bytes,
        signature: bytes,
        key_id: str | None = None,
    ) -> bool:
        """
        Verify RSA signature.

        Args:
            data: Data that was signed
            signature: RSA signature
            key_id: Key identifier

        Returns:
            True if signature is valid
        """
        if key_id:
            public_key = self._public_keys.get(key_id)
        else:
            public_key = next(
                (k for k in self._public_keys.values() if isinstance(k, rsa.RSAPublicKey)),
                None,
            )

        if not public_key or not isinstance(public_key, rsa.RSAPublicKey):
            log.error("rsa_key_not_found", key_id=key_id)
            return False

        try:
            public_key.verify(
                signature,
                data,
                padding.PKCS1v15(),
                hashes.SHA256(),
            )
            return True
        except InvalidSignature:
            return False

    def verify_ecdsa(
        self,
        data: bytes,
        signature: bytes,
        key_id: str | None = None,
    ) -> bool:
        """
        Verify ECDSA signature.

        Args:
            data: Data that was signed
            signature: ECDSA signature
            key_id: Key identifier

        Returns:
            True if signature is valid
        """
        if key_id:
            public_key = self._public_keys.get(key_id)
        else:
            public_key = next(
                (k for k in self._public_keys.values() if isinstance(k, ec.EllipticCurvePublicKey)),
                None,
            )

        if not public_key or not isinstance(public_key, ec.EllipticCurvePublicKey):
            log.error("ecdsa_key_not_found", key_id=key_id)
            return False

        try:
            public_key.verify(
                signature,
                data,
                ec.ECDSA(hashes.SHA256()),
            )
            return True
        except InvalidSignature:
            return False

    def verify(
        self,
        data: bytes,
        signature_info: SignatureInfo,
    ) -> bool:
        """
        Verify a signature using the appropriate algorithm.

        Args:
            data: Data that was signed
            signature_info: Signature information

        Returns:
            True if signature is valid
        """
        algorithm = signature_info.algorithm.lower()
        key_id = signature_info.key_id
        signature = signature_info.signature

        if algorithm.startswith("hmac"):
            hash_alg = algorithm.replace("hmac-", "").replace("hmac", "sha256")
            return self.verify_hmac(data, signature, key_id, hash_alg)
        elif algorithm in ("rsa", "rs256", "rsa-sha256"):
            return self.verify_rsa(data, signature, key_id)
        elif algorithm in ("ecdsa", "es256", "ecdsa-sha256"):
            return self.verify_ecdsa(data, signature, key_id)
        else:
            log.error("unsupported_signature_algorithm", algorithm=algorithm)
            return False


def parse_signature_header(header: str) -> SignatureInfo:
    """
    Parse a signature header.

    Expected format: algorithm=<alg>,keyId=<id>,signature=<base64>

    Args:
        header: Signature header string

    Returns:
        SignatureInfo

    Raises:
        ValidationError: If header is malformed
    """
    parts = {}
    for part in header.split(","):
        if "=" in part:
            key, value = part.split("=", 1)
            parts[key.strip()] = value.strip().strip('"')

    algorithm = parts.get("algorithm", "hmac-sha256")
    key_id = parts.get("keyId")
    signature_b64 = parts.get("signature")

    if not signature_b64:
        raise ValidationError("Missing signature in header")

    try:
        signature = base64.b64decode(signature_b64)
    except Exception as e:
        raise ValidationError(f"Invalid signature encoding: {e}", cause=e) from e

    timestamp = None
    if "timestamp" in parts:
        with contextlib.suppress(ValueError):
            timestamp = int(parts["timestamp"])

    return SignatureInfo(
        algorithm=algorithm,
        key_id=key_id,
        signature=signature,
        timestamp=timestamp,
    )


def compute_hmac(data: bytes, key: bytes, algorithm: str = "sha256") -> bytes:
    """
    Compute HMAC signature.

    Args:
        data: Data to sign
        key: HMAC key
        algorithm: Hash algorithm

    Returns:
        HMAC signature bytes
    """
    hash_func = getattr(hashlib, algorithm, hashlib.sha256)
    return hmac.new(key, data, hash_func).digest()
