use std::borrow::Cow;

use anyhow::Context;

use sequoia_openpgp as openpgp;
use openpgp::crypto::mpi;
use openpgp::crypto;
use openpgp::packet::key::PublicParts;
use openpgp::packet::key::UnspecifiedRole;
use openpgp::packet;
use openpgp::types::HashAlgorithm;
use openpgp::types::PublicKeyAlgorithm;

use cryptoki::mechanism::Mechanism;
use cryptoki::mechanism::eddsa::EddsaParams;
use cryptoki::mechanism::eddsa::EddsaSignatureScheme;
use cryptoki::mechanism::dsa;
use cryptoki::object::KeyType;
use cryptoki::session::UserType;
use cryptoki::types::AuthPin;

use crate::Result;
use crate::Prompt;
use crate::key_info::KeyInfo;
use crate::objects::Key;

/// A signer for a PKCS#11 key.
///
/// This implements [`sequoia_openpgp::crypto::Signer`] and hence can
/// be used wherever a Sequoia method needs to create a signature.
///
///   [`sequoia_openpgp::crypto::Signer`]: https://docs.rs/sequoia-openpgp/latest/sequoia_openpgp/crypto/trait.Signer.html
pub struct Signer<'key, 'prompt> {
    key: Cow<'key, Key>,
    openpgp_key: packet::Key<PublicParts, UnspecifiedRole>,
    user_pin: Option<AuthPin>,
    prompt: Box<dyn Prompt + 'prompt>,
}

impl<'key, 'prompt> Signer<'key, 'prompt> {
    /// Creates a `Signer`.
    ///
    /// This is a bit ugly, because we can't stash a PKCS#11
    /// `Session`: a struct that implements `Signer` has to implement
    /// `Sync`, but a `Session` is not `Sync`.
    ///
    /// If the user PIN is known, pass it through to avoid prompting
    /// the user.  (If `user_pin` is wrong, we don't prompt the user,
    /// but immediately return an error.)  Otherwise, the user is
    /// prompted.
    pub fn new<P>(prompt: P, key: &'key Key, user_pin: Option<AuthPin>)
        -> Result<Self>
    where
        P: Prompt + 'prompt
    {
        // Before we prompt for the PIN, make sure we can convert the
        // public key to an OpenPGP key.
        let openpgp_key = key.to_openpgp_key()
            .context("Converting public key to OpenPGP format")?
            .role_into_unspecified();

        Ok(Self {
            key: Cow::Borrowed(key),
            openpgp_key,
            user_pin,
            prompt: Box::new(prompt),
        })
    }

    /// Creates a `Signer`.
    ///
    /// Like [`Signer::new`], but takes ownership of `key`.
    pub fn new_owned<P>(prompt: P, key: Key, user_pin: Option<AuthPin>)
        -> Result<Self>
    where
        P: Prompt + 'prompt
    {
        // Before we prompt for the PIN, make sure we can convert the
        // public key to an OpenPGP key.
        let openpgp_key = key.to_openpgp_key()
            .context("Converting public key to OpenPGP format")?
            .role_into_unspecified();

        Ok(Self {
            key: Cow::Owned(key),
            openpgp_key,
            user_pin,
            prompt: Box::new(prompt),
        })
    }
}

impl crypto::Signer for Signer<'_, '_> {
    fn public(&self) -> &packet::Key<PublicParts, UnspecifiedRole> {
        &self.openpgp_key
    }

