use std::collections::HashMap;
use std::collections::hash_map::Entry;
use std::sync::Arc;
use std::sync::Mutex;
use std::sync::MutexGuard;
use std::sync::Weak;
use std::ops::Deref;
use std::ops::DerefMut;
use std::path::Path;

use anyhow::Context;

use sequoia_openpgp as openpgp;
use openpgp::Cert;
use openpgp::cert::prelude::*;
use openpgp::crypto::Decryptor;
use openpgp::crypto::mpi;
use openpgp::crypto::Password;
use openpgp::crypto::SessionKey;
use openpgp::crypto::Signer;
use openpgp::Result;
use openpgp::Fingerprint;
use openpgp::packet;
use openpgp::parse::Parse;
use openpgp::types::HashAlgorithm;
use openpgp::types::PublicKeyAlgorithm;

use sequoia_keystore_backend as backend;
use backend::Backend as _;
use backend::KeyHandle as _;
use backend::Error;
use backend::utils::Directory;

/// Known soft keys.
#[derive(Clone)]
pub struct Backend(Arc<Mutex<BackendInternal>>);

struct BackendInternal {
    home: Directory,

    // Map from certificate fingerprint to Device.
    certs: HashMap<Fingerprint, Device>,

    // Map from key to certificate fingerprint.
    keys: HashMap<Fingerprint, Key>,
}

// If we do: backend.0.lock().unwrap().method(), then method doesn't
// have a reference to the `Arc`, which we sometimes need.  This
// conveniently encapsulates an `Arc` and a `MutexGuard`.
struct BackendRef<'a>(MutexGuard<'a, BackendInternal>, Arc<Mutex<BackendInternal>>);

impl<'a> Deref for BackendRef<'a> {
    type Target = MutexGuard<'a, BackendInternal>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<'a> DerefMut for BackendRef<'a> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl Backend {
    fn lock(&self) -> BackendRef {
        BackendRef(self.0.lock().unwrap(), self.0.clone())
    }
}

/// A Device corresponds to a certificate.
#[derive(Clone)]
pub struct Device(Arc<Mutex<DeviceInternal>>);

struct DeviceInternal {
    #[allow(dead_code)]
    backend: Weak<Mutex<BackendInternal>>,

    // The device's id (this is just the certificate's fingerprint).
    id: String,

    cert: Cert,

    keys: Vec<Weak<Mutex<KeyInternal>>>,
}

#[derive(Clone)]
pub struct Key(Arc<Mutex<KeyInternal>>);

/// A secret key.
struct KeyInternal {
    #[allow(dead_code)]
    backend: Weak<Mutex<BackendInternal>>,

    // The key's id (this is just the key's fingerprint).
    id: String,

    // The Key data structure.
    //
    // The first one is the variant as read from the TSK.  The second
    // is the unlocked variant (if any).
    variants: Vec<packet::Key<packet::key::SecretParts,
                              packet::key::UnspecifiedRole>>,

    // The unlocked variant, if any.
    unlocked: Option<packet::Key<packet::key::SecretParts,
                                 packet::key::UnspecifiedRole>>,
}

impl Backend {
    /// Initializes a SoftKeys backend.
    ///
    /// `home` is the directory where the backend will look for its
    /// configuration.  If this is `None`, then it will look in
    /// `$HOME/.sq/keystore/softkeys`.
    ///
    /// Currently, this backend looks for certificates containing
    /// secret key material in files with an 'asc' or 'pgp' extension.
    /// Other files are silently ignored.
    pub async fn init<P: AsRef<Path>>(home: Option<P>) -> Result<Box<Self>> {
        log::trace!("Backend::init");

        let home = if let Some(home) = home {
            home.as_ref().into()
        } else if let Some(home) = dirs::home_dir() {
            home.join(".sq").join("keystore").join("softkeys")
        } else {
            return Err(anyhow::anyhow!("Failed get the home directory"));
        };

        log::info!("Home directory is: {:?}", home);

        let mut backend = Backend(Arc::new(Mutex::new(BackendInternal {
            home: home.clone().into(),
            keys: Default::default(),
            certs: Default::default(),
        })));

        if let Err(err) = backend.scan().await {
            log::error!("Scanning {:?} for keys: {}", home, err);
        }

        Ok(Box::new(backend))
    }

