//! OpenPGP Shard functionality.

use std::{
    collections::HashMap,
    io::{Read, Write},
    path::Path,
    str::FromStr,
};

use keyfork_derive_openpgp::{
    derive_util::{DerivationPath, VariableLengthSeed},
    XPrv,
};
use openpgp::{
    armor::{Kind, Writer},
    cert::{Cert, CertParser, ValidCert},
    packet::{Packet, Tag, UserID, PKESK, SEIP},
    parse::{
        stream::{DecryptionHelper, DecryptorBuilder, VerificationHelper},
        Parse,
    },
    policy::{NullPolicy, Policy, StandardPolicy},
    serialize::{
        stream::{ArbitraryWriter, Encryptor2, LiteralWriter, Message, Recipient, Signer},
        Marshal,
    },
    types::KeyFlags,
    KeyID, PacketPile,
};
pub use sequoia_openpgp as openpgp;
use sharks::Share;

mod keyring;
use keyring::Keyring;

mod smartcard;
use smartcard::SmartcardManager;

/// Shard metadata verson 1:
/// 1 byte: Version
/// 1 byte: Threshold
/// Packet Pile of Certs
const SHARD_METADATA_VERSION: u8 = 1;
const SHARD_METADATA_OFFSET: usize = 2;

use super::{Format, KeyDiscovery, SharksError};

