use core::{ops::{Add, Sub, Mul, Rem, Neg}, fmt::Debug};

use cryptix_bigint::{
    ops::{BigIntOps, BigIntOpsExt, slice::{slice_mul_dig, slice_add_inplace}}, 
    property::IsBigInt,
    digit::Digit, BigUInt
};

use crate::{PrimeModular, Element, field::montgomery::MontgomeryOps};
use crate::group::*;
use crate::ring::*;
use crate::field::*;
use crate::field::montgomery::Montgomery;


/// Element in field F_p
#[derive(PartialEq, Eq, PartialOrd, Ord)]
pub struct FpElement<I, M: PrimeModular<I>>(pub Element<I, M>);

impl<I: Debug, M: PrimeModular<I>> Debug for FpElement<I, M> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        self.0.0.fmt(f)
    }
}

impl<I: Copy, M: PrimeModular<I>> Clone for FpElement<I, M> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<I: Copy, M: PrimeModular<I>> Copy for FpElement<I, M> { }

impl<I, M: PrimeModular<I>> From<FpElement<I, M>> for Element<I, M> {
    #[inline(always)]
    fn from(value: FpElement<I, M>) -> Self {
        value.0
    }
}

impl<I, M: PrimeModular<I>> FpElement<I, M> {
    /// # Safety
    /// 
    /// The safey is the same as `Element::new_unchecked`
    #[inline(always)]
    pub fn new_unchecked(value: I) -> Self {
        Self(Element::new_unchecked(value))
    }
}

impl<I: Copy, M: PrimeModular<I>> FpElement<I, M> {
    pub fn repr(self) -> I {
        self.0.0
    }
}

impl<I, M: PrimeModular<I>> From<I> for FpElement<I, M> 
where
    I: Rem<Output = I>
{
    /// convert a integer to a mod p field element, this value will be converted 
    /// into the canonical representative
    fn from(value: I) -> Self {
        FpElement(Element::new(value))
    }
}

impl<I, M> AbelianGroup for FpElement<I, M> 
where 
    M: PrimeModular<I>,
    Self: Group + CommunicativeAdd
{ }

impl<I, M> Group for FpElement<I, M> 
where 
    M: PrimeModular<I>,
    I: BigIntOps
{ }

impl<I, M> Add for FpElement<I, M> 
where
    M: PrimeModular<I>,
    I: BigIntOps
{
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self(self.0 + rhs.0)
    }
}

impl<I, M> Sub for FpElement<I, M>
where
    M: PrimeModular<I>,
    I: BigIntOps
{
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        Self(self.0 - rhs.0)
    }
}

impl<I, M> AddIdentity for FpElement<I, M>
where
    M: PrimeModular<I>,
    I: BigIntOps + IsBigInt
{
    const ADD_IDENTITY: Self = Self(Element::ZERO);
}

impl<I, M> Neg for FpElement<I, M> 
where
    M: PrimeModular<I>,
    I: BigIntOps
{   
    type Output = Self;

    fn neg(self) -> Self::Output {
        Self(-self.0)
    }
}

/// # Safety
/// 
/// Element is backed by biguint, which is associative under addition
impl<I, M> AssociativeAdd for FpElement<I, M> 
where 
    M: PrimeModular<I>,
    Self: Add<Output = Self> 
{ }

/// # Safety
/// 
/// Element is backed by biguint, which is communicative under addition
impl<I, M> CommunicativeAdd for FpElement<I, M> 
where
    M: PrimeModular<I>,
    Self: Add<Output = Self>
{ }


impl<I, M> Ring for FpElement<I, M> 
where 
    M: PrimeModular<I>,
    Self: Mul<Output = Self> + AssociativeMul + DistributiveMul + AbelianGroup
{ }

impl<I, M> Mul for FpElement<I, M> 
where
    M: Montgomery<I>,
    I: BigIntOpsExt,
    [(); I::DIG_LEN + 1]: ,
    [(); I::DIG_LEN * 2 + 1]: ,
{
    type Output = Self;

    /// calculate a * b mod m
    /// this can be achieved more efficiently with montgomery multiplication
    fn mul(self, rhs: Self) -> Self::Output {
        self.mont_mul(rhs).mont_form()
    }
}

