use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
use std::fs::File;
use std::io::Write;
use std::path::Path;
use std::sync::Arc;
use std::sync::Weak;
use std::time::SystemTime;

use anyhow::Context;

use futures::lock::Mutex;

use sequoia_openpgp as openpgp;
use openpgp::Cert;
use openpgp::Fingerprint;
use openpgp::Packet;
use openpgp::Result;
use openpgp::crypto::Password;
use openpgp::crypto::SessionKey;
use openpgp::crypto::mpi;
use openpgp::fmt::hex;
use openpgp::packet::key::Key4;
use openpgp::packet::key::PublicParts;
use openpgp::packet::key::UnspecifiedRole;
use openpgp::packet;
use openpgp::serialize::SerializeInto;
use openpgp::serialize::stream::Armorer;
use openpgp::serialize::stream::Message;
use openpgp::types::HashAlgorithm;
use openpgp::types::PublicKeyAlgorithm;

use sequoia_ipc::Keygrip;

use sequoia_tpm::Description;

use sequoia_keystore_backend as backend;
use backend::Error;
use backend::ImportStatus;
use backend::PasswordSource;
use backend::Protection;
use backend::utils::Directory;

mod certd;

#[derive(Clone)]
pub struct Backend {
    inner: Arc<Mutex<BackendInternal>>,
}

struct BackendInternal {
    home: Directory,

    // OpenPGP TPM don't store the OpenPGP key data structures nor
    // do they (always) have all the information required to
    // reconstruct them (in particular, it's missing the ECC algorithm
    // parameters).  Since looking up a key by subkey means doing a
    // full scan, we cache the results.
    certd: Arc<certd::CertD>,

    // One device per TPM.
    devices: Vec<Device>,
}

/// A Device exposes the (usable) keys managed by a particular TPM,
/// e.g., `device:/dev/tpmrm0`.
///
/// A key is usable if we have the OpenPGP data structure.  If we
/// don't have the OpenPGP key, we ignore the key.
#[derive(Clone)]
pub struct Device {
    // Device's URI.
    id: String,
    inner: Arc<Mutex<DeviceInternal>>
}

impl Device {
    /// Creates a new `WeakDevice` from this `Device`.
    fn downgrade(&self) -> WeakDevice {
        WeakDevice {
            id: self.id.clone(),
            inner: Arc::downgrade(&self.inner),
        }
    }
}

/// A `Device`, but with a weak reference to the data.
///
/// Before you use this, you need to upgrade this to a `Device`.
#[derive(Clone)]
pub struct WeakDevice {
    /// The device's URI, e.g., `device:/dev/tpmrm0`.
    id: String,
    inner: Weak<Mutex<DeviceInternal>>,
}

impl WeakDevice {
    /// Upgrades the `WeakDevice` to a `Device`, if possible.
    fn upgrade(&self) -> Option<Device> {
        Some(Device {
            id: self.id.clone(),
            inner: self.inner.upgrade()?,
        })
    }
}

struct DeviceInternal {
    /// The device's URI, e.g., `device:/dev/tpmrm0`.
    id: String,

    /// A map from id (fingerprint) to key.
    keys: BTreeMap<Fingerprint, Key>,
}

#[derive(Clone)]
pub struct Key {
    // The fingerprint is also the id.
    fpr: Fingerprint,
    inner: Arc<Mutex<KeyInternal>>,
}

/// A secret key.
struct KeyInternal {
    device: WeakDevice,

    fingerprint: Fingerprint,
    public_key: packet::Key<packet::key::PublicParts,
                            packet::key::UnspecifiedRole>,

    desc: Description,
}

impl Backend {
    /// Initializes the TPM backend.
    ///
    /// `home` is the directory where the backend will look for its
    /// configuration, e.g., `$HOME/.sq/keystore/tpm`.
    pub async fn init<P: AsRef<Path>>(home: P, _default: bool)
        -> Result<Self>
    {
        log::trace!("Backend::init");

        let home = Directory::from(home.as_ref());

        let certd = Arc::new(certd::CertD::new()?);

        let mut inner = BackendInternal {
            home,
            certd,
            devices: vec![],
        };

        let _ = inner.scan_internal(true).await;

        Ok(Backend {
            inner: Arc::new(Mutex::new(inner)),
        })
    }

    /// Initializes an ephemeral TPM backend.
    ///
    /// This is primarily useful for testing.
    pub async fn init_ephemeral() -> Result<Self> {
        log::trace!("Backend::init_ephemeral");

        let home = Directory::ephemeral()?;
        Self::init(home, false).await
    }
}