    /// Initializes a SoftKeys backend.
    ///
    /// `home` is the directory where the backend will look for its
    /// configuration.  If this is `None`, then it will look in
    /// `$HOME/.sq/keystore/softkeys`.
    ///
    /// Currently, this backend looks for certificates containing
    /// secret key material in files with an 'asc' or 'pgp' extension.
    /// Other files are silently ignored.
    pub async fn init_ephemeral() -> Result<Self> {
        log::trace!("Backend::init_ephemeral");

        let home = Directory::ephemeral()?;

        log::info!("Home directory is: {:?}", home.display());

        let mut backend = Backend(Arc::new(Mutex::new(BackendInternal {
            home: home.clone().into(),
            keys: Default::default(),
            certs: Default::default(),
        })));

        if let Err(err) = backend.scan().await {
            log::error!("Scanning {:?} for keys: {}", home, err);
        }

        Ok(backend)
    }
}

impl BackendRef<'_> {
    /// Ingests a certificate.
    ///
    /// This merges the secret key of a certificate into the backend.
    /// If the certificate does not contain any secret key material,
    /// then an error is returned.
    pub fn ingest(&mut self, cert: Cert) -> Result<()> {
        log::trace!("BackendRef::ingest");

        if ! cert.is_tsk() {
            return Err(anyhow::anyhow!(
                "{} does not contain any secret key material",
                cert.fingerprint()));
        }

        // Add entries for the secret keys.
        let mut keys = Vec::with_capacity(cert.keys().secret().count());
        for ka in cert.keys().secret() {
            log::info!("Adding {} ({})", ka.keyid(), cert.keyid());

            match self.0.keys.entry(ka.fingerprint()) {
                Entry::Occupied(oe) => {
                    log::info!("{} ({}) already present",
                               ka.keyid(), cert.keyid());

                    keys.push(Arc::downgrade(&oe.get().0));

                    // We collect all secret key variants, but avoid
                    // duplicates.
                    let variants = &mut oe.get().0.lock().unwrap().variants;
                    if ! variants.contains(ka.key()) {
                        log::info!("{} ({}): new variant",
                                   ka.keyid(), cert.keyid());

                        variants.push(ka.key().clone());
                    }
                }
                Entry::Vacant(ve) => {
                    let sk = Key(Arc::new(Mutex::new(KeyInternal {
                        backend: Arc::downgrade(&self.1),
                        id: ka.fingerprint().to_hex(),
                        variants: vec![ ka.key().clone() ],
                        unlocked: None,
                    })));
                    keys.push(Arc::downgrade(&sk.0));
                    ve.insert(sk);
                }
            }
        }

        // Add an entry for the certificate.
        match self.0.certs.entry(cert.fingerprint()) {
            Entry::Occupied(oe) => {
                // Merge cert.
                let mut e = oe.get().0.lock().unwrap();

                e.cert = e.cert.clone().merge_public_and_secret(cert)
                    .expect("same certificate");

                e.keys.append(&mut keys);

                e.keys.sort_by_key(|k| {
                    k.upgrade().expect("valid").lock().unwrap().id.clone()
                });
                e.keys.dedup_by_key(|k| {
                    k.upgrade().expect("valid").lock().unwrap().id.clone()
                });
            }
            Entry::Vacant(ve) => {
                ve.insert(
                    Device(Arc::new(Mutex::new(DeviceInternal {
                        backend: Arc::downgrade(&self.1),
                        id: cert.fingerprint().to_hex(),
                        cert: cert,
                        keys,
                    }))));
            }
        }

        Ok(())
    }
}

#[async_trait::async_trait]
impl backend::Backend for Backend {
    fn id(&self) -> String {
        "softkeys".into()
    }

    async fn scan(&mut self) -> Result<()> {
        log::trace!("Backend::scan");

        let mut backend = self.lock();

        log::info!("Scanning: {:?}", backend.home.display());
        for entry in backend.home.read_dir()
            .with_context(|| format!("{:?}", backend.home.display()))?
        {
            let entry = match entry {
                Ok(entry) => entry,
                Err(err) => {
                    log::debug!("While listing {:?}: {}",
                                backend.home.display(), err);
                    continue;
                }
            };

            let filename = entry.path();

            let read = if let Some(extension) = filename.extension() {
                if extension == "asc" || extension == "pgp" {
                    true
                } else {
                    false
                }
            } else {
                false
            };

            if ! read {
                log::debug!("Ignoring {:?}", filename);
                continue;
            }
            log::debug!("Considering: {:?}", filename);

            let parser = CertParser::from_file(&filename)?;
            for certr in parser {
                let cert = match certr {
                    Ok(cert) => cert,
                    Err(err) => {
                        log::info!("{:?} partially corrupted: {}",
                                   filename, err);
                        continue;
                    }
                };

                log::debug!("Found certificate {} in {:?}",
                            cert.keyid(), filename);

                backend.ingest(cert)?;
            }
        }

        Ok(())
    }

