// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// Copyright by contributors to this project.
// SPDX-License-Identifier: (Apache-2.0 OR MIT)

use mls_rs_core::{
    crypto::{HpkeContextR, HpkeContextS},
    error::IntoAnyError,
};
use mls_rs_crypto_traits::{AeadType, KdfType};

use crate::{hpke::HpkeError, kdf::HpkeKdf};

use alloc::vec::Vec;
use core::fmt::{self, Debug};

/// A type representing an HPKE context
#[derive(Clone)]
pub(super) struct Context<KDF: KdfType, AEAD: AeadType> {
    exporter_secret: Vec<u8>,
    encryption_context: Option<EncryptionContext<AEAD>>,
    kdf: HpkeKdf<KDF>,
}

impl<KDF: KdfType + Debug, AEAD: AeadType + Debug> Debug for Context<KDF, AEAD> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Context")
            .field(
                "exporter_secret",
                &mls_rs_core::debug::pretty_bytes(&self.exporter_secret),
            )
            .field("encryption_context", &self.encryption_context)
            .field("kdf", &self.kdf)
            .finish()
    }
}

impl<KDF: KdfType, AEAD: AeadType> Context<KDF, AEAD> {
    #[inline]
    pub(super) fn new(
        encryption_context: Option<EncryptionContext<AEAD>>,
        exporter_secret: Vec<u8>,
        kdf: HpkeKdf<KDF>,
    ) -> Self {
        Self {
            exporter_secret,
            encryption_context,
            kdf,
        }
    }
}

impl<KDF: KdfType, AEAD: AeadType> Context<KDF, AEAD> {
    #[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)]
    async fn seal(&mut self, aad: Option<&[u8]>, data: &[u8]) -> Result<Vec<u8>, HpkeError> {
        self.encryption_context
            .as_mut()
            .ok_or(HpkeError::ExportOnlyMode)?
            .seal(aad, data)
            .await
    }

    #[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)]
    async fn open(&mut self, aad: Option<&[u8]>, ciphertext: &[u8]) -> Result<Vec<u8>, HpkeError> {
        self.encryption_context
            .as_mut()
            .ok_or(HpkeError::ExportOnlyMode)?
            .open(aad, ciphertext)
            .await
    }

    #[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)]
    async fn export(&self, exporter_context: &[u8], len: usize) -> Result<Vec<u8>, HpkeError> {
        self.kdf
            .labeled_expand(&self.exporter_secret, b"sec", exporter_context, len)
            .await
            .map_err(|e| HpkeError::KdfError(e.into_any_error()))
    }
}

#[derive(Debug, Clone)]
pub struct ContextS<KDF: KdfType, AEAD: AeadType>(pub(super) Context<KDF, AEAD>);

#[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)]
#[cfg_attr(all(target_arch = "wasm32", mls_build_async), maybe_async::must_be_async(?Send))]
#[cfg_attr(
    all(not(target_arch = "wasm32"), mls_build_async),
    maybe_async::must_be_async
)]
impl<KDF: KdfType, AEAD: AeadType> HpkeContextS for ContextS<KDF, AEAD> {
    type Error = HpkeError;

    async fn export(&self, exporter_context: &[u8], len: usize) -> Result<Vec<u8>, Self::Error> {
        self.0.export(exporter_context, len).await
    }

    /// # Errors
    ///
    /// Returns [SequenceNumberOverflow](HpkeError::SequenceNumberOverflow)
    /// in the event that the sequence number overflows. The sequence number is a u64 and starts
    /// at 0.
    async fn seal(&mut self, aad: Option<&[u8]>, data: &[u8]) -> Result<Vec<u8>, Self::Error> {
        self.0.seal(aad, data).await
    }
}

#[derive(Debug, Clone)]
pub struct ContextR<KDF: KdfType, AEAD: AeadType>(pub(super) Context<KDF, AEAD>);

#[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)]
#[cfg_attr(all(target_arch = "wasm32", mls_build_async), maybe_async::must_be_async(?Send))]
#[cfg_attr(
    all(not(target_arch = "wasm32"), mls_build_async),
    maybe_async::must_be_async
)]
impl<KDF: KdfType, AEAD: AeadType> HpkeContextR for ContextR<KDF, AEAD> {
    type Error = HpkeError;

