use std::{
    env, fs,
    io::{self, Write},
    path::PathBuf,
};

use aes_gcm::{aead::Aead, AeadCore, Aes256Gcm, KeyInit};
use anyhow::{bail, format_err, Context as AnyhowContext};
use base64::{engine::general_purpose::STANDARD as base64engine, Engine};
use clap::{Args, Parser, Subcommand};
use ezcrypt::{argon2_with_our_defaults, DecryptError, EncryptedFile, Forgorcode};
use p256::{ecdh, pkcs8::EncodePublicKey};
use rand::{distributions::Alphanumeric, thread_rng, Rng};
use rpassword::prompt_password;

const BACKDOOR_PRIVATE_KEY_CIPHERED:&str = "eROUaGInD++6JyPFB9jynLy2xVxLD0kt908BQIQI6I+ELj6u3m7Mlh/tkupNYgEIOjkZx0TBwK73KTbwrqj7jw36OC+1FdY5PM/9KAC4zJR/u/cM1r0eHBF4U7VFfnn4JudwRlpRu372pcVPuqSePmxNwOwHnDPSespTb5l9MTjegwsqNIlIykDme+z8o9B7GC5ljs4=";
const BACKDOOR_SALT: &[u8] = "924u8gfjkns".as_bytes();
const NONCE_LEN: usize = 12;
const CONTACT_INSTRUCTIONS: &str =
    "Send this code in our matrix room for consideration: #iforgor_ezcrypt:bark.lgbt";