fn derive_key(pk: &sequoia_tpm::PublicKeyBytes)
    -> Result<packet::Key<PublicParts, UnspecifiedRole>>
{
    match pk {
        sequoia_tpm::PublicKeyBytes::RSA(pk) => {
            let bytes = hex::decode(&pk.bytes);
            match bytes {
                Ok(bytes) => {
                    Key4::import_public_rsa(
                        &[1, 0, 1], &bytes, SystemTime::UNIX_EPOCH)
                        .map(|k| packet::Key::V4(k))
                }
                Err(err) => {
                    Err(err.into())
                }
            }
        }
        _ => {
            return Err(anyhow::anyhow!("Unsupported algorithm"));
        }
    }
}

impl BackendInternal {
    async fn scan_internal(&mut self, _force: bool) -> Result<()> {
        log::trace!("Backend::scan");

        let certd = Arc::clone(&self.certd);

        // let home = &backend.home;
        let home = &self.home;

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

            let filename = entry.path();

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

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

            let file = match File::open(&filename) {
                Ok(file) => file,
                Err(err) => {
                    log::debug!("Failed to open {}: {}",
                                filename.display(), err);
                    continue;
                }
            };

            let mut desc: Description = match serde_yaml::from_reader(file) {
                Ok(desc) => desc,
                Err(err) => {
                    log::debug!("Failed to parse {}: {}",
                                filename.display(), err);
                    continue;
                }
            };

            log::trace!("{}:\n{}",
                        filename.display(),
                        serde_yaml::to_string(&desc)?);

            match sequoia_tpm::read_key(&mut desc.spec) {
                Ok(_key) => (),
                Err(err) => {
                    log::debug!("Error configured key ({}) on TPM: {}\n{}",
                                filename.display(), err,
                                serde_yaml::to_string(&desc)?);
                    continue;
                }
            }

            let pk = if let Some(pk) = desc.spec.provider.tpm.unique.as_ref() {
                pk
            } else {
                log::debug!("{}: No public key, skipping.",
                            filename.display());
                continue;
            };

            let key = match derive_key(pk) {
                Ok(key) => key,
                Err(err) => {
                    log::debug!("{}: Failed to derive a null OpenPGP key: {}",
                                filename.display(), err);
                    continue;
                }
            };

            let keygrip = match Keygrip::of(key.mpis()) {
                Ok(keygrip) => keygrip,
                Err(err) => {
                    log::debug!("{}: computing keygrip: {}",
                                filename.display(), err);
                    continue;
                }
            };

            let key = match certd.find(&keygrip).await {
                Ok(key) => key,
                Err(err) => {
                    log::debug!("{}: looking for OpenPGP key: {}",
                                filename.display(), err);

                    let packet = Packet::from(
                        key.role_into_primary());

                    let packet = packet.to_vec()
                        .expect("serializing to a vec is infallible");

                    let mut bytes = Vec::new();
                    let message = Message::new(&mut bytes);
                    let mut message = Armorer::new(message)
                        .kind(openpgp::armor::Kind::PublicKey)
                        .build()
                        .expect("can build armorer");
                    message.write_all(&packet)
                        .expect("can serialize to vec");
                    message.finalize()
                        .expect("can serialize to vec");

                    log::debug!("Null key:\n{}",
                                String::from_utf8_lossy(&bytes));
                    continue;
                }
            };

            let fpr = key.fingerprint();

            // Insert the key.
            let provider = desc.spec.provider.tpm.tcti.clone();
            let mut device = None;
            for d in self.devices.iter_mut() {
                if d.id == provider {
                    device = Some(d);
                }
            }

            let device = if let Some(device) = device {
                device
            } else {
                self.devices.push(Device {
                    id: provider.clone(),
                    inner: Arc::new(Mutex::new(DeviceInternal {
                        id: provider.clone(),
                        keys: Default::default(),
                    })),
                });
                self.devices.last_mut().unwrap()
            };

            let mut device_internal = device.inner.lock().await;
            match device_internal.keys.entry(fpr.clone()) {
                Entry::Occupied(_oe) => {
                    // Already have it.
                }
                Entry::Vacant(ve) => {
                    // Need to add it.
                    log::debug!("Found key {}", fpr);

                    let key = Key {
                        fpr: fpr.clone(),
                        inner: Arc::new(Mutex::new(KeyInternal {
                            device: device.downgrade(),
                            fingerprint: fpr,
                            // XXX
                            public_key: key,
                            desc,
                        })),
                    };

                    ve.insert(key);
                }
            }
        }