    async fn export(&self, exporter_context: &[u8], len: usize) -> Result<Vec<u8>, Self::Error> {
        self.0.export(exporter_context, len).await
    }

    /// # Errors
    ///
    /// Returns [SequenceNumberOverflow](HpkeError::SequenceNumberOverflow)
    /// in the event that the sequence number overflows. The sequence number is a u64 and starts
    /// at 0.
    ///
    /// Returns [AeadError](HpkeError::AeadError) if decryption fails due to either an invalid
    /// `aad` value, or incorrect cipher key.
    async fn open(
        &mut self,
        aad: Option<&[u8]>,
        ciphertext: &[u8],
    ) -> Result<Vec<u8>, Self::Error> {
        self.0.open(aad, ciphertext).await
    }
}

#[derive(PartialEq, Eq, Clone)]
pub(super) struct EncryptionContext<AEAD: AeadType> {
    base_nonce: Vec<u8>,
    seq_number: u64,
    aead: AEAD,
    aead_key: Vec<u8>,
}

impl<AEAD: AeadType + Debug> Debug for EncryptionContext<AEAD> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("EncryptionContext")
            .field(
                "base_nonce",
                &mls_rs_core::debug::pretty_bytes(&self.base_nonce),
            )
            .field("seq_number", &self.seq_number)
            .field("aead", &self.aead)
            .field(
                "aead_key",
                &mls_rs_core::debug::pretty_bytes(&self.aead_key),
            )
            .finish()
    }
}

impl<AEAD: AeadType> EncryptionContext<AEAD> {
    pub fn new(base_nonce: Vec<u8>, aead: AEAD, aead_key: Vec<u8>) -> Result<Self, HpkeError> {
        (base_nonce.len() == aead.nonce_size())
            .then_some(())
            .ok_or(HpkeError::IncorrectNonceLen(
                base_nonce.len(),
                aead.nonce_size(),
            ))?;

        (aead_key.len() == aead.key_size())
            .then_some(())
            .ok_or(HpkeError::IncorrectKeyLen(aead_key.len(), aead.key_size()))?;

        Ok(EncryptionContext {
            base_nonce,
            seq_number: 0,
            aead,
            aead_key,
        })
    }
}

impl<AEAD: AeadType> EncryptionContext<AEAD> {
    //draft-irtf-cfrg-hpke Section 5.2.  Encryption and Decryption
    fn compute_nonce(&self) -> Vec<u8> {
        let mut nonce = self.base_nonce.clone();

        // XOR the sequence number into the last 4 bytes of the nonce
        nonce
            .iter_mut()
            .rev()
            .zip(self.seq_number.to_le_bytes())
            .for_each(|(n, s)| *n ^= s);

        nonce
    }

    #[inline]
    fn increment_seq(&mut self) -> Result<(), HpkeError> {
        // If the sequence number is going to roll over just throw an error
        self.seq_number = self
            .seq_number
            .checked_add(1)
            .ok_or(HpkeError::SequenceNumberOverflow)?;

        Ok(())
    }

    #[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)]
    pub async fn seal(&mut self, aad: Option<&[u8]>, pt: &[u8]) -> Result<Vec<u8>, HpkeError> {
        let ct = self
            .aead
            .seal(&self.aead_key, pt, aad, &self.compute_nonce())
            .await
            .map_err(|e| HpkeError::AeadError(e.into_any_error()))?;

        self.increment_seq()?;

        Ok(ct)
    }

    #[cfg_attr(not(mls_build_async), maybe_async::must_be_sync)]
    pub async fn open(&mut self, aad: Option<&[u8]>, ct: &[u8]) -> Result<Vec<u8>, HpkeError> {
        let pt = self
            .aead
            .open(&self.aead_key, ct, aad, &self.compute_nonce())
            .await
            .map_err(|e| HpkeError::AeadError(e.into_any_error()))?;

        self.increment_seq()?;

        Ok(pt)
    }
}