    async fn list<'a>(&'a self)
        -> Box<dyn Iterator<Item=Box<dyn backend::DeviceHandle + Send + Sync + 'a>>
               + Send + Sync + 'a>
    {
        log::trace!("Backend::list");

        let backend = self.0.lock().unwrap();

        Box::new(
            backend.certs.values()
                .map(|sc| {
                    Box::new(sc.clone())
                        as Box<dyn backend::DeviceHandle + Send + Sync + 'a>
                })
                .collect::<Vec<_>>()
                .into_iter())
    }

    async fn find_device<'a>(&self, id: &str)
        -> Result<Box<dyn backend::DeviceHandle + Send + Sync + 'a>>
    {
        log::trace!("Backend::find_device");

        let backend = self.0.lock().unwrap();

        backend.certs.get(&id.parse()?)
            .map(|sc| {
                Box::new(sc.clone())
                    as Box<dyn backend::DeviceHandle + Send + Sync + 'a>
            })
            .ok_or_else(|| Error::NotFound(id.into()).into())
    }

    async fn find_key<'a>(&self, id: &str)
        -> Result<Box<dyn backend::KeyHandle + Send + Sync + 'a>>
    {
        log::trace!("Backend::find_key");

        let backend = self.lock();

        backend.keys.get(&id.parse()?)
            .map(|sk| {
                Box::new(sk.clone()) as Box<dyn backend::KeyHandle + Send + Sync + 'a>
            })
            .ok_or_else(|| Error::NotFound(id.into()).into())
    }
}

#[async_trait::async_trait]
impl backend::DeviceHandle for Device {
    fn id(&self) -> String {
        self.0.lock().unwrap().id.clone()
    }

    /// Certificates cannot be registered so if it exists, it is
    /// available.
    async fn available(&self) -> bool {
        true
    }

    /// Certificates cannot be unconfigured.
    async fn configured(&self) -> bool {
        true
    }

    /// Certificates cannot be registered.  (They are always available.)
    async fn registered(&self) -> bool {
        false
    }

    async fn lock(&mut self) -> Result<()> {
        log::trace!("DeviceHandle::lock");

        let mut dh = self.0.lock().unwrap();

        for k in dh.keys.iter_mut() {
            if let Some(kh) = k.upgrade() {
                let mut kh = Key(kh);
                if let Err(err) = kh.lock_sync() {
                    log::debug!("Error locking key {}: {}",
                                kh.id(), err);
                }
            }
        }

        Ok(())
    }

    async fn list<'a>(&'a self)
        -> Box<dyn Iterator<Item=Box<dyn backend::KeyHandle + Send + Sync + 'a>>
               + Send + Sync + 'a>
    {
        log::trace!("DeviceHandle::list");

        let cert = self.0.lock().unwrap();

        Box::new(
            cert
                .keys
                .iter()
                .filter_map(|k| {
                    Some(Box::new(Key(k.upgrade()?))
                         as Box<dyn sequoia_keystore_backend::KeyHandle
                                + Send + Sync>)
                })
                .collect::<Vec<_>>()
                .into_iter())
    }
}

// The traits make all functions async, but a number of
// implementations are sync in the sense that they don't block.  Using
// the async variants is not possible when holding a lock (as, e.g.,
// `MutexGuard` is not send).  Instead of doing acrobatics, we just
// provide sync versions of some functions.
impl Key {
    fn lock_sync(&mut self) -> Result<()> {
        log::trace!("KeyHandle::lock");

        let mut kh = self.0.lock().unwrap();
        kh.unlocked = None;
        Ok(())
    }
}

#[async_trait::async_trait]
impl backend::KeyHandle for Key {
    fn id(&self) -> String {
        self.0.lock().unwrap().id.clone()
    }

    fn fingerprint(&self) -> Fingerprint {
        log::trace!("KeyHandle::fingerprint");

        self.0.lock().unwrap().variants[0].fingerprint()
    }

    async fn available(&self) -> bool {
        true
    }

    async fn locked(&self) -> bool {
        let kh = self.0.lock().unwrap();
        let unlocked = kh.unlocked.is_some()
            || kh.variants.iter().any(|k| k.has_unencrypted_secret());
        ! unlocked
    }