        Ok(())
    }
}

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

    async fn scan(&mut self) -> Result<()> {
        let mut backend = self.inner.lock().await;
        backend.scan_internal(false).await
    }

    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 mut backend = self.inner.lock().await;

        if let Err(err) = backend.scan_internal(false).await {
            log::debug!("Scanning TPMs: {}", err);
        }

        Box::new(
            backend.devices.iter()
                .map(|device| {
                    Box::new(device.clone())
                        as Box<dyn backend::DeviceHandle + Send + Sync>
                })
                .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 mut backend = self.inner.lock().await;

        // The first time through we look for the key without
        // scanning.  If we don't find it, then we rescan.
        for scan in [false, true] {
            if scan {
                log::trace!("Rescanning");
                if let Err(err) = backend.scan_internal(true).await {
                    log::debug!("Scanning TPMs: {}", err);
                }
            }

            for device in backend.devices.iter() {
                if device.id == id {
                    return Ok(Box::new(device.clone())
                              as Box<dyn backend::DeviceHandle + Send + Sync>);
                }
            }
        }

        Err(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 mut backend = self.inner.lock().await;

        // The first time through we look for the key without
        // scanning.  If we don't find it, then we rescan.
        for scan in [false, true] {
            if scan {
                log::trace!("Rescanning");
                if let Err(err) = backend.scan_internal(true).await {
                    log::debug!("Scanning TPMs: {}", err);
                }
            }

            for device in backend.devices.iter() {
                let device = device.inner.lock().await;

                for (key_id, key) in device.keys.iter() {
                    if &key_id.to_string() == id {
                        return Ok(Box::new(key.clone())
                                  as Box<dyn backend::KeyHandle + Send + Sync>);
                    }
                }
            }
        }

        Err(Error::NotFound(id.into()).into())
    }

    async fn import<'a>(&self, _cert: Cert)
        -> Result<Vec<(ImportStatus,
                       Box<dyn backend::KeyHandle + Send + Sync + 'a>)>>
    {
        log::trace!("Backend::import");

        Err(Error::ExternalImportRequired(Some(
            "Use an external tool to manage TPM keys.".into()).into()).into())
    }
}

#[async_trait::async_trait]
impl backend::DeviceHandle for Device {
    fn id(&self) -> String {
        log::trace!("Device::id");

        self.id.clone()
    }

    async fn available(&self) -> bool {
        log::trace!("Device::available");

        // XXX: Right now, we only support configured and available
        // keys.
        true
    }

    async fn configured(&self) -> bool {
        log::trace!("Device::configured");

        // XXX: Right now, we only support configured and available
        // keys.
        true
    }

    async fn registered(&self) -> bool {
        log::trace!("Device::registered");

        // XXX: Right now, we only support configured and available
        // keys.
        true
    }

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

        // We manage passwords at the key level; a device can't be
        // locked.

        Ok(())
    }

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

        let device = self.inner.lock().await;
        let keys = device.keys.values()
            .map(|key| {
                Box::new(key.clone())
                    as Box<dyn backend::KeyHandle + Send + Sync>
            })
            .collect::<Vec<_>>();

        Box::new(keys.into_iter())
    }
}

#[async_trait::async_trait]
impl backend::KeyHandle for Key {
    fn id(&self) -> String {
        log::trace!("Key::id");

        self.fpr.to_string()
    }

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

