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

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

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

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

/// 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

/// 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 wordlist = Wordlist::default();

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

    let mut threshold = 0;

    while iter_count.is_none() || iter_count.is_some_and(|i| i > 0) {
        let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
        let nonce_mnemonic =
            unsafe { Mnemonic::from_raw_entropy(nonce.as_slice(), Default::default()) };
        let our_key = EphemeralSecret::random();
        let key_mnemonic =
            Mnemonic::from_entropy(PublicKey::from(&our_key).as_bytes(), Default::default())?;

        #[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::Data(qrcode))?;
            }
        }

        pm.prompt_message(PromptMessage::Text(format!(
            "Our words: {nonce_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(
                    "Unable to detect QR code, falling back to text".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("Their words: ", &wordlist, 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(())
}