#![doc = include_str!("../README.md")]

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

use aes_gcm::{
    aead::{consts::U12, Aead, AeadCore, OsRng},
    Aes256Gcm, KeyInit, Nonce,
};
use hkdf::Hkdf;
use keyfork_mnemonic_util::{English, Mnemonic};
use keyfork_prompt::{
    validators::{mnemonic::MnemonicSetValidator, Validator},
    Message as PromptMessage, PromptHandler, Terminal,
};
use sha2::Sha256;
use sharks::{Share, Sharks};
use x25519_dalek::{EphemeralSecret, PublicKey};

// 256 bit share encrypted is 49 bytes, couple more bytes before we reach max size
const ENC_LEN: u8 = 4 * 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],
    ) -> 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],
    ) -> 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,
    ) -> 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)?;

        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,
    ) -> Result<(), Box<dyn std::error::Error>> {
        let mut pm = Terminal::new(stdin(), stdout())?;

        // 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 nonce_data: Option<[u8; 12]> = None;
        let mut pubkey_data: Option<[u8; 32]> = None;

        // receive remote data via scanning QR code from camera
        #[cfg(feature = "qrcode")]
        {
            pm.prompt_message(PromptMessage::Text(QRCODE_PROMPT.to_string()))?;
            if let Ok(Some(hex)) =
                keyfork_qrcode::scan_camera(std::time::Duration::from_secs(30), 0)
            {
                let decoded_data = smex::decode(&hex)?;
                nonce_data = Some(decoded_data[..12].try_into().map_err(|_| InvalidData)?);
                pubkey_data = Some(decoded_data[12..].try_into().map_err(|_| InvalidData)?)
            } else {
                pm.prompt_message(PromptMessage::Text(QRCODE_ERROR.to_string()))?;
            };
        }

        // if QR code scanning failed or was unavailable, read from a set of mnemonics
        let (nonce, their_pubkey) = match (nonce_data, pubkey_data) {
            (Some(nonce), Some(pubkey)) => (nonce, pubkey),
            _ => {
                let validator = MnemonicSetValidator {
                    word_lengths: [9, 24],
                };
                let [nonce_mnemonic, pubkey_mnemonic] = pm
                    .prompt_validated_wordlist::<English, _>(
                        QRCODE_COULDNT_READ,
                        3,
                        validator.to_fn(),
                    )?;

                let nonce = nonce_mnemonic
                    .as_bytes()
                    .try_into()
                    .map_err(|_| InvalidData)?;
                let pubkey = pubkey_mnemonic
                    .as_bytes()
                    .try_into()
                    .map_err(|_| InvalidData)?;
                (nonce, pubkey)
            }
        };

        // 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))
            .to_bytes();
        let hkdf = Hkdf::<Sha256>::new(None, &shared_secret);
        let mut hkdf_output = [0u8; 256 / 8];
        hkdf.expand(&[], &mut hkdf_output)?;
        let shared_key = Aes256Gcm::new_from_slice(&hkdf_output)?;

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

        // encrypt data
        let nonce = Nonce::<U12>::from_slice(&nonce);
        let payload_bytes = shared_key.encrypt(nonce, payload.as_slice())?;

        // convert data to a static-size payload
        // NOTE: Padding length is less than u8::MAX because ENC_LEN < u8::MAX
        #[allow(clippy::assertions_on_constants)]
        {
            assert!(ENC_LEN < u8::MAX, "padding byte can be u8");
        }
        #[allow(clippy::cast_possible_truncation)]
        let mut out_bytes = [payload_bytes.len() as u8; ENC_LEN as usize];
        assert!(
            payload_bytes.len() < out_bytes.len(),
            "encrypted payload larger than acceptable limit"
        );
        out_bytes[..payload_bytes.len()].clone_from_slice(&payload_bytes);

        // NOTE: This previously used a single repeated value as the padding byte, but resulted in
        // difficulty when entering in prompts manually, as one's place could be lost due to
        // repeated keywords. This is resolved below by having sequentially increasing numbers up to
        // but not including the last byte.
        #[allow(clippy::cast_possible_truncation)]
        for (i, byte) in (out_bytes[payload_bytes.len()..(ENC_LEN as usize - 1)])
            .iter_mut()
            .enumerate()
        {
            *byte = (i % u8::MAX as usize) as u8;
        }

        // safety: size of out_bytes is constant and always % 4 == 0
        let payload_mnemonic = unsafe { Mnemonic::from_raw_bytes(&out_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(&smex::encode(&qrcode_data), ErrorCorrection::Highest) {
                pm.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(),
                ))?;
                pm.prompt_message(PromptMessage::Data(qrcode))?;
            }
        }

        pm.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 = 1;
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 nonce = Aes256Gcm::generate_nonce(&mut OsRng);
        let nonce_mnemonic = unsafe { Mnemonic::from_raw_bytes(nonce.as_slice()) };
        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 mut qrcode_data = nonce_mnemonic.to_bytes();
            qrcode_data.extend(key_mnemonic.as_bytes());
            if let Ok(qrcode) = qrencode(&smex::encode(&qrcode_data), ErrorCorrection::Highest) {
                pm.prompt_message(PromptMessage::Text(format!(
                    concat!(
                        "A QR code will be displayed after this prompt. ",
                        "Send the QR code to only shardholder {iter}. ",
                        "Nobody else should scan this 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 shardholder {iter}: ",
                "{nonce_mnemonic} {key_mnemonic}"
            ),
            iter = iter,
            nonce_mnemonic = nonce_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(hex)) =
                keyfork_qrcode::scan_camera(std::time::Duration::from_secs(QRCODE_TIMEOUT), 0)
            {
                let decoded_data = smex::decode(&hex)?;
                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, 48],
                };

                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)
            }
        };

        let shared_secret = our_key.diffie_hellman(&PublicKey::from(pubkey)).to_bytes();
        let hkdf = Hkdf::<Sha256>::new(None, &shared_secret);
        let mut hkdf_output = [0u8; 256 / 8];
        hkdf.expand(&[], &mut hkdf_output)?;
        let shared_key = Aes256Gcm::new_from_slice(&hkdf_output)?;

        let payload =
            shared_key.decrypt(&nonce, &payload[..payload[payload.len() - 1] as usize])?;
        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);
            }
        }

        shares.push(payload[HUNK_OFFSET..].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(())
}