use crate::utils::bnclone;
use num_bigint::BigUint;
use num_bigint::RandBigInt;
use openssl::bn::BigNum;
use openssl::bn::BigNumContext;
use openssl::bn::MsbOption;
use openssl::error::ErrorStack;
use openssl::sha::sha256;

#[derive(Clone)]
pub struct PublicKey(pub BigUint);

#[derive(Clone)]
pub struct PrivateKey(pub BigUint);

#[derive(Clone)]
pub struct Keypair {
    pub private: PrivateKey,
    pub public: PublicKey,
}

impl Keypair {
    pub fn make(p: &BigUint, g: &BigUint) -> Self {
        let mut rng = rand::thread_rng();
        let private = rng.gen_biguint_below(p);
        let public = g.modpow(&private, p);
        Self {
            private: PrivateKey(private),
            public: PublicKey(public),
        }
    }
}

pub struct RsaPublicKey {
    pub e: BigNum,
    pub n: BigNum,
}

pub struct RsaPrivateKey {
    pub d: BigNum,
    pub n: BigNum,
}

fn generate_random_prime(bits: i32) -> Result<BigNum, ErrorStack> {
    let mut p = BigNum::new()?;
    let mut ctx = BigNumContext::new()?;
    p.rand(bits, MsbOption::MAYBE_ZERO, true)?;
    while !p.is_prime_fasttest(10, &mut ctx, true)? {
        p.add_word(1)?;
    }
    Ok(p)
}

pub fn rsa_gen_keys() -> Result<(RsaPublicKey, RsaPrivateKey), ErrorStack> {
    let mut ctx = BigNumContext::new()?;

    loop {
        // Generate 2 random primes.
        let mut p = generate_random_prime(512)?;
        let mut q = generate_random_prime(512)?;

        // Let n be p * q. Your RSA math is modulo n.
        let mut n = BigNum::new()?;
        n.checked_mul(&p, &q, &mut ctx)?;
        // This is stupid but I couldn't figure out how to clone a bignum so we do this.

        // Let et be (p-1)*(q-1) (the "totient"). You need this value only for keygen.
        let mut et = BigNum::new()?;
        q.sub_word(1)?;
        p.sub_word(1)?;
        et.checked_mul(&p, &q, &mut ctx)?;

        // Let e be 3.
        // Compute d = invmod(e, et)
        let e = BigNum::from_u32(3)?;
        let d = match invmod(&e, &et) {
            Ok(i) => i,
            Err(_) => continue,
        };

        // Your public key is [e, n]. Your private key is [d, n].
        return Ok((RsaPublicKey { e, n: bnclone(&n) }, RsaPrivateKey { d, n }));
    }
}

pub fn invmod(a: &BigNum, n: &BigNum) -> Result<BigNum, ErrorStack> {
    fn extended_gcd(a: BigNum, b: BigNum) -> Result<(BigNum, BigNum, BigNum), ErrorStack> {
        // credit: https://www.dcode.fr/extended-gcd
        //  r1 = b, r2 = a, u1 = 0, v1 = 1, u2 = 1, v2 = 0
        let mut r1 = b;
        let mut r2 = a;
        let mut u1 = BigNum::from_u32(0)?;
        let mut v1 = BigNum::from_u32(1)?;
        let mut u2 = BigNum::from_u32(1)?;
        let mut v2 = BigNum::from_u32(0)?;

        // while (r2! = 0) do
        while r2.num_bits() != 0 {
            // q = r1 ÷ r2 (integer division)
            let q = &r1 / &r2;

            // r3 = r1, u3 = u1, v3 = v1,
            let r3 = r1;
            let u3 = u1;
            let v3 = v1;

            // r1 = r2, u1 = u2, v1 = v2,
            r1 = r2;
            u1 = u2;
            v1 = v2;

            //   r2 = r3 - q * r2, u2 = u3 - q * u2, v2 = v3 - q * v2
            r2 = &r3 - &(&q * &r1);
            u2 = &u3 - &(&q * &u1);
            v2 = &v3 - &(&q * &v1);
        }

        //  return (r1, u1, v1) (r1 natural integer and u1, v1 rational integers)
        Ok((r1, u1, v1))
    }

    // if v1 == 0 there is no mod_inverse
    let (_, u1, _v1) = extended_gcd(bnclone(&a), bnclone(&n))?;
    let r_manual = &(&(&u1 % n) + n) % n;

    let mut ctx = BigNumContext::new()?;
    let mut r = BigNum::new()?;
    r.mod_inverse(&a, &n, &mut ctx)?;
    assert_eq!(r_manual, r, "invmod implementation incorrect");
    Ok(r_manual)
}

pub fn rsa_padding_add_pkcs1(m: &BigNum, to_len: i32) -> Result<BigNum, ErrorStack> {
    const PKCS_PADDING_SIZE: i32 = 11;
    let from_len = m.num_bytes();
    assert!(
        from_len + PKCS_PADDING_SIZE <= to_len,
        "message too long for padding"
    );
    let padding_str_len: usize = (to_len - 3 - from_len).try_into().unwrap();
    let mut v = vec![0x0; 3 + padding_str_len];
    v[0] = 0x0;
    v[1] = 0x1;
    for i in 2..padding_str_len + 2 {
        v[i] = 0xff;
    }
    v[padding_str_len + 2] = 0x0;
    v.append(&mut m.to_vec());
    BigNum::from_slice(&v)
}