    fn sign(&mut self, hash_algo: HashAlgorithm, digest: &[u8])
            -> Result<mpi::Signature>
    {
        let (session, _object) = self.key.open_ro_session()?;

        // Login as the user.
        if let Some(user_pin) = self.user_pin.as_ref() {
            // PIN was supplied.
            session.login(UserType::User, Some(&user_pin))
                .context("Logging in to token")?;
        } else {
            // Prompt the user for the PIN.
            loop {
                let user_pin = self.prompt.prompt_to_unlock("the user PIN")?;
                let user_pin = String::from_utf8(user_pin.map(|p| p.to_vec()))
                    .context("Parsing user PIN")?;
                let user_pin = AuthPin::new(user_pin.into());

                match session.login(UserType::User, Some(&user_pin)) {
                    Ok(_) => {
                        self.user_pin = Some(user_pin);
                        break;
                    }
                    Err(err) => {
                        eprintln!("Error logging in: {}", err);
                    }
                }
            }
        }

        let key_info = KeyInfo::new(&session, &self.key)?;

        let composite_info = if let Some(composite)
            = self.key.composite_key().as_ref()
        {
            Some(KeyInfo::new(&session, composite)?)
        } else {
            None
        };
        match self.openpgp_key.pk_algo() {
            PublicKeyAlgorithm::Ed25519 => {
                let key_info = key_info.check_standalone(
                    &self.openpgp_key, composite_info, KeyType::EC_EDWARDS)?;

                // RFC 9580, Section 12.7:
                //
                // [W]hile [RFC8032] describes different variants of EdDSA,
                // OpenPGP uses the "pure" variant (PureEdDSA). The hashing
                // that happens with OpenPGP is done as part of the standard
                // OpenPGP signature process, and that hash itself is fed as
                // the input message to the PureEdDSA algorithm.
                //
                // https://www.rfc-editor.org/rfc/rfc9580.html#section-12.7
                let eddsa_params = EddsaParams::new(EddsaSignatureScheme::Pure);

                let raw = session.sign(&Mechanism::Eddsa(eddsa_params),
                                       key_info.private_key.clone(), digest)
                    .with_context(|| {
                        format!("Signing data")
                    })?;

                if raw.len() != 64 {
                    return Err(anyhow::anyhow!(
                        "Token returned an invalid signature: expected \
                         64 bytes of signature data, got {} bytes",
                        raw.len()));
                }

                Ok(mpi::Signature::Ed25519 {
                    s: raw.try_into().expect("checked length above"),
                })
            },

            PublicKeyAlgorithm::Ed448 => {
                let key_info = key_info.check_standalone(
                    &self.openpgp_key, composite_info, KeyType::EC_EDWARDS)?;

                // RFC 9580, Section 12.7:
                //
                // [W]hile [RFC8032] describes different variants of EdDSA,
                // OpenPGP uses the "pure" variant (PureEdDSA). The hashing
                // that happens with OpenPGP is done as part of the standard
                // OpenPGP signature process, and that hash itself is fed as
                // the input message to the PureEdDSA algorithm.
                //
                // As specified in [RFC8032], Ed448 also expects a "context
                // string". In OpenPGP, Ed448 is used with the empty string as
                // a context string.
                //
                // https://www.rfc-editor.org/rfc/rfc9580.html#section-12.7
                let eddsa_params =
                    EddsaParams::new(EddsaSignatureScheme::Ed448(&[]));

                let raw = session.sign(&Mechanism::Eddsa(eddsa_params),
                                       key_info.private_key.clone(), digest)
                    .with_context(|| {
                        format!("Signing data")
                    })?;

                if raw.len() != 114 {
                    return Err(anyhow::anyhow!(
                        "Token returned an invalid signature: expected \
                         114 bytes of signature data, got {} bytes",
                        raw.len()));
                }

                Ok(mpi::Signature::Ed448 {
                    s: raw.try_into().expect("checked length above"),
                })
            },

            // MLDSA65_Ed25519.
            PublicKeyAlgorithm::MLDSA65_Ed25519 => {
                let (mldsa, eddsa) = key_info.check_composite(
                    &self.openpgp_key, composite_info,
                    KeyType::ML_DSA, KeyType::EC_EDWARDS)?;

                let eddsa = session.sign(
                    &Mechanism::Eddsa(EddsaParams::new(
                        EddsaSignatureScheme::Pure)),
                    eddsa.private_key, digest)
                    .context("Generating the Ed25519 signature")?
                    .try_into()
                    .map_err(|_| anyhow::anyhow!("signature has wrong size"))?;

                let mldsa = session.sign(
                    &Mechanism::MlDsa(dsa::SignAdditionalContext::new(
                        dsa::HedgeType::Required, None)),
                    mldsa.private_key, digest)
                    .context("Generating the ML-DSA signature")?
                    .try_into()
                    .map_err(|_| anyhow::anyhow!("signature has wrong size"))?;

                Ok(mpi::Signature::MLDSA65_Ed25519 {
                    eddsa: Box::new(eddsa),
                    mldsa: Box::new(mldsa),
                })
            },

            // MLDSA87_Ed448.
            PublicKeyAlgorithm::MLDSA87_Ed448 => {
                let (mldsa, eddsa) = key_info.check_composite(
                    &self.openpgp_key, composite_info,
                    KeyType::ML_DSA, KeyType::EC_EDWARDS)?;

                let eddsa = session.sign(
                    &Mechanism::Eddsa(EddsaParams::new(
                        EddsaSignatureScheme::Ed448(&[]))),
                    eddsa.private_key, digest)
                    .context("Generating the Ed448 signature")?
                    .try_into()
                    .map_err(|_| anyhow::anyhow!("signature has wrong size"))?;

                let mldsa = session.sign(
                    &Mechanism::MlDsa(dsa::SignAdditionalContext::new(
                        dsa::HedgeType::Required, None)),
                    mldsa.private_key, digest)
                    .context("Generating the ML-DSA signature")?
                    .try_into()
                    .map_err(|_| anyhow::anyhow!("signature has wrong size"))?;

                Ok(mpi::Signature::MLDSA87_Ed448 {
                    eddsa: Box::new(eddsa),
                    mldsa: Box::new(mldsa),
                })
            },

            // SLH-DSA-128s.
            PublicKeyAlgorithm::SLHDSA128s => {
                let key_info = key_info.check_standalone(
                    &self.openpgp_key, composite_info, KeyType::SLH_DSA)?;

                let sig = session.sign(
                    &Mechanism::SlhDsa(dsa::SignAdditionalContext::new(
                        dsa::HedgeType::Required, None)),
                    key_info.private_key, digest)?
                    .try_into()
                    .map_err(|_| anyhow::anyhow!("signature has wrong size"))?;

                Ok(mpi::Signature::SLHDSA128s {
                    sig,
                })
            },

            // SLH-DSA-128f.
            PublicKeyAlgorithm::SLHDSA128f => {
                let key_info = key_info.check_standalone(
                    &self.openpgp_key, composite_info, KeyType::SLH_DSA)?;

                let sig = session.sign(
                    &Mechanism::SlhDsa(dsa::SignAdditionalContext::new(
                        dsa::HedgeType::Required, None)),
                    key_info.private_key, digest)?
                    .try_into()
                    .map_err(|_| anyhow::anyhow!("signature has wrong size"))?;

                Ok(mpi::Signature::SLHDSA128f {
                    sig,
                })
            },

            // SLH-DSA-256s.
            PublicKeyAlgorithm::SLHDSA256s => {
                let key_info = key_info.check_standalone(
                    &self.openpgp_key, composite_info, KeyType::SLH_DSA)?;

                let sig = session.sign(
                    &Mechanism::SlhDsa(dsa::SignAdditionalContext::new(
                        dsa::HedgeType::Required, None)),
                    key_info.private_key, digest)?
                    .try_into()
                    .map_err(|_| anyhow::anyhow!("signature has wrong size"))?;

                Ok(mpi::Signature::SLHDSA256s {
                    sig,
                })
            },

            // Ed25519Legacy.
            PublicKeyAlgorithm::EdDSA => {
                let key_info = key_info.check_standalone(
                    &self.openpgp_key, composite_info, KeyType::EC_EDWARDS)?;

                // RFC 9580, Section 12.7:
                //
                // [W]hile [RFC8032] describes different variants of EdDSA,
                // OpenPGP uses the "pure" variant (PureEdDSA). The hashing
                // that happens with OpenPGP is done as part of the standard
                // OpenPGP signature process, and that hash itself is fed as
                // the input message to the PureEdDSA algorithm.
                //
                // https://www.rfc-editor.org/rfc/rfc9580.html#section-12.7
                let eddsa_params = EddsaParams::new(EddsaSignatureScheme::Pure);

                let raw = session.sign(&Mechanism::Eddsa(eddsa_params),
                                       key_info.private_key.clone(), digest)
                    .with_context(|| {
                        format!("Signing data")
                    })?;

                if raw.len() != 64 {
                    return Err(anyhow::anyhow!(
                        "Token returned an invalid signature: expected \
                         64 bytes of signature data, got {} bytes",
                        raw.len()));
                }

                let r = mpi::MPI::new(&raw[..32]);
                let s = mpi::MPI::new(&raw[32..]);
                Ok(mpi::Signature::EdDSA { r, s, })
            },

            #[allow(deprecated)]
            PublicKeyAlgorithm::RSASign | PublicKeyAlgorithm::RSAEncryptSign => {
                let key_info = key_info.check_standalone(
                    &self.openpgp_key, composite_info, KeyType::RSA)?;

                // We use the raw RSA signing mechanism because we
                // handle the hashing in the OpenPGP layer.  But, that
                // means we have to do part of EMSA-PKCS1-v1_5
                // ourselves.

                // Form the DigestInfo sequence.
                const MAX_OID_SIZE: usize = 20;
                let mut v = Vec::with_capacity(MAX_OID_SIZE + digest.len());
                v.extend(hash_algo.oid()?);
                v.extend(digest);

                // Then sign that.
                let raw = session.sign(&Mechanism::RsaPkcs,
                                       key_info.private_key.clone(),
                                       &v)?;

                Ok(mpi::Signature::RSA { s: mpi::MPI::new(&raw[..]) })
            },

            algo => {
                Err(anyhow::anyhow!("{} keys are not supported", algo))
            },
        }
    }
}
