// SPDX-License-Identifier: MIT
//
// Copyright IBM Corp. 2024

use enum_dispatch::enum_dispatch;
use pv::request::openssl::pkey::{PKey, Public};
use pv::request::{decrypt_aead, encrypt_aead, Confidential, SymKey, SymKeyType};

use super::se_hdr::{SeHdrBinV1, SeHdrBinV2, SeHdrData, SeHdrVersioned};
use crate::pv_utils::error::{Error, Result};
use crate::pv_utils::serializing::deserialize_from_bytes;

/// Trait to be used for Authenticated Encryption with Associated Data (AEAD)
/// data structures.
#[enum_dispatch]
pub trait AeadCipherTrait {
    /// Returns the AEAD key type used by the data structure.
    fn aead_key_type(&self) -> SymKeyType;

    /// Returns the AEAD tag size used by the data structure.
    fn aead_tag_size(&self) -> usize;

    /// Returns the initialization vector (IV) used for AEAD
    /// encryption/decryption.
    fn iv(&self) -> &[u8];
}

/// Trait to be used for AEAD cipher data
#[enum_dispatch]
pub trait AeadDataTrait {
    /// Returns the authenticated associated data.
    fn aad(&self) -> Result<Vec<u8>>;

    /// Returns the encrypted data.
    fn data(&self) -> Vec<u8>;

    /// Returns the tag data.
    fn tag(&self) -> Vec<u8>;
}

/// Trait to be used for AEAD plaintext data
#[enum_dispatch]
pub trait AeadPlainDataTrait {
    /// Returns the authenticated associated data.
    fn aad(&self) -> Result<Vec<u8>>;

    /// Returns the unencrypted data.
    fn data(&self) -> Result<Confidential<Vec<u8>>>;

    /// Returns the tag data.
    fn tag(&self) -> Vec<u8>;
}

/// Key exchange related methods
pub trait KeyExchangeTrait {
    type TargetKeyType;
    type PrivateKeyType: ToOwned;

    /// Checks if a public target key was used.
    ///
    /// # Errors
    ///
    /// This function will return an error if the public key cannot be converted
    /// into a hash.
    fn contains<K>(&self, key: K) -> Result<bool>
    where
        K: AsRef<Self::TargetKeyType>;

    /// Checks if the hash of a public key was used.
    fn contains_hash<H: AsRef<[u8]>>(&self, hash: H) -> bool;

    /// Returns the public key from the party (e.g. guest owner) that is stored
    /// in the data structed.
    ///
    /// # Errors
    ///
    /// This function will return an error if the public key cannot be
    /// reconstructed.
    fn cust_pub_key(&mut self) -> Result<PKey<Public>>;

    /// Returns the key type of the exchanged key.
    fn key_type(&self) -> SymKeyType;

    // TODO Implement it
    // /// Derive the key.
    // ///
    // /// # Errors
    // ///
    // /// This function will return an error if there is no customer public key is
    // /// available or the key derivations fails.
    // fn derive_key<K>(&mut self, other_priv_key: K) -> Result<SymKey> where
    // K: AsRef<Self::PrivateKeyType>{
    //     match self.key_type() {
    //         SymKeyType::Aes256Gcm =>
    //             match Self::TargetKeyType {
    //                 PKeyRef::<Public> => {todo!()},
    //                 HybridPKey => {todo!()},
    //             }
    //         _ => unreachable!("BUG"),
    //     }
    // }
}

/// Trait to be used for plain UV data.
#[enum_dispatch]
pub trait UvDataPlainTrait:
    AeadPlainDataTrait + AeadCipherTrait + KeyExchangeTrait + Clone
{
    /// Returned type by [`Self::encrypt`].
    type C: UvDataTrait;

    /// Encrypt the plain data.
    ///
    /// # Errors
    ///
    /// This function will return an error if the passed `key` has the wrong
    /// key type or the encryption fails.
    fn encrypt(&self, key: &SymKey) -> Result<Self::C>
    where
        <Self as UvDataPlainTrait>::C: for<'a> deku::DekuContainerRead<'a>,
    {
        if key.key_type() != self.aead_key_type() {
            return Err(Error::UnexpectedKeyType {
                given: self.key_type().to_string(),
                expected: self.aead_key_type().to_string(),
            });
        }
        let aad = self.aad().map_err(|err| match err {
            Error::Deku(_) => Error::InvalidSeHdr,
            err => err,
        })?;
        let unecrypted_data = self.data().map_err(|err| match err {
            Error::Deku(_) => Error::InvalidSeHdr,
            err => err,
        })?;
        let iv = self.iv();
        let result = encrypt_aead(key, iv, &aad, unecrypted_data.value())?;
        Self::C::try_from_data(&result.into_buf())
    }

    /// Parses and converts the data into an instance of [`Self`] if possible.
    ///
    /// # Errors
    ///
    /// This function will return an error if the data could not parsed or
    /// converted.
    fn try_from_data<'a>(data: &'a [u8]) -> Result<Self>
    where
        Self: deku::DekuContainerRead<'a> + Sized,
    {
        deserialize_from_bytes(data)
    }
}

/// Trait to be used for (cipher) UV data.
#[enum_dispatch]
pub trait UvDataTrait: AeadDataTrait + AeadCipherTrait + KeyExchangeTrait + Clone {
    /// Returned type by [`Self::decrypt`].
    type P: UvDataPlainTrait;

    /// Decrypt the UV data.
    ///
    /// # Errors
    ///
    /// This function will return an error if the passed `key` has the wrong key
    /// type or the decryption fails.
    fn decrypt(&self, key: &SymKey) -> Result<Self::P>
    where
        <Self as UvDataTrait>::P: for<'a> deku::DekuContainerRead<'a>,
    {
        if key.key_type() != self.aead_key_type() {
            return Err(Error::UnexpectedKeyType {
                given: key.key_type().to_string(),
                expected: self.aead_key_type().to_string(),
            });
        }

        let tag_size = self.aead_tag_size();
        let aad = self.aad()?;
        let unecrypted_data = self.data();
        let iv = self.iv();
        let tag = self.tag();
        assert_eq!(tag.len(), tag_size);
        let result = decrypt_aead(key, iv, &aad, &unecrypted_data, &tag)?;
        Self::P::try_from_data(result.into_buf().value())
    }

    /// Parses and converts the data into an instance of [`Self`] if possible.
    ///
    /// # Errors
    ///
    /// This function will return an error if the data could not parsed or
    /// converted.
    fn try_from_data<'a>(data: &'a [u8]) -> Result<Self>
    where
        Self: deku::DekuContainerRead<'a> + Sized,
    {
        deserialize_from_bytes(data)
    }
}