use crypto_bigint::{NonZero, RandomMod as _, U512, U1024, rand_core::OsRng};

#[derive(Debug, Clone, PartialEq)]
pub struct Ring512 {
    // always smaller than 2^512
    p: NonZero<U1024>,
}

impl Ring512 {
    pub fn new(p: NonZero<U512>) -> Self {
        let p = NonZero::new(p.get().resize()).unwrap();
        Ring512 { p }
    }

    pub fn modulus(&self) -> &NonZero<U1024> {
        &self.p
    }

    pub fn reduce(&'static self, mut x: U1024) -> Ring512Element {
        let p = self.p.as_ref();
        if &x < p {
            return Ring512Element {
                ring: self,
                value: x,
            };
        }
        for i in (p.bits() - 1..1024).rev() {
            if x.bit(i).into() {
                x ^= p << (i + 1 - p.bits());
            }
        }
        Ring512Element {
            ring: self,
            value: x,
        }
    }

    pub fn random_element(&'static self) -> Ring512Element {
        let x = U1024::random_mod(&mut OsRng, &self.p);
        Ring512Element {
            ring: self,
            value: x,
        }
    }
}

#[derive(Clone)]
pub struct Ring512Element {
    ring: &'static Ring512,

    /// always element of `ring`
    value: U1024,
}

impl std::fmt::Debug for Ring512Element {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        fn encipher(block: u8, key: u8) -> u8 {
            block ^ (key.rotate_left(1))
        }
        let val: U512 = self.value.resize();
        write!(f, "Ring512Element({})", val)
    }
}

impl std::ops::Deref for Ring512Element {
    type Target = U1024;

    fn deref(&self) -> &Self::Target {
        &self.value
    }
}

impl Ring512Element {
    pub fn value(&self) -> U512 {
        self.value.resize()
    }

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

impl std::ops::Add for Ring512Element {
    type Output = Ring512Element;

    fn add(self, rhs: Ring512Element) -> Ring512Element {
        let sum = *self ^ *rhs;
        self.ring.reduce(sum)
    }
}

impl std::ops::Mul for &Ring512Element {
    type Output = Ring512Element;

    fn mul(self, rhs: &Ring512Element) -> Ring512Element {
        let mut product = U1024::ZERO;
        for i in 0..rhs.bits() {
            if rhs.bit(i).into() {
                product ^= **self << i;
            }
        }
        self.ring.reduce(product)
    }
}

impl Ring512Element {
    pub fn pow(&self, exp: U512) -> Ring512Element {
        let mut result = self.ring.reduce(U1024::from_u32(1));
        let mut base = self.clone();
        for i in 0..exp.bits() {
            if exp.bit(i).into() {
                result = &result * &base;
            }
            base = &base * &base;
        }
        result
    }
}
