#![doc = include_str!("../README.md")]
#![allow(clippy::expect_fun_call)]

use std::{
    io::{stdin, stdout, Read, Write},
    sync::{Arc, Mutex},
};

use aes_gcm::{
    aead::{consts::U12, Aead},
    Aes256Gcm, KeyInit, Nonce,
};
use base64::prelude::{Engine, BASE64_STANDARD};
use hkdf::Hkdf;
use keyfork_bug::{bug, POISONED_MUTEX};
use keyfork_mnemonic_util::{English, Mnemonic};
use keyfork_prompt::{
    validators::{
        mnemonic::{MnemonicSetValidator, MnemonicValidator, WordLength},
        Validator,
    },
    Message as PromptMessage, PromptHandler, Terminal,
};
use sha2::Sha256;
use sharks::{Share, Sharks};
use x25519_dalek::{EphemeralSecret, PublicKey};

// 32-byte share, 1-byte index, 1-byte threshold, 1-byte version == 36 bytes
// Encrypted, is 52 bytes
const PLAINTEXT_LENGTH: u8 = 36;
const ENCRYPTED_LENGTH: u8 = PLAINTEXT_LENGTH + 16;

#[cfg(feature = "openpgp")]
pub mod openpgp;

/// A trait to specify where keys can be discovered from, such as a Rust-native type or a path on
/// the filesystem that keys may be read from.
pub trait KeyDiscovery<F: Format + ?Sized> {
    /// Discover public keys for the associated format.
    ///
    /// # Errors
    /// The method may return an error if public keys could not be loaded from the given discovery
    /// mechanism. A discovery mechanism _must_ be able to detect public keys.
    fn discover_public_keys(&self) -> Result<Vec<F::PublicKey>, F::Error>;

    /// Discover private keys for the associated format.
    ///
    /// # Errors
    /// The method may return an error if private keys could not be loaded from the given
    /// discovery mechanism. Keys may exist off-system (such as with smartcards), in which case the
    /// PrivateKeyData type of the asssociated format should be either `()` (if the keys may never
    /// exist on-system) or an empty container (such as an empty Vec); in either case, this method
    /// _must not_ return an error if keys are accessible but can't be transferred into memory.
    fn discover_private_keys(&self) -> Result<F::PrivateKeyData, F::Error>;
}

/// A format to use for splitting and combining secrets.
pub trait Format {
    /// The error type returned from any failed operations.
    type Error: std::error::Error + 'static;

    /// A type encapsulating a single public key recipient.
    type PublicKey;

    /// A type encapsulating the private key recipients of shards.
    type PrivateKeyData;

    /// A type representing a Signer derived from the secret.
    type SigningKey;

    /// A type representing the parsed, but encrypted, Shard data.
    type EncryptedData;

    /// Derive a signer
    fn derive_signing_key(&self, seed: &[u8]) -> Self::SigningKey;

    /// Format a header containing necessary metadata. Such metadata contains a version byte, a
    /// threshold byte, a public version of the [`Format::SigningKey`], and the public keys used to
    /// encrypt shards. The public keys must be kept _in order_ to the encrypted shards. Keyfork
    /// will use the same key_data for both, ensuring an iteration of this method will match with
    /// iterations in methods called later.
    ///
    /// # Errors
    /// The method may return an error if encryption to any of the public keys fails.
    fn format_encrypted_header(
        &self,
        signing_key: &Self::SigningKey,
        key_data: &[Self::PublicKey],
        threshold: u8,
    ) -> Result<Self::EncryptedData, Self::Error>;

    /// Format a shard encrypted to the given public key, signing with the private key.
    ///
    /// # Errors
    /// The method may return an error if the public key used to encrypt the shard is unsuitable
    /// for encryption, or if an error occurs while encrypting.
    fn encrypt_shard(
        &self,
        shard: &[u8],
        public_key: &Self::PublicKey,
        signing_key: &mut Self::SigningKey,
    ) -> Result<Self::EncryptedData, Self::Error>;

    /// Parse the Shard file into a processable type.
    ///
    /// # Errors
    /// The method may return an error if the Shard file could not be read from or if the Shard
    /// file could not be properly parsed.
    fn parse_shard_file(
        &self,
        shard_file: impl Read + Send + Sync,
    ) -> Result<Vec<Self::EncryptedData>, Self::Error>;

