use super::{DataKeyWithTag, EncryptPayload};
use aes_gcm_siv::{aead::Payload, Nonce};
use recipher::key::{Iv, Key};
use uuid::Uuid;
use zerokms_protocol::DecryptionPolicy;

/// Metadata extracted from an encryption target after encryption.
pub(super) struct EncryptionMeta {
    pub iv: Iv,
    pub tag: Vec<u8>,
    pub descriptor: String,
    pub keyset_id: Option<Uuid>,
    pub decryption_policy: Option<DecryptionPolicy>,
}

/// Build the AAD (Additional Authenticated Data) for AES-GCM-SIV.
/// Shared between encryption and decryption to ensure they always match.
/// The tag already commits to the decryption policy (if any), so AAD
/// only needs descriptor + tag.
pub(super) fn build_aad(descriptor: &str, tag: &[u8]) -> Vec<u8> {
    let descriptor_bytes = descriptor.as_bytes();
    let mut aad = Vec::with_capacity(descriptor_bytes.len() + tag.len());
    aad.extend_from_slice(descriptor_bytes);
    aad.extend_from_slice(tag);
    aad
}

/// An encryption target is a payload that is ready to be encrypted.
pub struct EncryptionTarget<'e> {
    payload: EncryptPayload<'e>,
    key_with_tag: DataKeyWithTag,
    aad: Vec<u8>,
    keyset_id: Option<Uuid>,
}

impl<'e> EncryptionTarget<'e> {
    pub(crate) fn new(
        payload: EncryptPayload<'e>,
        key: DataKeyWithTag,
        keyset_id: Option<Uuid>,
    ) -> Self {
        let aad = build_aad(payload.descriptor, &key.tag);

        Self {
            payload,
            key_with_tag: key,
            aad,
            keyset_id,
        }
    }

    /// We're using the first 12 bytes of the IV as the nonce for AES-GCM-SIV.
    pub(super) fn nonce(&self) -> &Nonce {
        Nonce::from_slice(&self.key_with_tag.iv[..12])
    }

    pub(super) fn key(&self) -> &Key {
        self.key_with_tag.key()
    }

    /// Consumes self and returns the metadata that should be stored with the encrypted record.
    pub(super) fn into_meta(self) -> EncryptionMeta {
        let iv = self.key_with_tag.key.iv;
        let tag = self.key_with_tag.tag;

        EncryptionMeta {
            iv,
            tag,
            descriptor: self.payload.descriptor.to_string(),
            keyset_id: self.keyset_id,
            // The tagged policy comes from the server response (has MAC attached)
            decryption_policy: self.key_with_tag.decryption_policy,
        }
    }
}

impl<'e> From<&'e EncryptionTarget<'e>> for Payload<'e, 'e> {
    fn from(target: &'e EncryptionTarget<'e>) -> Self {
        Payload {
            msg: target.payload.msg,
            aad: &target.aad,
        }
    }
}
