use std::sync::atomic::Ordering;

use anyhow::Context as _;

use cryptoki::object::Attribute;
use cryptoki::object::AttributeType;

use crate::Key;
use crate::Result;

impl Key {
    fn get_capabilities(&self) -> Result<()> {
        if self.got_capabilities.load(Ordering::Relaxed) {
            return Ok(());
        }

        let (session, object) = self.open_ro_session()?;

        // We have the public key object.  To get the private key, we
        // need to log in, which means we need a password.  To avoid
        // that, we assume that the private key's capabilities
        // correspond to the public key's capabilities.  If that is
        // not the case, all bets are off.
        let types = &[
            // Decryption related.
            AttributeType::Encrypt,
            AttributeType::Wrap,
            AttributeType::Encapsulate,
            // Signing related.
            AttributeType::Verify,
        ][..];

        let attributes = session.get_attributes(object, types)
            .context("Failed to lookup key attributes")?;

        let mut decryption_capable = false;
        let mut signing_capable = false;

        for attribute in attributes.into_iter() {
            match attribute {
                Attribute::Encrypt(yes)
                    | Attribute::Wrap(yes)
                    | Attribute::Encapsulate(yes)
                    =>
                {
                    if yes {
                        decryption_capable = true;
                    }
                }
                Attribute::Verify(yes) => {
                    if yes {
                        signing_capable = true;
                    }
                }
                _ => (),
            }
        }

        self.signing_capable.store(signing_capable, Ordering::Relaxed);
        self.decryption_capable.store(decryption_capable, Ordering::Relaxed);
        self.got_capabilities.store(true, Ordering::Relaxed);

        Ok(())
    }

    /// Returns whether the key supports signing.
    ///
    /// This is derived from the cryptoki objects and not from an
    /// OpenPGP self-signature, which is not available.
    pub fn signing_capable(&self) -> Result<bool> {
        self.get_capabilities()?;
        Ok(self.signing_capable.load(Ordering::Relaxed))
    }

    /// Returns whether the key supports decryption.
    ///
    /// This is derived from the cryptoki objects and not from an
    /// OpenPGP self-signature, which is not available.
    pub fn decryption_capable(&self) -> Result<bool> {
        self.get_capabilities()?;
        Ok(self.decryption_capable.load(Ordering::Relaxed))
    }
}