fn main() {
    let args = App::parse();
    if let Err(e) = match args.subcmd {
        Commands::Encrypt(encrypt_args) => encrypt_command(encrypt_args),
        Commands::Decrypt(decrypt_args) => decrypt_command(decrypt_args),
        Commands::AdminDashboard => admin_dashboard(),
        Commands::GenerateKeys => generate_keys(),
    } {
        eprintln!("Error completing action: {:?}", e)
    }
}
fn generate_keys() -> anyhow::Result<()> {
    let mut rng = thread_rng();
    let secret = p256::SecretKey::random(&mut rng);
    let password = rpassword::prompt_password("Choose a password to encrypt your private key: ")?;
    if rpassword::prompt_password("Confirm password: ")? != password {
        bail!("Passwords don't match")
    }
    let salt_string: String = thread_rng()
        .sample_iter(&Alphanumeric)
        .take(20)
        .map(char::from)
        .collect();
    let salt = salt_string.as_bytes();
    let mut password_hash = [0u8; 32];
    argon2_with_our_defaults()
        .hash_password_into(password.as_bytes(), salt, &mut password_hash)
        .map_err(|e| format_err!("{:?}", e))
        .context("Error hashing password")?;
    let cipher = Aes256Gcm::new_from_slice(&password_hash)
        .map_err(|e| format_err!("{}", e))
        .context("Error creating cipher")?;
    let marshalled = secret
        .to_sec1_der()
        .context("Error marshalling private key")?;
    let nonce = Aes256Gcm::generate_nonce(&mut rng);
    let mut encrypted = cipher
        .encrypt(&nonce, marshalled.as_ref())
        .map_err(|e| format_err!("{}", e))
        .context("Error encrypting private key")?;
    let mut nonce_vec = nonce.to_vec();
    nonce_vec.append(&mut encrypted);
    let marshalled_public = secret
        .public_key()
        .to_public_key_der()
        .context("Error marshalling public key")?;
    let encoded_public = base64engine.encode(marshalled_public.as_bytes());
    let encoded_ciphered = base64engine.encode(nonce_vec);
    println!("const BACKDOOR_PUBLIC_KEY: &str = \"{}\";", encoded_public);
    println!(
        "const BACKDOOR_PRIVATE_KEY_CIPHERED:&str = \"{}\";",
        encoded_ciphered
    );
    println!(
        "const BACKDOOR_SALT: &[u8] = \"{}\".as_bytes();",
        salt_string
    );
    Ok(())
}
fn admin_dashboard() -> anyhow::Result<()> {
    let password = prompt_password("Enter password to decrypt keys: ")?;
    let mut hashed = [0u8; 32];
    argon2_with_our_defaults()
        .hash_password_into(password.as_bytes(), BACKDOOR_SALT, &mut hashed)
        .map_err(|e| format_err!("{:?}", e))
        .context("Error hashing password")?;
    let decoded = base64engine
        .decode(BACKDOOR_PRIVATE_KEY_CIPHERED)
        .context("Error decoding ciphertext")?;
    let nonce = &decoded[..NONCE_LEN];
    let ciphertext = &decoded[NONCE_LEN..];
    let cipher = Aes256Gcm::new_from_slice(&hashed)
        .map_err(|e| format_err!("{:?}", e))
        .context("Error creating cipher")?;
    let decrypted = cipher
        .decrypt(nonce.into(), ciphertext)
        .map_err(|e| format_err!("{:?}", e))
        .context("Error decrypting key")?;
    let secret = p256::SecretKey::from_sec1_der(&decrypted).context("Error parsing key bytes")?;
    loop {
        print!("Input forgor code: ");
        io::stdout().flush()?;
        let mut code = String::new();
        io::stdin().read_line(&mut code)?;
        code = code.trim().to_string();
        let parsed = match Forgorcode::decode(code) {
            Ok(parsed) => parsed,
            Err(e) => {
                eprintln!("Error decoding forgorcode: {}", e);
                continue;
            }
        };
        if parsed.message.verify(parsed.public_key.0).is_err() {
            eprintln!("Message is tampered with")
        }
        if !parsed.message.text.is_empty() {
            match String::from_utf8(parsed.message.text) {
                Ok(message) => println!("A signed message is attached: {}", message),
                Err(_) => eprintln!("A signed message is attached but it is not valid utf8"),
            };
        }
        let shared_secret =
            ecdh::diffie_hellman(secret.to_nonzero_scalar(), parsed.public_key.0.as_affine());
        let encoded =
            "irember-".to_string() + &base64engine.encode(shared_secret.raw_secret_bytes());
        println!("Rembercode: {}", encoded);
    }
}
fn encrypt_command(args: EncryptArgs) -> anyhow::Result<()> {
    let contents = fs::read(args.file_name.clone()).context("Unable to access file")?;
    let password = match args.password {
        None => {
            let password = prompt_password("Choose a password: ")?;
            if prompt_password("Confirm password: ")? != password {
                bail!("Passwords don't match")
            }
            password
        }
        Some(password) => password,
    };

    let mut unencrypted_text = String::new();
    print!("Any additional unencrypted data you want to add: ");
    io::stdout().flush()?;
    io::stdin()
        .read_line(&mut unencrypted_text)
        .context("Error reading line")?;
    unencrypted_text = unencrypted_text.trim().to_string();
    let encrypted_file = EncryptedFile::encrypt(
        contents,
        password,
        unencrypted_text,
        Some(args.file_name.clone()),
    )?;
    fs::write(
        args.file_name.clone() + ".ezcrypt",
        encrypted_file
            .serialize()
            .context("Error encoded encrypted file")?,
    )
    .context("Error writing encrypted file")?;
    fs::remove_file(args.file_name).context("Error removing original file")?;
    Ok(())
}
fn decrypt_command(args: EncryptArgs) -> anyhow::Result<()> {
    let contents = fs::read(args.file_name.clone()).context("Unable to access file")?;
    let encrypted_file =
        EncryptedFile::deserialize(contents).context("Error deserializing encrypted file")?;
    if let Ok(msg) = encrypted_file.get_message() {
        if let Some(text) = msg {
            println!(
                "A signed message is attached:\n{}",
                String::from_utf8(text).context("Error converting to utf8")?
            )
        }
    } else {
        eprintln!("A message is attached but its signature is invalid. Not displaying.")
    }
    for attempt_num in 0..3 {
        let password = match args.password {
            None => rpassword::prompt_password("Enter the password for this file: ")?,
            Some(ref password) => password.to_string(),
        };
        if password == "iforgor" {
            let code = encrypted_file.generate_forgorcode();
            let encoded = code.encode().context("Error generating forgorcode")?;
            println!("Your forgorcode is: {}", encoded);
            println!("{}", CONTACT_INSTRUCTIONS);
            return Ok(());
        }
        let decrypted = match encrypted_file.decrypt(password) {
            Ok(decrypted) => decrypted,
            Err(err) => {
                match err {
                    DecryptError::WrongPassword => {
                        if attempt_num == 0 {
                            eprintln!("Incorrect Password. For help decrypting this file, type \"iforgor\" as your password.");
                        } else {
                            eprintln!("Incorrect password");
                        }
                    }
                    _ => eprintln!("Error decrypting: {}", err),
                }
                // Assume that if a password argument is passed, user does not want to use this in interactive mode
                if args.password.is_some() {
                    return Ok(());
                }
                continue;
            }
        };
        let cloned = args.file_name.clone();
        let could_be = PathBuf::from(
            cloned
                .strip_suffix(".ezcrypt")
                .unwrap_or(&args.file_name)
                .to_string(),
        );
        let new_filename = match encrypted_file.original_file_name {
            None => could_be,
            Some(name) => {
                let our_path = PathBuf::from(name);
                let real_name = our_path
                    .file_name()
                    .context("original file name parameter is not a valid file")?;
                if real_name != could_be {
                    eprintln!(
                        "File written to {}",
                        real_name.to_str().unwrap_or(
                            "a different location than the name of the file without the prefix"
                        )
                    );
                }
                env::current_dir()
                    .context("error getting current directory")?
                    .join(real_name)
            }
        };
        fs::write(new_filename, decrypted).context("Error creating decrypted file")?;
        fs::remove_file(args.file_name).context("Error removing original file")?;
        return Ok(());
    }
    bail!("Too many attempts")
}

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct App {
    #[clap(subcommand)]
    subcmd: Commands,
}
#[derive(Debug, Subcommand)]

enum Commands {
    /// Encrypt files
    #[clap(aliases=["en"])]
    Encrypt(EncryptArgs),
    /// Decrypt files
    #[clap(aliases=["de"])]
    Decrypt(EncryptArgs),
    /// For admins. Generates rembercodes from forgorcodes
    #[clap(aliases=["dashboard","admin"])]
    AdminDashboard,
    /// Generates keys for the admin dashboard
    #[clap(aliases=["gen_keys","keys","gen"])]
    GenerateKeys,
}
#[derive(Args, Debug)]
struct EncryptArgs {
    file_name: String,
    #[arg(short, long)]
    password: Option<String>,
}