    async fn decryption_capable(&self) -> bool {
        let kh = self.0.lock().unwrap();
        let key = &kh.variants[0];

        // XXX: Also look up the certificate in the cert store and
        // check the key flags.
        key.pk_algo().for_encryption()
    }

    async fn signing_capable(&self) -> bool {
        let kh = self.0.lock().unwrap();
        let key = &kh.variants[0];

        // XXX: Also look up the certificate in the cert store and
        // check the key flags.
        key.pk_algo().for_signing()
    }

    async fn unlock(&mut self, password: &Password) -> Result<()> {
        log::trace!("KeyHandle::unlock");

        let mut kh = self.0.lock().unwrap();

        if kh.unlocked.is_some() {
            return Err(Error::AlreadyUnlocked(kh.id.clone()).into());
        }

        let mut err = None;
        for v in kh.variants.iter() {
            let mut k = v.clone();
            match k.secret_mut().decrypt_in_place(v.pk_algo(), password) {
                Ok(()) => {
                    kh.unlocked = Some(k);
                    return Ok(());
                }
                Err(err_) => {
                    err = Some(err_);
                }
            }
        }

        Err(err.expect("at least one key variant"))
    }

    async fn lock(&mut self) -> Result<()> {
        self.lock_sync()
    }

    async fn public_key(&self)
        -> packet::Key<packet::key::PublicParts,
                       packet::key::UnspecifiedRole>
    {
        log::trace!("KeyHandle::public_key");

        let k = self.0.lock().unwrap().variants[0].clone();
        k.take_secret().0
    }

    async fn decrypt_ciphertext(&mut self,
                          ciphertext: &mpi::Ciphertext,
                          plaintext_len: Option<usize>)
        -> Result<SessionKey>
    {
        log::trace!("KeyHandle::decrypt_ciphertext");

        let kh = self.0.lock().unwrap();

        // We may have multiple variants.  We try them in turn unless
        // a key is already unlocked in which case we only try that
        // one.
        let variants: Box<dyn Iterator<Item=&packet::Key<_, _>>>
            = if let Some(k) = kh.unlocked.as_ref() {
                Box::new(std::iter::once(k))
            } else {
                Box::new(kh.variants.iter())
            };

        let mut err = None;
        for k in variants {
            if k.secret().is_encrypted() {
                if err.is_none() {
                    err = Some(Error::Locked(k.fingerprint().to_hex()).into());
                }
                continue;
            }

            // `into_keypair` fails if the key is encrypted, but we
            // already checked that it isn't.
            let mut keypair = k.clone().into_keypair().expect("decrypted");

            match keypair.decrypt(ciphertext, plaintext_len) {
                Ok(sk) => {
                    log::info!("Decrypted ciphertext using {}", k.keyid());
                    return Ok(sk);
                }
                Err(err_) => {
                    log::info!("Decrypting ciphertext using {}: {}",
                               k.keyid(), err_);
                    if err.is_none() {
                        err = Some(err_);
                    }
                }
            }
        }

        Err(err.expect("have key variants that failed"))
    }