        self.fpr.clone()
    }

    async fn device<'a>(&self)
        -> Box<dyn backend::DeviceHandle + Send + Sync + 'a>
    {
        log::trace!("Key::device");

        let key_internal = self.inner.lock().await;

        // Get the device that the key is on.
        let device = match key_internal.device.upgrade() {
            Some(device) => device,
            None => panic!("Device disappeared"),
        };

        Box::new(device)
    }

    async fn available(&self) -> bool {
        log::trace!("Key::available");

        true
    }

    async fn locked(&self) -> Protection {
        log::trace!("Key::locked");

        // XXX: We currently only support keys where the password is
        // embedded in the specficiation file.
        Protection::Unlocked
    }

    async fn password_source(&self) -> PasswordSource {
        log::trace!("Key::password_source");

        PasswordSource::Inline
    }

    async fn decryption_capable(&self) -> bool {
        log::trace!("Key::decryption_capable");

        let key_internal = self.inner.lock().await;
        key_internal.desc.spec.capabilities.contains(
            &sequoia_tpm::Capability::Decrypt)
    }

    async fn signing_capable(&self) -> bool {
        log::trace!("Key::signing_capable");

        let key_internal = self.inner.lock().await;
        key_internal.desc.spec.capabilities.contains(
            &sequoia_tpm::Capability::Sign)
    }

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

        let key_internal = self.inner.lock().await;

        // We currently only support keys where the password is
        // embedded in the specficiation file.
        let _password = if let Some(password) = password {
            password
        } else {
            return Err(Error::NoExternalPassword(None).into());
        };

        log::trace!("KeyHandle::unlock({}): Already unlocked", self.fpr);
        Err(Error::AlreadyUnlocked(key_internal.fingerprint.to_string()).into())
    }

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

        Ok(())
    }

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

        let key = self.inner.lock().await;
        key.public_key.clone()
    }

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

        let key_internal = self.inner.lock().await;

        let pk_algo = key_internal.public_key.pk_algo();

        #[allow(deprecated)]
        let session_key = match (ciphertext, pk_algo) {
            (mpi::Ciphertext::RSA { c },
             PublicKeyAlgorithm::RSAEncryptSign
             | PublicKeyAlgorithm::RSAEncrypt
             | PublicKeyAlgorithm::RSASign) =>
            {
                let session_key = sequoia_tpm::decrypt(&key_internal.desc.spec,
                                                       c.value())?;
                SessionKey::from(&session_key[..])
            }
            _ => {
                return Err(Error::OperationNotSupported(
                    format!("Unsupported public key algorithm: {}", pk_algo)).into());
            }
        };

        Ok(session_key)
    }

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

        let key_internal = self.inner.lock().await;

        let pk_algo = key_internal.public_key.pk_algo();

        let hash_algo = match hash_algo {
            HashAlgorithm::SHA1 => sequoia_tpm::HashingAlgorithm::Sha1,
            HashAlgorithm::SHA256 => sequoia_tpm::HashingAlgorithm::Sha256,
            HashAlgorithm::SHA384 => sequoia_tpm::HashingAlgorithm::Sha384,
            HashAlgorithm::SHA512 => sequoia_tpm::HashingAlgorithm::Sha512,
            _ => {
                return Err(Error::OperationNotSupported(
                    format!("Unsupported hash algorithm: {}", hash_algo)).into());
            }
        };

        let sig = sequoia_tpm::sign(&key_internal.desc.spec, hash_algo, digest)?;
        #[allow(deprecated)]
        let sig = match pk_algo {
            PublicKeyAlgorithm::RSAEncryptSign
                | PublicKeyAlgorithm::RSAEncrypt
                | PublicKeyAlgorithm::RSASign =>
            {
                openpgp::crypto::mpi::Signature::RSA { s: sig.into() }
            }
            _ => {
                return Err(Error::OperationNotSupported(
                    format!("Unsupported public key algorithm: {}", pk_algo)).into());
            }
        };

        Ok((pk_algo, sig))
    }

    async fn export(&mut self)
        -> Result<openpgp::packet::Key<
            openpgp::packet::key::SecretParts,
            openpgp::packet::key::UnspecifiedRole>>
    {
        Err(Error::OperationNotSupported(
            "Use an external tool to manage TPM keys.".into()).into())
    }

    async fn change_password(&mut self, password: Option<&Password>)
        -> Result<()>
    {
        log::trace!("KeyHandle::change_password({}, {})",
                    self.fingerprint(),
                    if let Some(password) = password {
                        if password.map(|p| p.is_empty()) {
                            "clear password"
                        } else {
                            "set password"
                        }
                    } else {
                        "ask for password"
                    });

        Err(Error::OperationNotSupported(
            "Use an external tool to manage TPM keys.".into()).into())
    }

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

        Err(Error::OperationNotSupported(
            "Use an external tool to manage TPM keys.".into()).into())
    }
}

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

    use openpgp::KeyHandle;
    use openpgp::parse::Parse;
    use openpgp::policy::StandardPolicy;

    use backend::test_framework;

    use backend::Backend as _;

    // const TEST_DEVICE_ID: &str = "device:/dev/tpmrm0";

    fn get_test_device() -> Option<()>
    {
        // XXX: Connect to a tpm server.
        None
    }

    fn preinit() -> bool {
        // Don't run the tests if the TPM device is not available.
        get_test_device().is_some()
    }

    async fn init_backend() -> Backend {
        let backend = Backend::init_ephemeral().await.expect("can init backend");
        backend
    }

    async fn import_cert(backend: &mut Backend, cert: &Cert) {
        log::debug!("Importing {}", cert.fingerprint());

        // Get the internal device.
        let mut backend_internal = backend.inner.lock().await;
        let mut device = None;
        for d in backend_internal.devices.iter_mut() {
            if d.id == "internal" {
                device = Some(d);
                break;
            }
        }
        let _device = if let Some(device) = device {
            device
        } else {
            backend_internal.devices.push(Device {
                id: "internal".to_string(),
                inner: Arc::new(Mutex::new(DeviceInternal {
                    id: "internal".to_string(),
                    keys: Default::default(),
                })),
            });
            backend_internal.devices.last_mut().unwrap()
        };

        // XXX: Convert each secret key into a Specification.

        todo!("import_cert")
    }

    sequoia_keystore_backend::generate_tests!(
        preinit,
        true, // Serialize tests.
        Backend, init_backend,
        import_cert,
        false, // Can import encrypted secret key material.
        None, // Supported key sets.
        None, // Default password.
        false, // Can export.
        false, // Can change password.
        false // Can delete secret key material.
    );
}
