use std::sync::LazyLock;

use aes::cipher::block_padding;
use aes::cipher::{BlockDecryptMut as _, BlockEncryptMut as _, KeyIvInit as _};
use crypto_bigint::{Integer, Random};
use crypto_bigint::{
    NonZero, RandomMod as _, U512,
    rand_core::{OsRng, RngCore as _},
};
use sha2::Digest as _;

use crate::ring::{Ring512, Ring512Element};

pub static DEFAULT_CRYPTO_PARAMS: LazyLock<CryptoParams> = LazyLock::new(|| {
    CryptoParams::new(
    crypto_bigint::U512::from_be_hex(
        "0000d7cc0f742ba809fdef300174fa10badf20510c55b0124958b921b6aeb9b7bb036c4f2cf3927f1832f0d5f73ebeee144099e517d2900884f2a2324e52c54d",
    ),
    crypto_bigint::U512::from_be_hex(
        "00006bebd28220130bcac04f37cf06ef602752f0718d30061b70f64c22dca385e19efc0c767518103430448ee70477a1e8df6aea6d3faab87b317f154e682f09",
    ),
).unwrap()
});

pub const EXPONENT: U512 = U512::from_u32(65537);

#[derive(Debug)]
pub struct CryptoParams {
    ring: &'static Ring512,
    modulus512: NonZero<U512>,
    base: Ring512Element,
    base2: Ring512Element,
}

impl CryptoParams {
    pub fn new(modulus: U512, base: U512) -> anyhow::Result<Self> {
        if !Self::check_if_strong_prime(&modulus) {
            anyhow::bail!("Modulus is not a strong prime");
        }
        let Some(modulus) = NonZero::new(modulus).into() else {
            anyhow::bail!("Modulus cannot be zero");
        };
        let ring = Box::leak(Box::new(Ring512::new(modulus)));
        let base = ring.reduce(base.resize());

        if base.pow(U512::from_u32(2)).value().bits() < 2 {
            anyhow::bail!("Base is bad");
        }

        let base2 = base.pow(EXPONENT);

        Ok(CryptoParams {
            ring,
            modulus512: modulus,
            base,
            base2,
        })
    }

    fn check_if_strong_prime(_modulus: &U512) -> bool {
        _modulus.is_odd().into()
    }

    pub fn ring(&self) -> &Ring512 {
        self.ring
    }

    pub fn modulus(&self) -> &NonZero<U512> {
        &self.modulus512
    }

    pub fn base(&self) -> &Ring512Element {
        &self.base
    }

    pub fn base2(&self) -> &Ring512Element {
        &self.base2
    }
}
struct Rsa {
    n: crypto_bigint::U1024,
    e: crypto_bigint::U1024,
    d: crypto_bigint::U1024,
}
#[allow(dead_code)]
fn generate_rsa() -> Rsa {
    let e = crypto_bigint::U1024::from_u8(3);
    let p = crypto_bigint::U1024::random(&mut OsRng);
    let q = crypto_bigint::U1024::random(&mut OsRng);
    let n = p * q;
    let phi = (p - crypto_bigint::U1024::ONE) * (q - crypto_bigint::U1024::ONE);
    let d = e.inv_mod(&phi).unwrap_or(crypto_bigint::U1024::ONE);
    Rsa { n, e, d }
}

impl PartialEq for CryptoParams {
    fn eq(&self, other: &Self) -> bool {
        self.modulus512 == other.modulus512
            && self.base.value() == other.base.value()
            && self.base2.value() == other.base2.value()
    }
}

#[derive(Debug, Clone)]
pub struct PrivateKey {
    params: &'static CryptoParams,
    k: U512,
}

impl PrivateKey {
    pub fn new_random(params: &'static CryptoParams) -> Self {
        let k = U512::random_mod(
            &mut OsRng,
            &NonZero::new(params.ring.modulus().resize()).unwrap(),
        );
        PrivateKey { params, k }
    }

    pub fn from_k(params: &'static CryptoParams, k: U512) -> anyhow::Result<Self> {
        if k > params.ring.modulus().resize() {
            anyhow::bail!("k is bigger than modulus");
        }
        Ok(PrivateKey { params, k })
    }

    pub fn params(&self) -> &'static CryptoParams {
        self.params
    }

    pub fn k(&self) -> &U512 {
        &self.k
    }

    pub fn decrypt(&self, ciphertext: &Ciphertext) -> anyhow::Result<Vec<u8>> {
        if self.params.ring() != ciphertext.c1.ring() {
            anyhow::bail!("Mismatched crypto parameters");
        }

        let s = ciphertext.c1.pow(self.k);

        let mut hasher = sha2::Sha256::new();
        hasher.update(s.value().to_be_bytes());
        let key_bytes = hasher.finalize();

        let decryptor =
            cbc::Decryptor::<aes::Aes128>::new(key_bytes[..16].into(), &ciphertext.iv.into());
        let plaintext = decryptor
            .decrypt_padded_vec_mut::<block_padding::Pkcs7>(&ciphertext.ct)
            .map_err(|_| anyhow::anyhow!("invalid ciphertext padding"))?;
        Ok(plaintext)
    }
}

#[derive(Clone)]
pub struct PublicKey {
    params: &'static CryptoParams,
    h: Ring512Element,
}

impl std::fmt::Debug for PublicKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "PublicKey({:?})", self.h)
    }
}

impl PublicKey {
    pub fn from_private_key(sk: &PrivateKey) -> Self {
        let h = sk.params().base().pow(*sk.k());
        PublicKey {
            params: sk.params(),
            h,
        }
    }

    pub fn from_h(params: &'static CryptoParams, h: Ring512Element) -> anyhow::Result<Self> {
        if params.ring() != h.ring() {
            anyhow::bail!("Mismatched crypto parameters");
        }
        Ok(PublicKey { params, h })
    }

    pub fn params(&self) -> &'static CryptoParams {
        self.params
    }

    pub fn h(&self) -> &Ring512Element {
        &self.h
    }

    pub fn encrypt(&self, plaintext: &[u8]) -> Ciphertext {
        let r = U512::random_mod(&mut OsRng, &self.params().modulus512);
        let c1 = self.params().base2().pow(r);
        let s = self.h.pow(r).pow(EXPONENT);

        let mut hasher = sha2::Sha256::new();
        hasher.update(s.value().to_be_bytes());
        let key_bytes = hasher.finalize();

        let mut iv = [0u8; 16];
        OsRng.fill_bytes(&mut iv);
        #[allow(unused_variables)]
        let ciphertext = std::iter::zip(plaintext.iter(), key_bytes.iter().cycle())
            .map(|(x, y)| x ^ y)
            .collect::<Vec<u8>>();

        let ciphertext = cbc::Encryptor::<aes::Aes128>::new(key_bytes[..16].into(), &iv.into())
            .encrypt_padded_vec_mut::<block_padding::Pkcs7>(plaintext);
        Ciphertext {
            c1,
            iv,
            ct: ciphertext,
        }
    }
}

#[derive(Debug, Clone)]
pub struct Ciphertext {
    pub c1: Ring512Element,
    pub iv: [u8; 16],
    pub ct: Vec<u8>,
}
