use std::{
    io::{Read, Write},
    path::Path,
};

use rsa::{
    pkcs1::{DecodeRsaPrivateKey, DecodeRsaPublicKey, EncodeRsaPrivateKey, EncodeRsaPublicKey},
    pkcs8::LineEnding,
    Pkcs1v15Encrypt, RsaPrivateKey, RsaPublicKey,
};

use crate::{
    AsymmetricDecryptOptions, AsymmetricEncryptOptions, AsymmetricKeyOptions, Error, Input, Output,
};

#[derive(Default, Debug, Clone, Copy, PartialEq, clap::ValueEnum)]
pub enum KeyFormat {
    #[default]
    #[clap(alias("pem"), help = "PKCS#1 PEM")]
    Pkcs1Pem,
    #[clap(alias("der"), help = "PKCS#1 DER")]
    Pkcs1Der,
}

#[derive(Debug, clap::Parser)]
pub struct RsaKeyOptions {
    #[clap(short = 'f', long, help = "Key output format", default_value = "pem")]
    pub format: KeyFormat,
}

pub fn read_private_key<P: AsRef<Path>>(keyfile: P) -> Result<RsaPrivateKey, Error> {
    let keyfile = keyfile.as_ref();
    if let Ok(key) = RsaPrivateKey::read_pkcs1_pem_file(keyfile) {
        return Ok(key);
    }
    let key = RsaPrivateKey::read_pkcs1_der_file(keyfile)?;
    Ok(key)
}

pub fn read_public_key<P: AsRef<Path>>(keyfile: P) -> Result<RsaPublicKey, Error> {
    let keyfile = keyfile.as_ref();
    if let Ok(key) = RsaPublicKey::read_pkcs1_pem_file(keyfile) {
        return Ok(key);
    }
    let key = RsaPublicKey::read_pkcs1_der_file(keyfile)?;
    Ok(key)
}

pub fn generate_rsa(
    options: &AsymmetricKeyOptions,
    rsa: &RsaKeyOptions,
    bits: usize,
) -> Result<(), Error> {
    if rsa.format == KeyFormat::Pkcs1Der && options.public_key.is_none() {
        return Err(Error::IncompatibleOptions(
            "DER output format requires -o/--public-key".into(),
        ));
    }

    println!("Generating a new RSA keypair, bits={}", bits);
    println!("This will take a while");
    let mut rng = rand::thread_rng();
    let private_key = RsaPrivateKey::new(&mut rng, bits)?;

    println!("Done, deriving a public key");
    let public_key = private_key.to_public_key();

    match rsa.format {
        KeyFormat::Pkcs1Pem => {
            private_key.write_pkcs1_pem_file(&options.private_key, LineEnding::LF)?;

            if let Some(public_key_path) = options.public_key.as_ref() {
                public_key.write_pkcs1_pem_file(public_key_path, LineEnding::LF)?;
            } else {
                let output = public_key.to_pkcs1_pem(LineEnding::LF)?;
                println!("Public key:");
                println!("{output}");
            }
        }
        KeyFormat::Pkcs1Der => {
            private_key.write_pkcs1_der_file(&options.private_key)?;

            if let Some(public_key_path) = options.public_key.as_ref() {
                public_key.write_pkcs1_der_file(public_key_path)?;
            }
        }
    }

    println!("Done!");
    Ok(())
}

pub fn encrypt_rsa(options: &AsymmetricEncryptOptions) -> Result<(), Error> {
    let key = read_public_key(&options.public_key)?;
    let mut input = Input::from_path_option(options.input.as_ref())?;
    let mut output = Output::from_path_option(options.output.as_ref())?;
    let mut buf = [0; 2048];
    let mut rng = rand::thread_rng();
    loop {
        let len = input.read(&mut buf)?;
        if len == 0 {
            break;
        }
        let ciphertext = key.encrypt(&mut rng, Pkcs1v15Encrypt, &buf[..len])?;
        output.write_all(&ciphertext)?;
    }

    Ok(())
}

pub fn decrypt_rsa(options: &AsymmetricDecryptOptions) -> Result<(), Error> {
    let key = read_private_key(&options.private_key)?;
    let mut input = Input::from_path_option(options.input.as_ref())?;
    let mut output = Output::from_path_option(options.output.as_ref())?;
    let mut buf = [0; 2048];
    loop {
        let len = input.read(&mut buf)?;
        if len == 0 {
            break;
        }
        let plaintext = key.decrypt(Pkcs1v15Encrypt, &buf[..len])?;
        output.write_all(&plaintext)?;
    }

    Ok(())
}
