#![cfg(target_os = "windows")]
use ncrypt::{Algorithm, ExportType, Key, ProviderName, RsaKeyBlob, StorageProvider};
use rsa::RsaPublicKey;

mod error;

pub use error::Error;

/// Handle to the Windows TPM Key Store Provider
pub struct Tpm(StorageProvider);

/// Handle to an RSA Key managed by the TPM
pub struct RsaKey(Key);

/// Iterator over all keys managed by the TPM
pub struct KeyIterator<'a>(ncrypt::KeyIterator<'a>, &'a Tpm);

pub struct KeyInfo<'a> {
    pub name: String,
    pub algorithm: String,
    tpm: &'a Tpm,
}

impl Tpm {
    /// Open a handle to the Windows TPM Key Store Provider
    ///
    /// Fails with [`Error::DeviceNotReady`] if no TPM is available
    pub fn open() -> Result<Self, Error> {
        let handle = StorageProvider::open(ProviderName::PlatformCrypto).map_err(|e| match e {
            ncrypt::Error::DeviceNotReady => Error::NotReady(e),
            _ => Error::Internal(e),
        })?;

        Ok(Self(handle))
    }

    pub fn create_rsa_key(&self, key_name: &str) -> Result<RsaKey, Error> {
        let key = self
            .0
            .create_persisted_key(key_name, Algorithm::Rsa)
            .map_err(|e| match e {
                ncrypt::Error::ObjectExists => Error::KeyExists(e),
                _ => Error::Internal(e),
            })?
            .finalize()
            .map_err(Error::Internal)?;

        Ok(RsaKey(key))
    }

    pub fn open_rsa_key(&self, key_name: &str) -> Result<RsaKey, Error> {
        let key = self.0.open_key(key_name).map_err(|e| match e {
            ncrypt::Error::BadKeyset => Error::KeyNotFound(e),
            _ => Error::Internal(e),
        })?;

        Ok(RsaKey(key))
    }

    /// Obtain an iterator over all the keys currently stored on the tpm
    pub fn enum_keys(&self) -> KeyIterator {
        KeyIterator(self.0.enum_keys(), self)
    }
}

impl RsaKey {
    /// Exports the public component of the rsa key from the key store
    ///
    /// Note: The private key used for decryption never leaves the key store,
    /// it can't be exported
    pub fn public_key(&self) -> Result<Box<RsaPublicKey>, Error> {
        let export = self
            .0
            .export(ExportType::RsaPublicKey)
            .map_err(Error::Internal)?;

        let RsaKeyBlob::PublicKey(key) = export else {
            return Err(Error::UnexpectedKeyType);
        };

        Ok(key)
    }

    /// Encrypts the given data and applies PKCS1 padding
    pub fn encrypt(&self, data: &[u8]) -> Result<Vec<u8>, Error> {
        self.0.encrypt(data).map_err(Error::Internal)
    }

    /// Decrypts the given data (data MUST have valid PKCS1 padding)
    pub fn decrypt(&self, data: &[u8]) -> Result<Vec<u8>, Error> {
        self.0.decrypt(data).map_err(|e| match e {
            ncrypt::Error::Other(inner) if inner as u32 == 0x80280095 => Error::DecryptionFailed,
            _ => Error::Internal(e),
        })
    }

    /// Permanently deletes the key from the keystore
    ///
    /// THIS CANNOT BE UNDONE
    pub fn delete(self) {
        self.0.delete()
    }
}

impl<'a> Iterator for KeyIterator<'a> {
    type Item = KeyInfo<'a>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(item) = self.0.next() {
            return Some(KeyInfo {
                name: item.name,
                algorithm: item.algorithm,
                tpm: self.1,
            });
        }

        None
    }
}

impl<'a> KeyInfo<'a> {
    /// Try to open the key. This will fail if the key is not an rsa key
    pub fn open(&self) -> Result<RsaKey, Error> {
        if self.algorithm != "RSA" {
            return Err(Error::UnexpectedKeyType);
        }

        self.tpm.open_rsa_key(&self.name)
    }
}