/// # Safety
/// 
/// our element type is backed by biguint, so mod mul is associative
impl<I, M> AssociativeMul for FpElement<I, M> 
where 
    M: Montgomery<I>,
    I: BigIntOpsExt,
    [(); I::DIG_LEN + 1]: ,
    [(); I::DIG_LEN * 2 + 1]: ,
{ }

/// # Safety
/// 
/// our element type is backed by biguint, so mod mul is distributive over add
impl<I, M> DistributiveMul for FpElement<I, M> 
where
    M: Montgomery<I>,
    I: BigIntOpsExt,
    [(); I::DIG_LEN + 1]: ,
    [(); I::DIG_LEN * 2 + 1]: ,
{ }


/// # Safety
/// 
/// 1 is the multiplicative ideneity for biguint
impl<I, M> MulIdentity for FpElement<I, M> 
where
    I: BigIntOpsExt + IsBigInt,
    M: Montgomery<I>,
    [(); I::DIG_LEN + 1]: ,
    [(); I::DIG_LEN * 2 + 1]: ,
{
    const MUL_IDENTITY: Self = Self(Element(I::ONE, core::marker::PhantomData));
}

/// # Safety
/// 
/// BigUInt mod mul is communicative
impl<I, M> CommunicativeMul for FpElement<I, M> 
where
    M: Montgomery<I>,
    I: BigIntOpsExt,
    [(); I::DIG_LEN + 1]: ,
    [(); I::DIG_LEN * 2 + 1]: ,
{ }

/// The montgomery trait bound restricts the modular to odd prime
impl<I, M> MulInverse for FpElement<I, M> 
where
    I: BigIntOpsExt,
    [(); I::DIG_LEN + 1]: ,
    [(); I::DIG_LEN * 2 + 1]: ,
    M: Montgomery<I>
{
    fn mul_inv(self) -> Self {
        // step 1. calculate a * R mod P = a * (R * R) * R^(-1) mod P
        let ar = self.mont_mul(M::RR_P);
        // step 2. calculate a^(-1) * R mod P = (a^(-1) * R^(-1)) * R * R mod P
        let am1r = ar.mont_inv();
        // step 3. calculate a^(-1) mod P = a^(-1) * R * R^(-1) mod P
        am1r.mont_rdc()
    }
}

impl<I, M> Field for FpElement<I, M> 
where
    I: BigIntOpsExt,
    [(); I::DIG_LEN + 1]: ,
    [(); I::DIG_LEN * 2 + 1]: ,
    M: Montgomery<I>
{
    fn hlv(self) -> Self {
        if !self.repr().is_odd() {
            /*
             * # Safety
             * 
             * obviously i >> 1 < i
             */
            return Self::new_unchecked(self.repr() >> 1)
        }

        // here c can only be 0 or 1
        let (tmp, c) = self.repr() + M::P;
        
        // now tmp must be even, we shift it right and set the MSB if overflow occurs during addition
        if c.is_zero() {
            /*
             * # Safety
             * 
             * its safe to directly construct an Element from the result, because self < P => (self + P) / 2 < P
             */
            Self::new_unchecked(tmp >> 1)
        } else {
            /*
             * # Safety
             * 
             * its safe to directly construct an Element from the result, because self < P => (self + P) / 2 < P
             */
            Self::new_unchecked((tmp >> 1).set_bit(I::BIT_LEN - 1, true))
        }
    }

    fn is_zero(&self) -> bool {
        self.0.repr().is_zero()
    } 
}