    /// Write the Shard data to a Shard file.
    ///
    /// # Errors
    /// The method may return an error if the Shard data could not be properly serialized or if the
    /// Shard file could not be written to.
    fn format_shard_file(
        &self,
        encrypted_data: &[Self::EncryptedData],
        shard_file: impl Write + Send + Sync,
    ) -> Result<(), Self::Error>;

    /// Decrypt shares and associated metadata from a readable input. For the current version of
    /// Keyfork, the only associated metadata is a u8 representing the threshold to combine
    /// secrets.
    ///
    /// # Errors
    /// The method may return an error if the shardfile couldn't be read from, if all shards
    /// could not be decrypted, or if a shard could not be parsed from the decrypted data.
    fn decrypt_all_shards(
        &self,
        private_keys: Option<Self::PrivateKeyData>,
        encrypted_messages: &[Self::EncryptedData],
        prompt: Arc<Mutex<impl PromptHandler>>,
    ) -> Result<(Vec<Share>, u8), Self::Error>;

    /// Decrypt a single share and associated metadata from a reaable input. For the current
    /// version of Keyfork, the only associated metadata is a u8 representing the threshold to
    /// combine secrets.
    ///
    /// # Errors
    /// The method may return an error if the shardfile couldn't be read from, if a shard could not
    /// be decrypted, or if a shard could not be parsed from the decrypted data.
    fn decrypt_one_shard(
        &self,
        private_keys: Option<Self::PrivateKeyData>,
        encrypted_data: &[Self::EncryptedData],
        prompt: Arc<Mutex<impl PromptHandler>>,
    ) -> Result<(Share, u8), Self::Error>;

    /// Decrypt multiple shares and combine them to recreate a secret.
    ///
    /// # Errors
    /// The method may return an error if the shares can't be decrypted or if the shares can't
    /// be combined into a secret.
    fn decrypt_all_shards_to_secret(
        &self,
        private_key_discovery: Option<impl KeyDiscovery<Self>>,
        reader: impl Read + Send + Sync,
        prompt: impl PromptHandler,
    ) -> Result<Vec<u8>, Box<dyn std::error::Error>> {
        let private_keys = private_key_discovery
            .map(|p| p.discover_private_keys())
            .transpose()?;
        let encrypted_messages = self.parse_shard_file(reader)?;
        let (shares, threshold) = self.decrypt_all_shards(
            private_keys,
            &encrypted_messages,
            Arc::new(Mutex::new(prompt)),
        )?;

        let secret = Sharks(threshold)
            .recover(&shares)
            .map_err(|e| SharksError::CombineShare(e.to_string()))?;

        Ok(secret)
    }

