use anyhow::Context;

use sequoia_openpgp as openpgp;
use openpgp::packet;
use openpgp::packet::key::PublicParts;
use openpgp::packet::key::UnspecifiedRole;

use cryptoki::object::Attribute;
use cryptoki::object::AttributeType;
use cryptoki::object::KeyType;
use cryptoki::object::ObjectClass;
use cryptoki::object::ObjectHandle;
use cryptoki::session::Session;

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

fn get_private_key(session: &Session, id: Option<&[u8]>)
    -> Result<ObjectHandle>
{
    if let Some(key_id) = id {
        let attributes = &[
            Attribute::Class(ObjectClass::PRIVATE_KEY),
            Attribute::Id(key_id.to_vec()),
        ];

        let mut objects = match session.iter_objects(attributes) {
            Ok(objects) => objects,
            Err(err) => {
                return Err(anyhow::anyhow!(
                    "Error listing objects: {}", err));
            }
        };

        let first = objects.next();
        let second = objects.next();
        match (first, second) {
            (Some(object), None) => Ok(object.context("Getting private key")?),
            (Some(_), Some(_)) => {
                return Err(anyhow::anyhow!(
                    "Error: multiple private keys correspond to the public key ({})",
                    String::from_utf8_lossy(key_id)));
            }
            (None, _) => {
                return Err(anyhow::anyhow!(
                    "Error: no private key corresponds to the public key ({})",
                    String::from_utf8_lossy(key_id)));
            }
        }
    } else {
        Err(anyhow::anyhow!(
            "Failed to find corresponding private key: public key has no \
             ID attribute"))
    }
}

pub(crate) struct KeyInfo {
    pub private_key: ObjectHandle,
    pub key_type: KeyType,
}

impl KeyInfo {
    /// Check that a standalone OpenPGP key matches the addressed
    /// key.
    pub fn check_standalone(
        self,
        openpgp_key: &packet::Key<PublicParts, UnspecifiedRole>,
        composite_info: Option<KeyInfo>,
        expected_key_type: KeyType)
        -> Result<Self>
    {
        if self.key_type != expected_key_type {
            return Err(anyhow::anyhow!(
                "Unsupported key: can't use a {} with a {} OpenPGP key",
                self.key_type, openpgp_key.pk_algo()));
        }

        if let Some(composite_info) = composite_info {
            Err(anyhow::anyhow!(
                "{} is not a composite algorithm, but have a \
                 composite key ({} + {})",
                openpgp_key.pk_algo(),
                self.key_type,
                composite_info.key_type))
        } else {
            Ok(self)
        }
    }

    /// Check that a composite OpenPGP key matches the addressed
    /// keys.
    pub fn check_composite(
        self,
        openpgp_key: &packet::Key<PublicParts, UnspecifiedRole>,
        composite_info: Option<KeyInfo>,
        expected_key_type: KeyType,
        expected_key_type2: KeyType)
        -> Result<(KeyInfo, KeyInfo)>
    {
        let composite_info = if let Some(composite_info)
            = composite_info
        {
            composite_info
        } else {
            return Err(anyhow::anyhow!(
                "{} is a composite key, but only found the {} half",
                openpgp_key.pk_algo(),
                self.key_type));
        };

        if self.key_type == expected_key_type
            && composite_info.key_type == expected_key_type2
        {
            Ok((self, composite_info))
        } else if composite_info.key_type == expected_key_type
            && self.key_type == expected_key_type2
        {
            Ok((composite_info, self))
        } else {
            Err(anyhow::anyhow!(
                "Invalid key: a {} key requires \
                 a {} key and a {} key, but found \
                 a {} key and a {} key",
                openpgp_key.pk_algo(),
                expected_key_type, expected_key_type2,
                self.key_type, composite_info.key_type))
        }
    }

    pub fn new(session: &Session, key: &Key) -> Result<KeyInfo> {
        // The default public session only allows enumerating public
        // key objects.  We now look up the private key object using
        // the user session.  By convention, the public and private
        // keys have the same ID (CKA_ID).
        let private_key = get_private_key(session, key.key_id())?;

        let typs = &[
            AttributeType::KeyType,
        ][..];
        let attributes = session.get_attributes(private_key.clone(), typs)?;

        let mut key_type = None;
        for attr in attributes.into_iter() {
            match attr {
                Attribute::KeyType(value) => key_type = Some(value),
                _ => {
                    panic!("PKCS#11 driver returned unrequested attribute: {:?}",
                           attr);
                }
            }
        }

        let key_type = if let Some(key_type) = key_type {
            key_type
        } else {
            return Err(anyhow::anyhow!(
                "Invalid private key object: missing key type"));
        };

        Ok(KeyInfo {
            private_key,
            key_type,
        })
    }
}

