use clap::Parser;

/// program that generates ed25519 keypairs seeded by a passphrase and an optional ID.
#[derive(Parser)]
#[clap(
    name = "duralumin-keygen",
    version = "0.2.0",
    author = "No One <noonebtw@nirgendwo.xyz>"
)]
struct Opts {
    #[clap(short, long, default_value = "duralumin")]
    file: String,
}

pub mod keygen {
    use std::str::FromStr;

    use osshkeys::{error::OsshResult, keys::ed25519::Ed25519KeyPair, KeyPair};
    use rand::{Rng, SeedableRng};
    use sha2::Digest;
    use thiserror::Error;
    use zeroize::Zeroizing;

    #[derive(Debug, Error)]
    pub enum Error {
        #[error("Failed to parse")]
        ParseError,
    }

    #[derive(Debug, Clone)]
    pub enum KeyType {
        SSH,
        PGP,
    }

    impl FromStr for KeyType {
        type Err = Error;

        fn from_str(s: &str) -> Result<Self, Self::Err> {
            match s.to_lowercase().as_str() {
                "ssh" => Ok(Self::SSH),
                "pgp" => Ok(Self::PGP),
                _ => Err(Error::ParseError),
            }
        }
    }

    #[derive(Debug, Clone)]
    pub enum HashType {
        Sha256,
        Argon2,
    }

    impl FromStr for HashType {
        type Err = Error;

        fn from_str(s: &str) -> Result<Self, Self::Err> {
            match s.to_lowercase().as_str() {
                "sha256" => Ok(Self::Sha256),
                "argon2" | "argon" => Ok(Self::Argon2),
                _ => Err(Error::ParseError),
            }
        }
    }

    #[derive(Debug, Clone)]
    pub struct SshKeyBuilder {
        hash_type: HashType,
        iterations: i16,
        tag: Zeroizing<String>,
        passphrase: Zeroizing<String>,
        encrypt_key: bool,
    }

    impl Default for SshKeyBuilder {
        fn default() -> Self {
            Self {
                hash_type: HashType::Argon2,
                iterations: 4,
                tag: Default::default(),
                passphrase: Default::default(),
                encrypt_key: true,
            }
        }
    }

    impl SshKeyBuilder {
        pub fn with_hash_type(mut self, hash_type: HashType) -> Self {
            self.hash_type = hash_type;
            self
        }

        pub fn with_encrypt(mut self, encrypt: bool) -> Self {
            self.encrypt_key = encrypt;
            self
        }

        pub fn with_hash_type_and_iterations(self, hash_type: HashType, iterations: i16) -> Self {
            self.with_hash_type(hash_type).with_iterations(iterations)
        }

        pub fn with_iterations(mut self, iterations: i16) -> Self {
            self.iterations = iterations;
            self
        }

        pub fn with_maybe_tag(self, tag: Option<String>) -> Self {
            match tag {
                Some(tag) => self.with_tag(tag),
                None => self,
            }
        }

        pub fn with_tag(mut self, tag: String) -> Self {
            self.tag = Zeroizing::new(tag);
            self
        }

        pub fn with_passphrase<Z: Into<Zeroizing<String>>>(mut self, passphrase: Z) -> Self {
            self.passphrase = passphrase.into();
            self
        }

        pub fn build_keypair(self) -> OsshResult<((Zeroizing<String>, String), KeyPair)> {
            let hash_seed = Zeroizing::new(
                self.passphrase
                    .chars()
                    .chain(self.tag.chars())
                    .collect::<String>(),
            );

            let mut seed = [0u8; 32];
            match self.hash_type {
                HashType::Sha256 => {
                    let mut digest = sha2::Sha256::digest(hash_seed.as_bytes());

                    assert!(self.iterations <= 1);
                    for _ in 0..(self.iterations - 1) {
                        digest = sha2::Sha256::digest(digest.as_slice());
                    }

                    seed.copy_from_slice(&digest[..32]);
                }
                HashType::Argon2 => {
                    // TODO: make more random salt
                    let salt = b"thissaltneedsupdating";

                    let argon = argon2::Argon2::new(
                        argon2::Algorithm::Argon2i,
                        argon2::Version::V0x13,
                        argon2::Params::new(65536, self.iterations as u32, 4, Some(32))
                            .expect("argon2 params"),
                    );

                    argon
                        .hash_password_into(hash_seed.as_bytes(), salt, &mut seed)
                        .expect("hash passphrase");
                }
            };

            let rng = rand_chacha::ChaChaRng::from_seed(seed).gen::<[u8; 32]>();

            let keypair = Ed25519KeyPair::from_seed(&rng).map(|key| Into::<KeyPair>::into(key))?;

            let private_key = Zeroizing::new(keypair.serialize_openssh(
                self.encrypt_key.then(|| self.passphrase.as_str()),
                osshkeys::cipher::Cipher::Aes256_Ctr,
            )?);

            Ok(((private_key, keypair.serialize_publickey()?), keypair))
        }
    }
}

fn main() -> anyhow::Result<()> {
    let opts = Opts::parse();
    println!("Generating ed25519 ssh keypair:");

    let desc = libduralumin::key_gen::cli::keygen_desc_from_stdin()?;
    let keypair = libduralumin::key_gen::generate_key(desc)?;

    println!(
        "Your key fingerprint is: Sha256:{}",
        keypair.fingerprint_base64()
    );

    println!(
        "RandomArt:\n{}",
        keypair.randomart().render("ED25519 256", "SHA256")?
    );

    let private_path = opts.file.clone();
    let public_path = opts.file.clone() + ".pub";

    let (private_key, public_key) = keypair.encode_keys()?;
    std::fs::write(&private_path, private_key)?;

    std::fs::write(&public_path, public_key)?;

    #[cfg(target_family = "unix")]
    {
        use std::fs::Permissions;
        use std::os::unix::fs::PermissionsExt;

        std::fs::set_permissions(private_path, Permissions::from_mode(0o0600))?;
        std::fs::set_permissions(public_path, Permissions::from_mode(0o0600))?;
    }

    Ok(())
}