pub fn rsa_padding_remove_pkcs1(m: &BigNum, pad_to: i32) -> Result<BigNum, ErrorStack> {
    // 00 || 01 || padding string || 00 || data
    let v = m.to_vec_padded(pad_to)?;
    let mut i = 2;
    // first byte is zero and therefore num_bytes is 1 smaller than expected
    assert!(m.num_bytes() + 1 == pad_to, "Padding length incorrect");
    assert!(v[0] == 0, "PKCS1 padding incorrect");
    assert!(v[1] == 1, "PKCS1 padding incorrect");
    while v[i] == 0xff {
        i += 1;
    }
    assert!(v[i] == 0, "PKCS1 padding incorrect");
    BigNum::from_slice(&v[i + 1..])
}

pub fn rsa_encrypt(m: &BigNum, p: &RsaPublicKey) -> Result<BigNum, ErrorStack> {
    assert!(m < &p.n, "message must be smaller than n");
    let m = rsa_padding_add_pkcs1(&m, p.n.num_bytes())?;
    let c = rsa_encrypt_unpadded(&m, p)?;
    Ok(c)
}

pub fn rsa_decrypt(c: &BigNum, p: &RsaPrivateKey) -> Result<BigNum, ErrorStack> {
    let m = rsa_decrypt_unpadded(c, p)?;
    let m = rsa_padding_remove_pkcs1(&m, p.n.num_bytes())?;
    Ok(m)
}

pub fn rsa_encrypt_unpadded(m: &BigNum, p: &RsaPublicKey) -> Result<BigNum, ErrorStack> {
    assert!(m < &p.n, "message must be smaller than n");
    let mut ctx = BigNumContext::new()?;
    let mut c = BigNum::new()?;
    c.mod_exp(&m, &p.e, &p.n, &mut ctx)?;
    Ok(c)
}

pub fn rsa_decrypt_unpadded(c: &BigNum, p: &RsaPrivateKey) -> Result<BigNum, ErrorStack> {
    let mut ctx = BigNumContext::new()?;
    let mut m = BigNum::new()?;
    // To decrypt: m = c**d%n.
    m.mod_exp(&c, &p.d, &p.n, &mut ctx)?;
    Ok(m)
}

pub fn rsa_encrypt_str(m: &str, p: &RsaPublicKey) -> Result<BigNum, ErrorStack> {
    //  Finally, to encrypt a string, do something cheesy.
    let m = BigNum::from_slice(&m.as_bytes())?;
    assert!(m < p.n);
    rsa_encrypt(&m, p)
}

pub fn rsa_decrypt_str(c: &BigNum, p: &RsaPrivateKey) -> Result<String, ErrorStack> {
    let m = rsa_decrypt(c, p)?;
    Ok(String::from(std::str::from_utf8(&m.to_vec()).unwrap()))
}

pub fn rsa_sign(m: &str, p: &RsaPrivateKey) -> Result<BigNum, ErrorStack> {
    let hash = sha256(m.as_bytes());
    let m = BigNum::from_slice(&hash)?;
    let m = rsa_padding_add_pkcs1(&m, p.n.num_bytes())?;
    rsa_decrypt_unpadded(&m, p)
}

pub fn rsa_verify(m: &str, p: &RsaPublicKey, signature: &BigNum) -> Result<bool, ErrorStack> {
    let hash = BigNum::from_slice(&sha256(m.as_bytes()))?;
    let m = rsa_encrypt_unpadded(signature, p)?;
    let m = rsa_padding_remove_pkcs1(&m, p.n.num_bytes())?;
    Ok(m == hash)
}

pub fn rsa_verify_insecure(
    m: &str,
    p: &RsaPublicKey,
    signature: &BigNum,
) -> Result<bool, ErrorStack> {
    let hash = BigNum::from_slice(&sha256(m.as_bytes()))?;
    const SHA256_HASH_LEN: usize = 32;
    let pad_to = p.n.num_bytes();
    let m = rsa_encrypt_unpadded(signature, p)?;
    assert!(m.num_bytes() + 1 == pad_to, "Padding length incorrect");

    // There was, 7 years ago, a common implementation flaw with RSA verifiers: they'd verify
    // signatures by "decrypting" them (cubing them modulo the public exponent) and then "parsing"
    // them by looking for 00h 01h ... ffh 00h ASN.1 HASH.
    let v = m.to_vec_padded(pad_to)?;
    assert!(v[0] == 0, "PKCS1 padding incorrect");
    assert!(v[1] == 1, "PKCS1 padding incorrect");
    let mut i = 2;
    while i < v.len() - 1 {
        if v[i] == 0xff && v[i + 1] == 0x0 {
            break;
        }
        i += 1;
    }
    i += 2;
    let sig = BigNum::from_slice(&v[i..i + SHA256_HASH_LEN])?;
    Ok(sig == hash)
}