use clap::Args;
use eyre::{bail, Result};
use tokio::{fs::File, io::AsyncWriteExt};

use atuin_client::{
    encryption::{decode_key, encode_key, generate_encoded_key, load_key, Key},
    record::sqlite_store::SqliteStore,
    record::store::Store,
    settings::Settings,
};

#[derive(Args, Debug)]
pub struct Rekey {
    /// The new key to use for encryption. Omit for a randomly-generated key
    key: Option<String>,
}

impl Rekey {
    pub async fn run(&self, settings: &Settings, store: SqliteStore) -> Result<()> {
        let key = if let Some(key) = self.key.clone() {
            println!("Re-encrypting store with specified key");

            let key = match bip39::Mnemonic::from_phrase(&key, bip39::Language::English) {
                Ok(mnemonic) => encode_key(Key::from_slice(mnemonic.entropy()))?,
                Err(err) => {
                    if let Some(err) = err.downcast_ref::<bip39::ErrorKind>() {
                        match err {
                            // assume they copied in the base64 key
                            bip39::ErrorKind::InvalidWord => key,
                            bip39::ErrorKind::InvalidChecksum => {
                                bail!("key mnemonic was not valid")
                            }
                            bip39::ErrorKind::InvalidKeysize(_)
                            | bip39::ErrorKind::InvalidWordLength(_)
                            | bip39::ErrorKind::InvalidEntropyLength(_, _) => {
                                bail!("key was not the correct length")
                            }
                        }
                    } else {
                        // unknown error. assume they copied the base64 key
                        key
                    }
                }
            };

            key
        } else {
            println!("Re-encrypting store with freshly-generated key");
            let (_, encoded) = generate_encoded_key()?;
            encoded
        };

        let current_key: [u8; 32] = load_key(settings)?.into();
        let new_key: [u8; 32] = decode_key(key.clone())?.into();

        store.re_encrypt(&current_key, &new_key).await?;

        println!("Store rewritten. Saving new key");
        let mut file = File::create(settings.key_path.clone()).await?;
        file.write_all(key.as_bytes()).await?;

        Ok(())
    }
}
