use std::{io::Write, str::FromStr};

use clap::Parser;
use libduralumin::ed25519::randomart;
use osshkeys::PublicParts;
use zeroize::Zeroizing;

use crate::keygen::HashType;

/// 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 = "id_ed25519")]
    file: String,
}

fn fix_newline_ref(line: &mut String) {
    if line.ends_with('\n') {
        line.pop();

        if line.ends_with('\r') {
            line.pop();
        }
    }
}

#[allow(dead_code)]
fn fix_newline(mut line: String) -> String {
    fix_newline_ref(&mut line);

    line
}

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 seed = 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);
                    }
                    digest.to_vec()
                }
                HashType::Argon2 => {
                    // TODO: make more random salt
                    let salt = b"thissaltneedsupdating";

                    let config = argon2::Config {
                        variant: argon2::Variant::Argon2i,
                        version: argon2::Version::Version13,
                        mem_cost: 65536,
                        time_cost: self.iterations as u32,
                        lanes: 4,
                        thread_mode: argon2::ThreadMode::Parallel,
                        hash_length: 32,
                        ..Default::default()
                    };

                    argon2::hash_raw(hash_seed.as_bytes(), salt, &config)
                        .expect("hash passphrase")
                }
            }
            .try_into()
            .expect("unwrap seed into [u8; 32]");

            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 read_line() -> std::io::Result<String> {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line)?;
    fix_newline_ref(&mut line);

    Ok(line)
}

fn read_non_empty_line() -> std::io::Result<Option<String>> {
    let line = read_line()?;

    Ok(if line.is_empty() { None } else { Some(line) })
}

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

    let ((private_key, public_key), keypair) = keygen::SshKeyBuilder::default()
        .with_passphrase({
            print!("Enter a passphrase: ");
            std::io::stdout().flush()?;
            let passphrase = Zeroizing::new(rpassword::read_password()?);
            print!("Re-enter the same passphrase: ");
            std::io::stdout().flush()?;
            let passphrase2 = Zeroizing::new(rpassword::read_password()?);

            if passphrase == passphrase2 {
                Ok(passphrase)
            } else {
                println!("passphrases do not match.");
                Err(std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    "passphrase did not match.",
                ))
            }
        }?)
        .with_maybe_tag({
            print!("Enter an optional ID []: ");
            std::io::stdout().flush()?;

            read_non_empty_line()?
        })
        .with_encrypt({
            print!("Encrypt keypair with passphrase? [Y/n]: ");
            std::io::stdout().flush()?;

            let encrypting = read_line()? != "n";
            if encrypting {
                println!("Encrypting keypair..")
            };

            encrypting
        })
        .with_hash_type({
            print!("Use hash algorithm (sha256, argon2)[argon2]: ");
            std::io::stdout().flush()?;

            read_non_empty_line()?
                .map(|line| HashType::from_str(&line))
                .unwrap_or(Ok(HashType::Argon2))?
        })
        .with_iterations({
            print!("Number of hashing iterations [4]: ");
            std::io::stdout().flush()?;

            read_non_empty_line()?
                .map(|line| i16::from_str(&line))
                .unwrap_or(Ok(4))?
        })
        .build_keypair()?;

    let fingerprint =
        keypair.fingerprint(osshkeys::keys::FingerprintHash::SHA256)?;

    println!(
        "Your key fingerprint is: Sha256:{}",
        base64::encode(&fingerprint)
    );

    let randomart = randomart::RandomArt::from_digest(&fingerprint)
        .render("ED25519 256", "SHA256")?;

    println!("RandomArt:\n{}", randomart);

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

    std::fs::write(&private_path, private_key)?;

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

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

    #[cfg(target_family = "unix")]
    std::fs::set_permissions(private_path, Permissions::from_mode(0o0600))?;
    #[cfg(target_family = "unix")]
    std::fs::set_permissions(public_path, Permissions::from_mode(0o0600))?;

    Ok(())
}