/// Errors encountered while performing operations using OpenPGP.
#[derive(Debug, thiserror::Error)]
pub enum Error {
    /// Errors encountered while creating or combining shares.
    #[error("{0}")]
    Sharks(#[from] SharksError),

    /// An error occurred while performing an OpenPGP operation.
    #[error("OpenPGP error: {0}")]
    Sequoia(#[source] anyhow::Error),

    /// An IO error occurred while performing an OpenPGP operation.
    #[error("OpenPGP IO error: {0}")]
    SequoiaIo(#[source] std::io::Error),

    /// An error occurred while using a keyring.
    #[error("Keyring error: {0}")]
    Keyring(#[from] keyring::Error),

    /// An error occurred while using a smartcard.
    #[error("Smartcard error: {0}")]
    Smartcard(#[from] smartcard::Error),

    /// An IO error occurred.
    #[error("IO error: {0}")]
    Io(#[source] std::io::Error),
}

#[allow(missing_docs)]
pub type Result<T, E = Error> = std::result::Result<T, E>;

/// An OpenPGP encrypted message and public-key-encrypted-secret-key packets.
#[derive(Debug, Clone)]
pub struct EncryptedMessage {
    pkesks: Vec<PKESK>,
    message: SEIP,
}

impl EncryptedMessage {
    /// Create a new EncryptedMessage from known parts.
    pub fn new(pkesks: &mut Vec<PKESK>, seip: SEIP) -> Self {
        Self {
            pkesks: std::mem::take(pkesks),
            message: seip,
        }
    }

    /// Parse OpenPGP packets for encrypted messages.
    ///
    /// # Errors
    /// The function may return an error if Sequoia is unable to parse packets.
    ///
    /// # Panics
    /// The function may panic if an unexpected packet is encountered.
    pub fn from_reader(input: impl Read + Send + Sync) -> openpgp::Result<Vec<Self>> {
        let mut pkesks = Vec::new();
        let mut encrypted_messages = vec![];

        for packet in PacketPile::from_reader(input)
            .map_err(Error::Sequoia)?
            .into_children()
        {
            match packet {
                Packet::PKESK(p) => pkesks.push(p),
                Packet::SEIP(s) => {
                    encrypted_messages.push(EncryptedMessage::new(&mut pkesks, s));
                }
                s => {
                    panic!("Invalid variant found: {}", s.tag());
                }
            }
        }

        Ok(encrypted_messages)
    }

    /// Serialize all contents of the message to a writer.
    ///
    /// # Errors
    /// The function may error for any condition in Sequoia's Serialize trait.
    fn serialize(&self, mut o: impl std::io::Write + Send + Sync) -> openpgp::Result<()> {
        for pkesk in &self.pkesks {
            let mut packet = vec![];
            pkesk.serialize(&mut packet).map_err(Error::Sequoia)?;
            let message = Message::new(&mut o);
            let mut message = ArbitraryWriter::new(message, Tag::PKESK).map_err(Error::Sequoia)?;
            message.write_all(&packet).map_err(Error::SequoiaIo)?;
            message.finalize().map_err(Error::Sequoia)?;
        }
        let mut packet = vec![];
        self.message
            .serialize(&mut packet)
            .map_err(Error::Sequoia)?;

        let message = Message::new(&mut o);
        let mut message = ArbitraryWriter::new(message, Tag::SEIP).map_err(Error::Sequoia)?;
        message.write_all(&packet).map_err(Error::SequoiaIo)?;
        message.finalize().map_err(Error::Sequoia)?;

        Ok(())
    }

    /// Decrypt the message with a Sequoia policy and decryptor.
    ///
    /// This method creates a container containing the packets and passes the serialized container
    /// to a DecryptorBuilder, which is used to decrypt the message.
    ///
    /// # Errors
    /// The method may return an error if it is unable to rebuild the message to decrypt or if it
    /// is unable to decrypt the message.
    pub fn decrypt_with<H>(&self, policy: &'_ dyn Policy, decryptor: H) -> Result<Vec<u8>>
    where
        H: VerificationHelper + DecryptionHelper,
    {
        let mut packets = vec![];
        self.serialize(&mut packets).map_err(Error::Sequoia)?;

        let mut decryptor = DecryptorBuilder::from_bytes(&packets)
            .map_err(Error::Sequoia)?
            .with_policy(policy, None, decryptor)
            .map_err(Error::Sequoia)?;

        let mut content = vec![];
        decryptor
            .read_to_end(&mut content)
            .map_err(Error::SequoiaIo)?;
        Ok(content)
    }
}

///
pub struct OpenPGP;

impl OpenPGP {
    /// Read all OpenPGP certificates in a path and return a [`Vec`] of them. Certificates are read
    /// from a file, or from files one level deep in a directory.
    ///
    /// # Errors
    /// The function may return an error if it is unable to read the directory or if Sequoia is unable
    /// to load certificates from the file.
    pub fn discover_certs(path: impl AsRef<Path>) -> Result<Vec<Cert>> {
        let path = path.as_ref();

        if path.is_file() {
            let mut vec = vec![];
            for cert in CertParser::from_file(path).map_err(Error::Sequoia)? {
                vec.push(cert.map_err(Error::Sequoia)?);
            }
            Ok(vec)
        } else {
            let mut vec = vec![];
            for entry in path
                .read_dir()
                .map_err(Error::Io)?
                .filter_map(Result::ok)
                .filter(|p| p.path().is_file())
            {
                vec.push(Cert::from_file(entry.path()).map_err(Error::Sequoia)?);
            }
            Ok(vec)
        }
    }
}

impl Format for OpenPGP {
    type Error = Error;
    type PublicKey = Cert;
    type PrivateKeyData = Vec<Cert>;
    type SigningKey = Cert;
    type EncryptedData = EncryptedMessage;

    /// Derive an OpenPGP Shard certificate from the given seed.
    fn derive_signing_key(&self, seed: &[u8]) -> Self::SigningKey {
        let seed = VariableLengthSeed::new(seed);
        // build cert to sign encrypted shares
        let userid = UserID::from("keyfork-sss");
        let path = DerivationPath::from_str("m/7366512'/0'").expect("valid derivation path");
        let xprv = XPrv::new(seed)
            .derive_path(&path)
            .expect("valid derivation");
        keyfork_derive_openpgp::derive(
            xprv,
            &[KeyFlags::empty().set_certification().set_signing()],
            &userid,
        )
        .expect("valid cert creation")
    }

    fn format_encrypted_header(
        &self,
        signing_key: &Self::SigningKey,
        key_data: &[Self::PublicKey],
        threshold: u8,
    ) -> Result<Self::EncryptedData, Self::Error> {
        let policy = StandardPolicy::new();
        let mut pp = vec![SHARD_METADATA_VERSION, threshold];
        // Note: Sequoia does not export private keys on a Cert, only on a TSK
        signing_key
            .serialize(&mut pp)
            .expect("serialize cert into bytes");
        for cert in key_data {
            cert.serialize(&mut pp)
                .expect("serialize pubkey into bytes");
        }

        // verify packet pile
        let mut iter = openpgp::cert::CertParser::from_bytes(&pp[SHARD_METADATA_OFFSET..])
            .expect("should have certs");
        let first_cert = iter.next().transpose().ok().flatten().expect("first cert");
        assert_eq!(signing_key, &first_cert);

        for (packet_cert, cert) in iter.zip(key_data) {
            assert_eq!(
                &packet_cert.expect("parsed packet cert"),
                cert,
                "packet pile could not recreate cert: {}",
                cert.fingerprint(),
            );
        }

        let valid_certs = key_data
            .iter()
            .map(|c| c.with_policy(&policy, None))
            .collect::<openpgp::Result<Vec<_>>>()
            .map_err(Error::Sequoia)?;
        let recipients = valid_certs.iter().flat_map(|vc| {
            get_encryption_keys(vc).map(|key| Recipient::new(KeyID::wildcard(), key.key()))
        });

        // Process is as follows:
        // * Any OpenPGP message
        // * An encrypted message
        // * A literal message
        // * The packet pile
        //
        // When decrypting, OpenPGP will see:
        // * A message, and parse it
        // * An encrypted message, and decrypt it
        // * A literal message, and extract it
        // * The packet pile
        let mut output = vec![];
        let message = Message::new(&mut output);
        let encrypted_message = Encryptor2::for_recipients(message, recipients)
            .build()
            .map_err(Error::Sequoia)?;
        let mut literal_message = LiteralWriter::new(encrypted_message)
            .build()
            .map_err(Error::Sequoia)?;
        literal_message.write_all(&pp).map_err(Error::SequoiaIo)?;
        literal_message.finalize().map_err(Error::Sequoia)?;

        // Parse it into an EncryptedMessage. Yes, this takes a serialized message
        // and deserializes it. Don't think about it too hard. It's easier this way.

        let mut pkesks = vec![];
        for packet in PacketPile::from_reader(output.as_slice())
            .map_err(Error::Sequoia)?
            .into_children()
        {
            match packet {
                Packet::PKESK(p) => pkesks.push(p),
                Packet::SEIP(s) => return Ok(EncryptedMessage::new(&mut pkesks, s)),
                s => panic!("Invalid variant found: {}", s.tag()),
            }
        }

        panic!("Unable to build EncryptedMessage from PacketPile");
    }

    fn encrypt_shard(
        &self,
        shard: &[u8],
        public_key: &Cert,
        signing_key: &mut Self::SigningKey,
    ) -> Result<EncryptedMessage> {
        let policy = StandardPolicy::new();
        let valid_cert = public_key
            .with_policy(&policy, None)
            .map_err(Error::Sequoia)?;
        let encryption_keys = get_encryption_keys(&valid_cert).collect::<Vec<_>>();

        let signing_key = signing_key
            .primary_key()
            .parts_into_secret()
            .map_err(Error::Sequoia)?
            .key()
            .clone()
            .into_keypair()
            .map_err(Error::Sequoia)?;

        // Process is as follows:
        // * Any OpenPGP message
        // * An encrypted message
        // * A signed message
        // * A literal message
        // * The shard itself
        //
        // When decrypting, OpenPGP will see:
        // * A message, and parse it
        // * An encrypted message, and decrypt it
        // * A signed message, and verify it
        // * A literal message, and extract it
        // * The shard itself
        let mut message_output = vec![];
        let message = Message::new(&mut message_output);
        let encrypted_message = Encryptor2::for_recipients(
            message,
            encryption_keys
                .iter()
                .map(|k| Recipient::new(KeyID::wildcard(), k.key())),
        )
        .build()
        .map_err(Error::Sequoia)?;
        let signed_message = Signer::new(encrypted_message, signing_key)
            .build()
            .map_err(Error::Sequoia)?;
        let mut message = LiteralWriter::new(signed_message)
            .build()
            .map_err(Error::Sequoia)?;
        message.write_all(shard).map_err(Error::SequoiaIo)?;
        message.finalize().map_err(Error::Sequoia)?;

        let message = EncryptedMessage::from_reader(message_output.as_slice())
            .map_err(Error::Sequoia)?
            .into_iter()
            .next()
            .expect("serialized message should be parseable");

        Ok(message)
    }

    fn parse_shard_file(
        &self,
        shard_file: impl Read + Send + Sync,
    ) -> Result<Vec<Self::EncryptedData>, Self::Error> {
        EncryptedMessage::from_reader(shard_file).map_err(Error::Sequoia)
    }

    fn format_shard_file(
        &self,
        encrypted_data: &[Self::EncryptedData],
        shard_file: impl Write + Send + Sync,
    ) -> Result<(), Self::Error> {
        let mut writer = Writer::new(shard_file, Kind::Message).map_err(Error::SequoiaIo)?;
        for message in encrypted_data {
            message.serialize(&mut writer).map_err(Error::Sequoia)?;
        }
        writer.finalize().map_err(Error::SequoiaIo)?;
        Ok(())
    }

    fn decrypt_all_shards(
        &self,
        private_keys: Option<Self::PrivateKeyData>,
        encrypted_data: &[Self::EncryptedData],
    ) -> std::result::Result<(Vec<Share>, u8), Self::Error> {
        // Be as liberal as possible when decrypting.
        // We don't want to invalidate someone's keys just because the old sig expired.
        let policy = NullPolicy::new();
        let mut keyring = Keyring::new(private_keys.unwrap_or_default())?;
        let mut manager = SmartcardManager::new()?;

        let mut encrypted_messages = encrypted_data.iter();

        let metadata = encrypted_messages.next().expect("metdata");
        let metadata_content = decrypt_metadata(metadata, &policy, &mut keyring, &mut manager)?;

        let (threshold, root_cert, certs) = decode_metadata_v1(&metadata_content)?;

        keyring.set_root_cert(root_cert.clone());
        manager.set_root_cert(root_cert.clone());

        // Generate a controlled binding from certificates to encrypted messages. This is stable
        // because we control the order packets are encrypted and certificates are stored.

        // TODO: remove alloc, convert EncryptedMessage to &EncryptedMessage
        let mut messages: HashMap<KeyID, EncryptedMessage> = certs
            .iter()
            .map(Cert::keyid)
            .zip(encrypted_messages.cloned())
            .collect();
        let mut decrypted_messages =
            decrypt_with_keyring(&mut messages, &certs, &policy, &mut keyring)?;

        // clean decrypted messages from encrypted messages
        messages.retain(|k, _v| !decrypted_messages.contains_key(k));

        let left_from_threshold = threshold as usize - decrypted_messages.len();
        if left_from_threshold > 0 {
            #[allow(clippy::cast_possible_truncation)]
            let new_messages = decrypt_with_manager(
                left_from_threshold as u8,
                &mut messages,
                &certs,
                &policy,
                &mut manager,
            )?;
            decrypted_messages.extend(new_messages);
        }

        let shares = decrypted_messages
            .values()
            .map(|message| Share::try_from(message.as_slice()))
            .collect::<Result<Vec<_>, &str>>()
            .map_err(|e| SharksError::Share(e.to_string()))?;
        Ok((shares, threshold))
    }

    fn decrypt_one_shard(
        &self,
        private_keys: Option<Self::PrivateKeyData>,
        encrypted_data: &[Self::EncryptedData],
    ) -> std::result::Result<(Share, u8), Self::Error> {
        let policy = NullPolicy::new();
        let mut keyring = Keyring::new(private_keys.unwrap_or_default())?;
        let mut manager = SmartcardManager::new()?;

        let mut encrypted_messages = encrypted_data.iter();

        let metadata = encrypted_messages.next().expect("metadata");
        let metadata_content = decrypt_metadata(metadata, &policy, &mut keyring, &mut manager)?;

        let (threshold, root_cert, certs) = decode_metadata_v1(&metadata_content)?;

        keyring.set_root_cert(root_cert.clone());
        manager.set_root_cert(root_cert.clone());
        let mut messages: HashMap<KeyID, EncryptedMessage> = certs
            .iter()
            .map(Cert::keyid)
            .zip(encrypted_messages.cloned())
            .collect();

        let decrypted_messages =
            decrypt_with_keyring(&mut messages, &certs, &policy, &mut keyring)?;

        if let Some(message) = decrypted_messages.into_values().next() {
            let share = Share::try_from(message.as_slice())
                .map_err(|e| SharksError::Share(e.to_string()))?;
            return Ok((share, threshold));
        }

        let decrypted_messages =
            decrypt_with_manager(1, &mut messages, &certs, &policy, &mut manager)?;

        if let Some(message) = decrypted_messages.into_values().next() {
            let share = Share::try_from(message.as_slice())
                .map_err(|e| SharksError::Share(e.to_string()))?;
            return Ok((share, threshold));
        }

        panic!("unable to decrypt shard");
    }
}

impl KeyDiscovery<OpenPGP> for &Path {
    fn discover_public_keys(&self) -> Result<Vec<<OpenPGP as Format>::PublicKey>> {
        OpenPGP::discover_certs(self)
    }

    fn discover_private_keys(&self) -> Result<<OpenPGP as Format>::PrivateKeyData> {
        todo!()
    }
}

impl KeyDiscovery<OpenPGP> for &[Cert] {
    fn discover_public_keys(&self) -> Result<Vec<<OpenPGP as Format>::PublicKey>> {
        Ok(self.to_vec())
    }

    fn discover_private_keys(&self) -> Result<<OpenPGP as Format>::PrivateKeyData> {
        Ok(self.to_vec())
    }
}

fn get_encryption_keys<'a>(
    cert: &'a ValidCert,
) -> openpgp::cert::prelude::ValidKeyAmalgamationIter<
    'a,
    openpgp::packet::key::PublicParts,
    openpgp::packet::key::UnspecifiedRole,
> {
    cert.keys()
        .alive()
        .revoked(false)
        .supported()
        .for_storage_encryption()
}

fn get_decryption_keys<'a>(
    cert: &'a ValidCert,
) -> openpgp::cert::prelude::ValidKeyAmalgamationIter<
    'a,
    openpgp::packet::key::SecretParts,
    openpgp::packet::key::UnspecifiedRole,
> {
    cert.keys()
        /*
        .alive()
        .revoked(false)
        .supported()
        */
        .for_storage_encryption()
        .secret()
}

fn decode_metadata_v1(buf: &[u8]) -> Result<(u8, Cert, Vec<Cert>)> {
    assert_eq!(
        SHARD_METADATA_VERSION, buf[0],
        "Incompatible metadata version"
    );
    let threshold = buf[1];

    let mut cert_parser =
        CertParser::from_bytes(&buf[SHARD_METADATA_OFFSET..]).map_err(Error::Sequoia)?;

    let root_cert = match cert_parser.next() {
        Some(Ok(c)) => c,
        Some(Err(e)) => return Err(Error::Sequoia(e)),
        None => panic!("No data found"),
    };

    let certs = cert_parser
        .collect::<openpgp::Result<Vec<_>>>()
        .map_err(Error::Sequoia)?;

    Ok((threshold, root_cert, certs))
}

// NOTE: When using single-decryptor mechanism, use this method with `threshold = 1` to return a
// single message.
fn decrypt_with_manager(
    threshold: u8,
    messages: &mut HashMap<KeyID, EncryptedMessage>,
    certs: &[Cert],
    policy: &dyn Policy,
    manager: &mut SmartcardManager,
) -> Result<HashMap<KeyID, Vec<u8>>> {
    let mut decrypted_messages = HashMap::new();

    while threshold as usize - decrypted_messages.len() > 0 {
        // Build list of fingerprints that haven't yet been used for decrypting
        let mut cert_by_fingerprint = HashMap::new();
        let mut unused_fingerprints = vec![];
        for valid_cert in certs
            .iter()
            .filter(|cert| !decrypted_messages.contains_key(&cert.keyid()))
            .map(|cert| cert.with_policy(policy, None))
        {
            let valid_cert = valid_cert.map_err(Error::Sequoia)?;
            let fp = valid_cert
                .keys()
                .for_storage_encryption()
                .map(|k| k.fingerprint())
                .collect::<Vec<_>>();
            for fp in &fp {
                cert_by_fingerprint.insert(fp.clone(), valid_cert.keyid());
            }
            unused_fingerprints.extend(fp.into_iter());
        }

        // Iterate over all fingerprints and use key_by_fingerprints to assoc with Enc. Message
        if let Some(fp) = manager.load_any_fingerprint(unused_fingerprints)? {
            let cert_keyid = cert_by_fingerprint.get(&fp).unwrap().clone();
            if let Some(message) = messages.remove(&cert_keyid) {
                let message = message.decrypt_with(policy, &mut *manager)?;
                decrypted_messages.insert(cert_keyid, message);
            }
        }
    }

    Ok(decrypted_messages)
}

// NOTE: When using single-decryptor mechanism, only a single key should be provided in Keyring to
// decrypt messages with.
fn decrypt_with_keyring(
    messages: &mut HashMap<KeyID, EncryptedMessage>,
    certs: &[Cert],
    policy: &NullPolicy,
    keyring: &mut Keyring,
) -> Result<HashMap<KeyID, Vec<u8>>, Error> {
    let mut decrypted_messages = HashMap::new();

    for valid_cert in certs.iter().map(|cert| cert.with_policy(policy, None)) {
        let valid_cert = valid_cert.map_err(Error::Sequoia)?;
        let Some(secret_cert) = keyring.get_cert_for_primary_keyid(&valid_cert.keyid()) else {
            continue;
        };
        let secret_cert = secret_cert
            .with_policy(policy, None)
            .map_err(Error::Sequoia)?;
        let keys = get_decryption_keys(&secret_cert).collect::<Vec<_>>();
        if !keys.is_empty() {
            if let Some(message) = messages.get_mut(&valid_cert.keyid()) {
                for (pkesk, key) in message.pkesks.iter_mut().zip(keys) {
                    pkesk.set_recipient(key.keyid());
                }
                // we have a pkesk, decrypt via keyring
                decrypted_messages.insert(
                    valid_cert.keyid(),
                    message.decrypt_with(policy, &mut *keyring)?,
                );
            }
        }
    }

    Ok(decrypted_messages)
}

fn decrypt_metadata(
    message: &EncryptedMessage,
    policy: &NullPolicy,
    keyring: &mut Keyring,
    manager: &mut SmartcardManager,
) -> Result<Vec<u8>> {
    Ok(if keyring.is_empty() {
        manager.load_any_card()?;
        message.decrypt_with(policy, manager)?
    } else {
        message.decrypt_with(policy, keyring)?
    })
}