    async fn sign(&mut self, hash_algo: HashAlgorithm, digest: &[u8])
        -> Result<(PublicKeyAlgorithm, mpi::Signature)>
    {
        log::trace!("KeyHandle::sign");

        let kh = self.0.lock().unwrap();

        // We may have multiple variants.  We try them in turn unless
        // a key is already unlocked in which case we only try that
        // one.
        let variants: Box<dyn Iterator<Item=&packet::Key<_, _>>>
            = if let Some(k) = kh.unlocked.as_ref() {
                Box::new(std::iter::once(k))
            } else {
                Box::new(kh.variants.iter())
            };

        // Remember the first error.
        let mut err = None;
        for k in variants {
            if k.secret().is_encrypted() {
                if err.is_none() {
                    err = Some(Error::Locked(k.fingerprint().to_hex()).into());
                }
                continue;
            }

            // `into_keypair` fails if the key is encrypted, but we
            // already checked that it isn't.
            let mut keypair = k.clone().into_keypair().expect("decrypted");

            match keypair.sign(hash_algo, digest) {
                Ok(sig) => return Ok((k.pk_algo(), sig)),
                Err(err_) => err = Some(err_),
            }
        }

        Err(err.expect("At least one key variant"))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use backend::test_framework;

    async fn init_backend() -> Backend {
        let backend = Backend::init_ephemeral().await
            .expect("creating an ephemeral backend");

        backend
    }

    async fn import_cert(backend: &mut Backend, cert: &Cert) {
        let mut backend = backend.lock();
        backend.ingest(cert.clone()).expect("can import");
    }

    sequoia_keystore_backend::generate_tests!(
        Backend, init_backend, import_cert);

    #[tokio::test]
    pub async fn decrypt_multiple_variants() -> Result<()> {
        // Load two variants of a certificate: one where the keys are
        // encrypted with the password "bob" and the other where they
        // are encrypted with the password "smithy".  Make sure we can
        // decrypt the message using either passwords.

        let mut backend = init_backend().await;

        let cert = Cert::from_bytes(
            test_framework::data::key("bob-password-bob.asc"))
            .expect("valid cert");
        import_cert(&mut backend, &cert).await;

        let cert = Cert::from_bytes(
            test_framework::data::key("bob-password-smithy.asc"))
            .expect("valid cert");
        import_cert(&mut backend, &cert).await;

        let msg: &[u8] = test_framework::data::message("bob.asc");
        let keyid: sequoia_openpgp::KeyID
            = "0A4BA2249168D779".parse().expect("valid");

        test_framework::try_decrypt(
            &mut backend, msg, Some(b"hi bob\n"),
            &keyid, Some(Password::from("bob"))).await.unwrap();
        test_framework::try_decrypt(
            &mut backend, msg, Some(b"hi bob\n"),
            &keyid, Some(Password::from("smithy"))).await.unwrap();

        assert!(test_framework::try_decrypt(
            &mut backend, msg, Some(b"hi bob\n"),
            &keyid, Some(Password::from("sup3r s3cur3"))).await.is_err());

        Ok(())
    }

    #[tokio::test]
    pub async fn decrypt_multiple_variants2() -> Result<()> {
        // Load two variants of a certificate: one where the keys are
        // encrypted with the password "bob" and the other where they
        // are not encrypted.  Make sure we can decrypt the message
        // with and without the password.

        let mut backend = init_backend().await;

        let cert = Cert::from_bytes(
            test_framework::data::key("bob-password-bob.asc"))
            .expect("valid cert");
        import_cert(&mut backend, &cert).await;

        let cert2 = Cert::from_bytes(
            test_framework::data::key("bob-no-password.asc"))
            .expect("valid cert");
        import_cert(&mut backend, &cert2).await;

        assert_eq!(cert.fingerprint(), cert2.fingerprint());

        let msg: &[u8] = test_framework::data::message("bob.asc");
        let keyid: sequoia_openpgp::KeyID
            = "0A4BA2249168D779".parse().expect("valid");

        eprintln!("Decrypting with bob (correct)");
        test_framework::try_decrypt(
            &mut backend, msg, Some(b"hi bob\n"),
            &keyid, Some(Password::from("bob"))).await.unwrap();

        eprintln!("Decrypting with sup3r s3cur3 (incorrect)");
        assert!(test_framework::try_decrypt(
            &mut backend, msg, Some(b"hi bob\n"),
            &keyid, Some(Password::from("sup3r s3cur3"))).await.is_err());

        Ok(())
    }

    #[tokio::test]
    pub async fn verify_with_passwords2() -> Result<()> {
        let _ = env_logger::Builder::from_default_env().try_init();

        // Sign a message using a key that is password protected.  To
        // make it trickier, use a key with multiple variants.

        let mut backend = init_backend().await;

        let cert = Cert::from_bytes(
            test_framework::data::key("bob-password-bob.asc"))
            .expect("valid cert");
        import_cert(&mut backend, &cert).await;
        let cert = Cert::from_bytes(
            test_framework::data::key("bob-password-smithy.asc"))
            .expect("valid cert");
        import_cert(&mut backend, &cert).await;

        let signing_key = "9410E96E68643821794608269CB5006116833EE7";

        let mut k = backend.find_key(signing_key).await?;
        k.unlock(&Password::from("bob")).await?;

        let signing_fpr = Fingerprint::from_hex(signing_key)
            .expect("valid");
        let signer = backend.find_key(signing_key).await?;

        test_framework::sign_verify(
            vec![ signer.into(), ],
            vec![ cert.clone() ],
            Some(&[ signing_fpr ][..]))?;

        k.lock().await?;
        k.unlock(&Password::from("smithy")).await?;

        let signing_fpr = Fingerprint::from_hex(signing_key)
            .expect("valid");
        let signer = backend.find_key(signing_key).await?;

        test_framework::sign_verify(
            vec![ signer.into(), ],
            vec![ cert.clone() ],
            Some(&[ signing_fpr ][..]))?;

        Ok(())
    }
}