    /// Establish an AES-256-GCM transport key using ECDH, decrypt a single shard, and encrypt the
    /// shard to the AES key.
    ///
    /// # Errors
    /// The method may return an error if a share can't be decrypted. The method will not return an
    /// error if the camera is inaccessible or if a hardware error is encountered while scanning a
    /// QR code; instead, a mnemonic prompt will be used.
    fn decrypt_one_shard_for_transport(
        &self,
        private_key_discovery: Option<impl KeyDiscovery<Self>>,
        reader: impl Read + Send + Sync,
        prompt: impl PromptHandler,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let prompt = Arc::new(Mutex::new(prompt));

        // parse input
        let private_keys = private_key_discovery
            .map(|p| p.discover_private_keys())
            .transpose()?;
        let encrypted_messages = self.parse_shard_file(reader)?;

        // establish AES-256-GCM key via ECDH
        let mut pubkey_data: Option<[u8; 32]> = None;

        // receive remote data via scanning QR code from camera
        #[cfg(feature = "qrcode")]
        {
            prompt
                .lock()
                .expect(bug!(POISONED_MUTEX))
                .prompt_message(PromptMessage::Text(QRCODE_PROMPT.to_string()))?;
            if let Ok(Some(qrcode_content)) =
                keyfork_qrcode::scan_camera(std::time::Duration::from_secs(30), 0)
            {
                let decoded_data = BASE64_STANDARD
                    .decode(qrcode_content)
                    .expect(bug!("qrcode should contain base64 encoded data"));
                pubkey_data = Some(decoded_data.try_into().map_err(|_| InvalidData)?)
            } else {
                prompt
                    .lock()
                    .expect(bug!(POISONED_MUTEX))
                    .prompt_message(PromptMessage::Text(QRCODE_ERROR.to_string()))?;
            };
        }

        // if QR code scanning failed or was unavailable, read from a set of mnemonics
        let their_pubkey = match pubkey_data {
            Some(pubkey) => pubkey,
            None => {
                let validator = MnemonicValidator {
                    word_length: Some(WordLength::Count(24)),
                };
                prompt
                    .lock()
                    .expect(bug!(POISONED_MUTEX))
                    .prompt_validated_wordlist::<English, _>(
                        QRCODE_COULDNT_READ,
                        3,
                        validator.to_fn(),
                    )?
                    .as_bytes()
                    .try_into()
                    .map_err(|_| InvalidData)?
            }
        };

        // create our shared key
        let our_key = EphemeralSecret::random();
        let our_pubkey_mnemonic = Mnemonic::from_bytes(PublicKey::from(&our_key).as_bytes())?;
        let shared_secret = our_key.diffie_hellman(&PublicKey::from(their_pubkey));
        assert!(
            shared_secret.was_contributory(),
            bug!("shared secret might be insecure")
        );
        let hkdf = Hkdf::<Sha256>::new(None, shared_secret.as_bytes());

        let mut shared_key_data = [0u8; 256 / 8];
        hkdf.expand(b"key", &mut shared_key_data)?;
        let shared_key = Aes256Gcm::new_from_slice(&shared_key_data)?;

        let mut nonce_data = [0u8; 12];
        hkdf.expand(b"nonce", &mut nonce_data)?;
        let nonce = Nonce::<U12>::from_slice(&nonce_data);

        // decrypt a single shard and create the payload
        let (share, threshold) =
            self.decrypt_one_shard(private_keys, &encrypted_messages, prompt.clone())?;
        let mut payload = Vec::from(&share);
        payload.insert(0, HUNK_VERSION);
        payload.insert(1, threshold);
        assert!(
            payload.len() < PLAINTEXT_LENGTH as usize,
            "invalid share length (too long, max {PLAINTEXT_LENGTH} bytes)"
        );

        // convert plaintext to static-size payload
        #[allow(clippy::assertions_on_constants)]
        {
            assert!(PLAINTEXT_LENGTH < u8::MAX, "length byte can be u8");
        }

        // NOTE: Previous versions of Keyfork Shard would modify the padding bytes to avoid
        // duplicate mnemonic words. This version does not include that, and instead uses a
        // repeated length byte.
        #[allow(clippy::cast_possible_truncation)]
        let mut plaintext_bytes = [u8::try_from(payload.len()).expect(bug!(
            "previously asserted length must be < {PLAINTEXT_LENGTH}",
            PLAINTEXT_LENGTH = PLAINTEXT_LENGTH
        )); PLAINTEXT_LENGTH as usize];
        plaintext_bytes[..payload.len()].clone_from_slice(&payload);

        // encrypt data
        let encrypted_bytes = shared_key.encrypt(nonce, plaintext_bytes.as_slice())?;
        assert_eq!(
            encrypted_bytes.len(),
            ENCRYPTED_LENGTH as usize,
            bug!("encrypted bytes size != expected len"),
        );
        let mut mnemonic_bytes = [0u8; ENCRYPTED_LENGTH as usize];
        mnemonic_bytes.copy_from_slice(&encrypted_bytes);

        let payload_mnemonic = Mnemonic::from_nonstandard_bytes(mnemonic_bytes);

        #[cfg(feature = "qrcode")]
        {
            use keyfork_qrcode::{qrencode, ErrorCorrection};
            let mut qrcode_data = our_pubkey_mnemonic.to_bytes();
            qrcode_data.extend(payload_mnemonic.as_bytes());
            if let Ok(qrcode) = qrencode(
                &BASE64_STANDARD.encode(qrcode_data),
                ErrorCorrection::Highest,
            ) {
                prompt
                    .lock()
                    .expect(bug!(POISONED_MUTEX))
                    .prompt_message(PromptMessage::Text(
                        concat!(
                            "A QR code will be displayed after this prompt. ",
                            "Send the QR code back to the operator combining the shards. ",
                            "Nobody else should scan this QR code."
                        )
                        .to_string(),
                    ))?;
                prompt
                    .lock()
                    .expect(bug!(POISONED_MUTEX))
                    .prompt_message(PromptMessage::Data(qrcode))?;
            }
        }

        prompt
            .lock()
            .expect(bug!(POISONED_MUTEX))
            .prompt_message(PromptMessage::Text(format!(
            "Upon request, these words should be sent: {our_pubkey_mnemonic} {payload_mnemonic}"
        )))?;

        Ok(())
    }

