"""This module implements Elliptic Curve Diffie-Hellman based Key Management algorithms."""

from typing import Any, Type, Union

from binapy import BinaPy
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import ec, x448, x25519
from cryptography.hazmat.primitives.kdf.concatkdf import ConcatKDFHash

from ..base import BaseAsymmetricAlg, BaseKeyManagementAlg
from .aeskw import A128KW, A192KW, A256KW, BaseAesKeyWrap


class EcdhEs(
    BaseKeyManagementAlg,
    BaseAsymmetricAlg[
        Union[ec.EllipticCurvePrivateKey, x25519.X25519PrivateKey, x448.X448PrivateKey],
        Union[ec.EllipticCurvePublicKey, x25519.X25519PublicKey, x448.X448PublicKey],
    ],
):
    """Elliptic Curve Diffie-Hellman Ephemeral Static key agreement using Concat KDF."""

    name = "ECDH-ES"
    description = __doc__
    public_key_class = (
        ec.EllipticCurvePublicKey,
        x25519.X25519PublicKey,
        x448.X448PublicKey,
    )
    private_key_class = (
        ec.EllipticCurvePrivateKey,
        x25519.X25519PrivateKey,
        x448.X448PrivateKey,
    )

    @classmethod
    def otherinfo(cls, alg: str, apu: bytes, apv: bytes, keysize: int) -> BinaPy:
        """Build the "otherinfo" parameter for Concat KDF Hash.

        Args:
          alg: identifier for the encryption alg
          apu: Agreement PartyUInfo
          apv: Agreement PartyVInfo
          keysize: length of the generated key

        Returns:
            the "otherinfo" value
        """
        algorithm_id = BinaPy.from_int(len(alg), length=4) + BinaPy(alg)
        partyuinfo = BinaPy.from_int(len(apu), length=4) + apu
        partyvinfo = BinaPy.from_int(len(apv), length=4) + apv
        supppubinfo = BinaPy.from_int(keysize or keysize, length=4)
        otherinfo = b"".join((algorithm_id, partyuinfo, partyvinfo, supppubinfo))
        return BinaPy(otherinfo)

    @classmethod
    def ecdh(
        cls,
        private_key: Union[
            ec.EllipticCurvePrivateKey, x25519.X25519PrivateKey, x448.X448PrivateKey
        ],
        public_key: Union[
            ec.EllipticCurvePublicKey, x25519.X25519PublicKey, x448.X448PublicKey
        ],
    ) -> BinaPy:
        """This does an Elliptic Curve Diffie Hellman key exchange.

        This derives a shared key between a sender and a receiver, based on a public and a private key from each side.
        ECDH exchange produces the same key with either a sender private key and a recipient public key,
        or the matching sender public key and recipient private key.

        Args:
          private_key: a private EC key
          public_key: a public EC key

        Returns:
          a shared key
        """
        if isinstance(private_key, ec.EllipticCurvePrivateKey) and isinstance(
            public_key, ec.EllipticCurvePublicKey
        ):
            shared_key = private_key.exchange(ec.ECDH(), public_key)
        elif isinstance(private_key, x25519.X25519PrivateKey) and isinstance(
            public_key, x25519.X25519PublicKey
        ):
            shared_key = private_key.exchange(public_key)
        elif isinstance(private_key, x448.X448PrivateKey) and isinstance(
            public_key, x448.X448PublicKey
        ):
            shared_key = private_key.exchange(public_key)
        else:
            raise ValueError(
                "Invalid or unsupported private/public key combination for ECDH",
                private_key,
                public_key,
            )
        return BinaPy(shared_key)

    @classmethod
    def derive(
        cls,
        *,
        private_key: Union[
            ec.EllipticCurvePrivateKey, x25519.X25519PrivateKey, x448.X448PrivateKey
        ],
        public_key: Union[
            ec.EllipticCurvePublicKey, x25519.X25519PublicKey, x448.X448PublicKey
        ],
        otherinfo: bytes,
        key_size: int,
    ) -> BinaPy:
        """Derive a key using ECDH and Concat KDF Hash.

        Args:
          private_key: the private key
          public_key: the public key
          otherinfo: the Concat KDF "otherinfo" parameter
          key_size: the expected CEK key size

        Returns:
            the derived key
        """
        shared_key = cls.ecdh(private_key, public_key)
        ckdf = ConcatKDFHash(
            algorithm=hashes.SHA256(), length=key_size // 8, otherinfo=otherinfo
        )
        return BinaPy(ckdf.derive(shared_key))

    def generate_ephemeral_key(
        self,
    ) -> Union[
        ec.EllipticCurvePrivateKey, x25519.X25519PrivateKey, x448.X448PrivateKey
    ]:
        """Generate an ephemeral key that is suitable for use with this algorithm.

        Returns:
            a generated EllipticCurvePrivateKey, on the same curve as this algorithm key
        """
        if isinstance(
            self.key, (ec.EllipticCurvePrivateKey, ec.EllipticCurvePublicKey)
        ):
            return ec.generate_private_key(self.key.curve)
        elif isinstance(self.key, (x25519.X25519PrivateKey, x25519.X25519PublicKey)):
            return x25519.X25519PrivateKey.generate()
        elif isinstance(self.key, (x448.X448PublicKey, x448.X448PrivateKey)):
            return x448.X448PrivateKey.generate()

    def sender_key(
        self,
        ephemeral_private_key: Union[
            ec.EllipticCurvePrivateKey, x25519.X25519PrivateKey, x448.X448PrivateKey
        ],
        *,
        alg: str,
        key_size: int,
        **headers: Any,
    ) -> BinaPy:
        """Compute a CEK for encryption of a message. This method is meant for usage by a sender.

        Args:
          ephemeral_private_key: the EPK to use for this key
          alg: the content encryption algorithm identifier
          key_size: the expected CEK size
          **headers: additional headers to include for CEK derivation

        Returns:
            the CEK for encryption by the sender
        """
        with self.public_key_required() as key:
            apu = BinaPy(headers.get("apu", b"")).decode_from("b64u")
            apv = BinaPy(headers.get("apv", b"")).decode_from("b64u")
            otherinfo = self.otherinfo(alg, apu, apv, key_size)
            cek = self.derive(
                private_key=ephemeral_private_key,
                public_key=key,
                otherinfo=otherinfo,
                key_size=key_size,
            )
            return cek

    def recipient_key(
        self,
        ephemeral_public_key: Union[
            ec.EllipticCurvePublicKey, x25519.X25519PublicKey, x448.X448PublicKey
        ],
        *,
        alg: str,
        key_size: int,
        **headers: Any,
    ) -> BinaPy:
        """Compute a shared key. This method is meant for use by the recipient of an encrypted message.

        Args:
          ephemeral_public_key: the EPK, as received from sender
          alg: the content encryption algorithm identifier
          key_size: the CEK size
          **headers: additional headers as received from sender

        Returns:
            the CEK for decryption by the recipient
        """
        with self.private_key_required() as key:
            apu = BinaPy(headers.get("apu", b"")).decode_from("b64u")
            apv = BinaPy(headers.get("apv", b"")).decode_from("b64u")
            otherinfo = self.otherinfo(alg, apu, apv, key_size)
            cek = self.derive(
                private_key=key,
                public_key=ephemeral_public_key,
                otherinfo=otherinfo,
                key_size=key_size,
            )
            return cek


