//! This module contains the implementation of the `CryptoOracle` trait for the `Hsm` plugin.
//! The `HsmCryptoOracle` struct is a wrapper around an `HSM` instance and is responsible for
//! encrypting, decrypting, and signing data using the HSM.
//! This blanket implementation "glues" the `Hsm` interface with the `CryptoOracle` interface.

use std::sync::Arc;

use async_trait::async_trait;
use cosmian_logger::debug;
use zeroize::Zeroizing;

use crate::{
    CryptoAlgorithm, CryptoOracle, HSM, InterfaceError, InterfaceResult, KeyMetadata, KeyType,
    SigningAlgorithm, crypto_oracle::EncryptedContent,
};

pub struct HsmCryptoOracle {
    hsm: Arc<dyn HSM + Send + Sync>,
}

impl HsmCryptoOracle {
    pub fn new(hsm: Arc<dyn HSM + Send + Sync>) -> Self {
        Self { hsm }
    }
}

#[async_trait]
impl CryptoOracle for HsmCryptoOracle {
    async fn encrypt(
        &self,
        uid: &str,
        data: &[u8],
        cryptographic_algorithm: Option<CryptoAlgorithm>,
        authenticated_encryption_additional_data: Option<&[u8]>,
    ) -> InterfaceResult<EncryptedContent> {
        if authenticated_encryption_additional_data.is_some() {
            return Err(InterfaceError::InvalidRequest(
                "Additional authenticated data are not supported on HSMs for now".to_owned(),
            ));
        }
        let (mut slot_id, mut key_id) = parse_uid(uid)?;
        let supported_algorithms = self.hsm.get_supported_algorithms(slot_id).await?;
        let cryptographic_algorithm = if let Some(ca) = cryptographic_algorithm {
            ca
        } else {
            debug!("Using default algorithm to encrypt");
            match self.hsm.get_key_type(slot_id, &key_id).await? {
                None => {
                    return Err(InterfaceError::InvalidRequest(format!(
                        "The key type of key: {uid}, cannot be determined"
                    )));
                }
                Some(key_type) => match key_type {
                    KeyType::AesKey => CryptoAlgorithm::get_aes_algorithm(&supported_algorithms)?,
                    KeyType::RsaPublicKey => {
                        CryptoAlgorithm::get_rsa_algorithm(&supported_algorithms)?
                    }
                    KeyType::RsaPrivateKey => {
                        // try fetching the corresponding public key
                        let pk_uid = format!("{uid}_pk");
                        debug!(
                            "encrypt: an RSA private key {uid} was specified. Trying to use \
                             public key {pk_uid} for encryption"
                        );
                        (slot_id, key_id) = parse_uid(&pk_uid)?;
                        self.hsm
                            .get_key_type(slot_id, &key_id)
                            .await?
                            .ok_or_else(|| {
                                InterfaceError::InvalidRequest(format!(
                                    "The key {uid} is a private key, but no public key {pk_uid} \
                                     is available"
                                ))
                            })?;
                        CryptoAlgorithm::get_rsa_algorithm(&supported_algorithms)?
                    }
                },
            }
        };
        self.hsm
            .encrypt(slot_id, &key_id, cryptographic_algorithm, data)
            .await
    }

    async fn decrypt(
        &self,
        uid: &str,
        data: &[u8],
        cryptographic_algorithm: Option<CryptoAlgorithm>,
        authenticated_encryption_additional_data: Option<&[u8]>,
    ) -> InterfaceResult<Zeroizing<Vec<u8>>> {
        if authenticated_encryption_additional_data.is_some() {
            return Err(InterfaceError::InvalidRequest(
                "Additional authenticated data are not supported on HSMs for now".to_owned(),
            ));
        }
        let (slot_id, key_id) = parse_uid(uid)?;
        let supported_algorithms = self.hsm.get_supported_algorithms(slot_id).await?;
        let cryptographic_algorithm = if let Some(ca) = cryptographic_algorithm {
            ca
        } else {
            debug!("Using default algorithm to decrypt");
            match self.hsm.get_key_type(slot_id, &key_id).await? {
                None => {
                    return Err(InterfaceError::InvalidRequest(
                        "The key {}type is not known".to_owned(),
                    ));
                }
                Some(key_type) => match key_type {
                    KeyType::AesKey => CryptoAlgorithm::get_aes_algorithm(&supported_algorithms)?,
                    KeyType::RsaPrivateKey => {
                        CryptoAlgorithm::get_rsa_algorithm(&supported_algorithms)?
                    }
                    KeyType::RsaPublicKey => {
                        return Err(InterfaceError::Default(
                            "An RSA public key cannot be used to decrypt".to_owned(),
                        ));
                    }
                },
            }
        };
        self.hsm
            .decrypt(slot_id, &key_id, cryptographic_algorithm, data)
            .await
    }

    async fn get_key_type(&self, uid: &str) -> InterfaceResult<Option<KeyType>> {
        let (slot_id, key_id) = parse_uid(uid)?;
        self.hsm.get_key_type(slot_id, &key_id).await
    }

    async fn get_key_metadata(&self, uid: &str) -> InterfaceResult<Option<KeyMetadata>> {
        let (slot_id, key_id) = parse_uid(uid)?;
        self.hsm.get_key_metadata(slot_id, &key_id).await
    }

    async fn sign(
        &self,
        uid: &str,
        data: &[u8],
        cryptographic_parameters: Option<
            &cosmian_kmip::kmip_2_1::kmip_types::CryptographicParameters,
        >,
    ) -> InterfaceResult<Vec<u8>> {
        let (slot_id, key_id) = parse_uid(uid)?;
        let key_type = self.hsm.get_key_type(slot_id, &key_id).await?;
        match key_type {
            Some(KeyType::RsaPrivateKey) => {}
            Some(other) => {
                return Err(InterfaceError::InvalidRequest(format!(
                    "Sign: key {uid} is a {other:?}, expected an RSA private key"
                )));
            }
            None => {
                return Err(InterfaceError::InvalidRequest(format!(
                    "Sign: key {uid} not found on the HSM"
                )));
            }
        }
        let algorithm = SigningAlgorithm::from_kmip(cryptographic_parameters)?;
        debug!("sign: using algorithm {algorithm:?} for key {uid}");
        self.hsm.sign(slot_id, &key_id, algorithm, data).await
    }
}

/// Parse the `uid` into a `slot_id` and `key_id`
fn parse_uid(uid: &str) -> InterfaceResult<(usize, Vec<u8>)> {
    let (slot_id, key_id) = uid
        .trim_start_matches("hsm::")
        .split_once("::")
        .ok_or_else(|| {
            InterfaceError::InvalidRequest(
                format!("An HSM create request must have a uid in the form of 'hsm::<slot_id>::<key_id>'. Got {uid}")
            )
        })?;
    let slot_id = slot_id.parse::<usize>().map_err(|e| {
        InterfaceError::InvalidRequest(format!("The slot_id must be a valid unsigned integer: {e}"))
    })?;
    Ok((slot_id, key_id.as_bytes().to_vec()))
}