    /// Split a secret into a shard for every shard in keys, with the given Shamir's Secret Sharing
    /// threshold.
    ///
    /// # Errors
    /// The method may return an error if the shares can't be encrypted.
    fn shard_and_encrypt(
        &self,
        threshold: u8,
        max: u8,
        secret: &[u8],
        public_key_discovery: impl KeyDiscovery<Self>,
        writer: impl Write + Send + Sync,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let mut signing_key = self.derive_signing_key(secret);

        let sharks = Sharks(threshold);
        let dealer = sharks.dealer(secret);

        let public_keys = public_key_discovery.discover_public_keys()?;
        assert!(
            public_keys.len() < u8::MAX as usize,
            "must have less than u8::MAX public keys"
        );
        assert_eq!(
            max,
            public_keys.len() as u8,
            "max must be equal to amount of public keys"
        );
        let max = public_keys.len() as u8;
        assert!(max >= threshold, "threshold must not exceed max keys");

        let header = self.format_encrypted_header(&signing_key, &public_keys, threshold)?;
        let mut messages = vec![header];
        for (pk, share) in public_keys.iter().zip(dealer) {
            let shard = Vec::from(&share);
            messages.push(self.encrypt_shard(&shard, pk, &mut signing_key)?);
        }

        self.format_shard_file(&messages, writer)?;

        Ok(())
    }
}

/// Errors encountered while creating or combining shares using Shamir's Secret Sharing.
#[derive(thiserror::Error, Debug)]
pub enum SharksError {
    /// A Shamir Share could not be created.
    #[error("Error creating share: {0}")]
    Share(String),

    /// The Shamir shares could not be combined.
    #[error("Error combining shares: {0}")]
    CombineShare(String),
}

/// The mnemonic or QR code used to transport an encrypted shard did not store the correct amount
/// of data.
#[derive(thiserror::Error, Debug)]
#[error("Mnemonic or QR code did not store enough data")]
pub struct InvalidData;

/// Decrypt hunk version 1:
/// 1 byte: Version
/// 1 byte: Threshold
/// Data: &[u8]
pub(crate) const HUNK_VERSION: u8 = 2;
pub(crate) const HUNK_OFFSET: usize = 2;

const QRCODE_PROMPT: &str = "Press enter, then present QR code to camera.";
const QRCODE_TIMEOUT: u64 = 60; // One minute
const QRCODE_COULDNT_READ: &str = "A QR code could not be scanned. Please enter their words: ";
const QRCODE_ERROR: &str = "Unable to scan a QR code. Falling back to text entry.";