class BaseEcdhEs_AesKw(EcdhEs):
    """Base class for ECDH-ES+AESKW algorithms."""

    kwalg: Type[BaseAesKeyWrap]

    def wrap_key_with_epk(
        self,
        plainkey: bytes,
        ephemeral_private_key: Union[
            ec.EllipticCurvePrivateKey, x25519.X25519PrivateKey, x448.X448PrivateKey
        ],
        **headers: Any,
    ) -> BinaPy:
        """Wraps a key for content encryption.

        Args:
          plainkey: the key to wrap
          ephemeral_private_key: the EPK to use
          **headers: additional headers for CEK derivation

        Returns:
            the wrapped CEK
        """
        aes_key = self.sender_key(
            ephemeral_private_key, key_size=self.kwalg.key_size, **headers
        )
        return self.kwalg(aes_key).wrap_key(plainkey)

    def unwrap_key_with_epk(
        self,
        cipherkey: bytes,
        ephemeral_public_key: Union[
            ec.EllipticCurvePublicKey, x25519.X25519PublicKey, x448.X448PublicKey
        ],
        **headers: Any,
    ) -> BinaPy:
        """Unwrap a key for content decryption.

        Args:
          cipherkey: the wrapped key
          ephemeral_public_key: the EPK
          **headers: additional headers for CEK derivation

        Returns:
            the unwrapped key
        """
        aes_key = self.recipient_key(
            ephemeral_public_key, key_size=self.kwalg.key_size, **headers
        )
        return self.kwalg(aes_key).unwrap_key(cipherkey)


class EcdhEs_A128KW(BaseEcdhEs_AesKw):
    """ECDH-ES using Concat KDF and "A128KW" wrapping."""

    name = "ECDH-ES+A128KW"
    description = __doc__
    kwalg = A128KW


class EcdhEs_A192KW(BaseEcdhEs_AesKw):
    """ECDH-ES using Concat KDF and "A192KW" wrapping."""

    name = "ECDH-ES+A192KW"
    description = __doc__
    kwalg = A192KW


class EcdhEs_A256KW(BaseEcdhEs_AesKw):
    """ECDH-ES using Concat KDF and "A256KW" wrapping."""

    name = "ECDH-ES+A256KW"
    description = __doc__
    kwalg = A256KW