impl<I, M> MontgomeryOps<I, M> for FpElement<I, M> 
where
    I: BigIntOpsExt,
    [(); I::DIG_LEN + 1]: ,
    [(); I::DIG_LEN * 2 + 1]: ,
    M: Montgomery<I>,
{
    fn mont_mul(self, rhs: Self) -> Self {
        let y0 = rhs.repr().as_slice()[0];
        let mut buf = BigUInt([I::Dig::ZERO; I::DIG_LEN * 2 + 1]);

        self.repr().as_slice().iter().enumerate().for_each(
            // step 2: for i from 0 to n - 1, do
            |(idx, &xi)| {
                let a0 = buf[idx];

                // step 2.1: u_i = (a0 + xi * y0) * m' mod b
                // we use overflow mul and add here to simulate the process of mod b
                let ui = M::NEG_P_INV_B
                    .overflow_mul(xi.overflow_mul(y0).0.overflow_add(a0).0)
                    .0;

                // step 2.2: A = (A + xi * y + ui * m) / b
                // xi * y can have length up to LEN + 1, xi * y / b can have length up to LEN
                let mut tmp = [I::Dig::ZERO; I::DIG_LEN + 1];
                slice_mul_dig(&mut tmp, rhs.repr().as_slice(), xi);
                slice_add_inplace(&mut buf[idx..I::DIG_LEN + idx + 1], &tmp);

                slice_mul_dig(&mut tmp, M::P.as_slice(), ui);
                slice_add_inplace(&mut buf[idx..I::DIG_LEN + idx + 1], &tmp);
            },
        );

        let output = I::from_array(
            buf[I::DIG_LEN..I::DIG_LEN * 2].try_into().unwrap()
        );
        // step 3: if A >= m A -= m
        if !buf[I::DIG_LEN * 2].is_zero() || output >= M::P {
            /*
             * # Safety
             *
             * this operation is safe since we have ensured output - M::P < M::P
             */
            FpElement(Element::new_unchecked((output - M::P).0))
        } else {
            /*
             * # Safety
             *
             * this operation is safe since we have ensured output < Self::P
             */
            FpElement(Element::new_unchecked(output))
        }
    }

    fn mont_rdc(self) -> Self {
        let mut a = BigUInt([I::Dig::ZERO; I::DIG_LEN * 2]);
        a[..I::DIG_LEN].copy_from_slice(self.repr().as_slice());

        let a = (0..I::DIG_LEN).fold(a, |a, idx| {
            // step 2.1: u_i = ai * m' mod b
            // we use overflow mul and add here to simulate the process of mod b
            let ui = M::NEG_P_INV_B.overflow_mul(a[idx]).0;

            // step 2.2: A = (A + xi * y + ui * m) / b
            // xi * y can have length up to LEN + 1, xi * y / b can have length up to LEN
            let mut tmp = BigUInt([I::Dig::ZERO; I::DIG_LEN * 2]);
            slice_mul_dig(&mut tmp[idx..idx + I::DIG_LEN + 1], M::P.as_slice(), ui);

            (a + tmp).0
        });

        let output = I::from_array(a.0[I::DIG_LEN..].try_into().unwrap());
        // step 3: if A >= m A -= m
        if output >= M::P {
            /*
             * # Safety
             *
             * this operation is safe since we have ensured output - M::P < M::P
             */
            FpElement(Element::new_unchecked((output - M::P).0))
        } else {
            /*
             * # Safety
             *
             * this operation is safe since we have ensured output < Self::P
             */
            FpElement(Element::new_unchecked(output))
        }
    }

    fn mont_inv(self) -> Self {
        let mut r = Self::ZERO;
        let mut u = M::P;
        let mut s = M::RR_P;
        let mut v = self.repr();

        while !v.is_zero() {
            // L.6
            if !u.is_odd() {
                u = u >> 1;
                r = r.hlv();
            } else if !v.is_odd() {
                v = v >> 1;
                s = s.hlv();
            } else if u > v {
                u = (u - v).0 >> 1;
                r = r - s;
                r = r.hlv();
            } else {
                v = (v - u).0 >> 1;
                s = s - r;
                s = s.hlv();
            }
        }

        r
    }

    fn mont_mul_fp(self, rhs: FpElement<I, M>) -> Self {
        self.mont_mul(rhs)
    }
}

impl<I: IsBigInt, M: PrimeModular<I>> From<FpElement<I, M>> for [u8; I::BYTE_LEN] 
where
    [(); I::DIG_LEN]: 
{
    fn from(val: FpElement<I, M>) -> Self {
        unsafe {
            /*
             * # Safety
             * 
             * This is safe since I::BYTE_LEN = Self::DIG_LEN * Self::DIG_BYTE_LEN, thus this 
             * is just an array reshaping
             */
            core::intrinsics::transmute_unchecked(val.0.0.to_array())
        }
    }
}

#[cfg(feature = "rand")]
impl<I, M> FpElement<I, M> 
where
    I: IsBigInt + BigIntOpsExt,
    M: PrimeModular<I>,
    [(); I::DIG_LEN]:
{
    pub fn rand(rng: &mut impl rand_core::CryptoRngCore) -> Self {
        loop {
            let i = I::rand(rng);
            if i < M::P {
                return FpElement(Element::new_unchecked(i))
            }
        }
    }
}