/// Establish ECDH transport for remote operators, receive transport-encrypted shares, decrypt the
/// shares, and combine them.
///
/// # Errors
/// The function may error if:
/// * Prompting for transport-encrypted shards fails.
/// * Decrypting shards fails.
/// * Combining shards fails.
///
/// # Panics
/// The function may panic if it is given payloads generated using a version of Keyfork that is
/// incompatible with the currently running version.
pub fn remote_decrypt(w: &mut impl Write) -> Result<(), Box<dyn std::error::Error>> {
    let mut pm = Terminal::new(stdin(), stdout())?;

    let mut iter_count = None;
    let mut shares = vec![];

    let mut threshold = 0;
    let mut iter = 0;

    while iter_count.is_none() || iter_count.is_some_and(|i| i > 0) {
        iter += 1;
        let our_key = EphemeralSecret::random();
        let key_mnemonic = Mnemonic::from_bytes(PublicKey::from(&our_key).as_bytes())?;

        #[cfg(feature = "qrcode")]
        {
            use keyfork_qrcode::{qrencode, ErrorCorrection};
            let qrcode_data = key_mnemonic.to_bytes();
            if let Ok(qrcode) = qrencode(
                &BASE64_STANDARD.encode(qrcode_data),
                ErrorCorrection::Highest,
            ) {
                pm.prompt_message(PromptMessage::Text(format!(
                    concat!(
                        "QR code #{iter} will be displayed after this prompt. ",
                        "Send the QR code to the next shardholder. ",
                        "Only the next shardholder should scan the QR code."
                    ),
                    iter = iter
                )))?;
                pm.prompt_message(PromptMessage::Data(qrcode))?;
            }
        }

        pm.prompt_message(PromptMessage::Text(format!(
            concat!(
                "Upon request, these words should be sent to the shardholder: ",
                "{key_mnemonic}"
            ),
            key_mnemonic = key_mnemonic,
        )))?;

        let mut pubkey_data: Option<[u8; 32]> = None;
        let mut payload_data = None;

        #[cfg(feature = "qrcode")]
        {
            pm.prompt_message(PromptMessage::Text(QRCODE_PROMPT.to_string()))?;
            if let Ok(Some(qrcode_content)) =
                keyfork_qrcode::scan_camera(std::time::Duration::from_secs(QRCODE_TIMEOUT), 0)
            {
                let decoded_data = BASE64_STANDARD
                    .decode(qrcode_content)
                    .expect(bug!("qrcode should contain base64 encoded data"));
                assert_eq!(
                    decoded_data.len(),
                    ENCRYPTED_LENGTH as usize,
                    bug!("invalid payload data")
                );
                let _ = pubkey_data.insert(decoded_data[..32].try_into().map_err(|_| InvalidData)?);
                let _ = payload_data.insert(decoded_data[32..].to_vec());
            } else {
                pm.prompt_message(PromptMessage::Text(QRCODE_ERROR.to_string()))?;
            };
        }

        let (pubkey, payload) = match (pubkey_data, payload_data) {
            (Some(pubkey), Some(payload)) => (pubkey, payload),
            _ => {
                let validator = MnemonicSetValidator {
                    word_lengths: [24, 39],
                };

                let [pubkey_mnemonic, payload_mnemonic] = pm
                    .prompt_validated_wordlist::<English, _>(
                        QRCODE_COULDNT_READ,
                        3,
                        validator.to_fn(),
                    )?;
                let pubkey = pubkey_mnemonic
                    .as_bytes()
                    .try_into()
                    .map_err(|_| InvalidData)?;
                let payload = payload_mnemonic.to_bytes();
                (pubkey, payload)
            }
        };

        assert_eq!(
            payload.len(),
            ENCRYPTED_LENGTH as usize,
            bug!("invalid payload data")
        );

        let shared_secret = our_key.diffie_hellman(&PublicKey::from(pubkey));
        assert!(
            shared_secret.was_contributory(),
            bug!("shared secret might be insecure")
        );
        let hkdf = Hkdf::<Sha256>::new(None, shared_secret.as_bytes());

        let mut shared_key_data = [0u8; 256 / 8];
        hkdf.expand(b"key", &mut shared_key_data)?;
        let shared_key = Aes256Gcm::new_from_slice(&shared_key_data)?;

        let mut nonce_data = [0u8; 12];
        hkdf.expand(b"nonce", &mut nonce_data)?;
        let nonce = Nonce::<U12>::from_slice(&nonce_data);

        let payload = shared_key.decrypt(nonce, payload.as_slice())?;
        assert_eq!(HUNK_VERSION, payload[0], "Incompatible hunk version");

        match &mut iter_count {
            Some(n) => {
                // Must be > 0 to start loop, can't go lower
                *n -= 1;
            }
            None => {
                // NOTE: Should always be >= 1, < 256 due to Shamir constraints
                threshold = payload[1];
                let _ = iter_count.insert(threshold - 1);
            }
        }

        let payload_len = payload.last().expect(bug!("payload should not be empty"));
        shares.push(payload[HUNK_OFFSET..usize::from(*payload_len)].to_vec());
    }

    let shares = shares
        .into_iter()
        .map(|s| Share::try_from(s.as_slice()))
        .collect::<Result<Vec<_>, &str>>()
        .map_err(|e| SharksError::Share(e.to_string()))?;
    let secret = Sharks(threshold)
        .recover(&shares)
        .map_err(|e| SharksError::CombineShare(e.to_string()))?;

    /*
     * Verification would take up too much size, mnemonic would be very large
    let userid = UserID::from("keyfork-sss");
    let kdr = DerivationRequest::new(
        DerivationAlgorithm::Ed25519,
        &DerivationPath::from_str("m/7366512'/0'")?,
    )
    .derive_with_master_seed(secret.to_vec())?;
    let derived_cert = keyfork_derive_openpgp::derive(
        kdr,
        &[KeyFlags::empty().set_certification().set_signing()],
        userid,
    )?;

    // NOTE: Signatures on certs will be different. Compare fingerprints instead.
    let derived_fp = derived_cert.fingerprint();
    let expected_fp = root_cert.fingerprint();
    if derived_fp != expected_fp {
        return Err(Error::InvalidSecret(derived_fp, expected_fp));
    }
    */

    w.write_all(&secret)?;

    Ok(())
}