use crate::{
    core::circuits::boolean::{boolean_value::Boolean, byte::Byte},
    utils::{elliptic_curve::F25519, number::Number},
};
use std::ops::Not;

pub trait Equal<Other>: Sized {
    type Output: Not<Output = Self::Output>;

    fn eq(self, other: Other) -> Self::Output;
    fn ne(self, other: Other) -> Self::Output {
        Self::eq(self, other).not()
    }
}

pub trait IsZero {
    type Output;

    fn is_zero(&self) -> Self::Output;
}

pub trait GreaterEqual<Other>: Sized {
    type Output: Not<Output = Self::Output>;

    fn ge(self, other: Other) -> Self::Output;
    fn lt(self, other: Other) -> Self::Output {
        Self::ge(self, other).not()
    }
}

pub trait GreaterThan<Other>: Sized {
    type Output: Not<Output = Self::Output>;

    fn gt(self, other: Other) -> Self::Output;
    fn le(self, other: Other) -> Self::Output {
        Self::gt(self, other).not()
    }
}

/// Implement [`Equal`] for all types that implement [`Eq`].
impl<T> Equal<T> for T
where
    T: Eq,
{
    type Output = bool;

    fn eq(self, other: T) -> Self::Output {
        self == other
    }
}

/// Implement [`GreaterEqual`] for all types that implement [`PartialOrd`].
impl<T> GreaterEqual<T> for T
where
    T: PartialOrd,
{
    type Output = bool;

    fn ge(self, other: T) -> Self::Output {
        self >= other
    }
}

/// Implement [`GreaterThan`] for all types that implement [`PartialOrd`].
impl<T> GreaterThan<T> for T
where
    T: PartialOrd,
{
    type Output = bool;

    fn gt(self, other: T) -> Self::Output {
        self > other
    }
}

pub trait Selectable<T = Self> {
    type Conditional;
    type Output;

    fn construct_selection(condition: Self::Conditional, a: Self, b: T) -> Self::Output;
}

pub trait Select<T, U, V> {
    fn select(self, a: T, b: V) -> U;
}

pub trait Enc<T> {
    fn reveal(self) -> T;
}

pub trait FromLeBits<B: Boolean> {
    fn from_le_bits(bits: Vec<B>, signed: bool) -> Self;
}

pub trait GetBit {
    type Output: Boolean;

    fn get_bit(&self, index: usize, signed: bool) -> Self::Output;
}

pub trait FromLeBytes {
    fn from_le_bytes(bytes: [u8; 32]) -> Self;
}

pub trait ToLeBytes {
    type BooleanOutput: Boolean;

    fn to_le_bytes(self) -> [Byte<Self::BooleanOutput>; 32];
}

pub trait Random {
    fn random() -> Self;
}

pub trait RandomBit {
    fn random() -> Self;
}

pub trait Reveal {
    fn reveal(self) -> Self;
}

pub trait Invert {
    fn invert(self, is_expected_non_zero: bool) -> Self;
}

pub trait Pow {
    fn pow(self, e: &Number, is_expected_non_zero: bool) -> Self;
}

pub trait Keccak {
    fn f1600(state: [Byte<Self>; 200]) -> [Byte<Self>; 200]
    where
        Self: Boolean;

    fn sponge<const N: usize>(
        rate: usize,
        capacity: usize,
        input_bytes: Vec<Byte<Self>>,
    ) -> [Byte<Self>; N]
    where
        Self: Boolean,
    {
        if input_bytes.len() > 1 << 20 {
            panic!(
                "sha3 not supported on inputs of more than 2^20 bytes (found {})",
                input_bytes.len()
            );
        }
        if rate + capacity != 1600 || rate % 8 != 0 {
            panic!("rate + capacity must equal 1600 and rate must be a multiple of 8 (found rate: {rate}, capacity: {capacity})");
        }
        let mut state = [Byte::from(0u8); 200];
        let rate_in_bytes = rate / 8;
        // absorb the input blocks
        input_bytes.chunks(rate_in_bytes).for_each(|chunk| {
            chunk.iter().copied().enumerate().for_each(|(i, c)| {
                state[i] ^= c;
            });
            if chunk.len() == rate_in_bytes {
                state = Keccak::f1600(state);
            }
        });
        // do the padding
        let block_size = input_bytes.len() % rate_in_bytes;
        state[block_size] ^= Byte::from(0x06);
        state[rate_in_bytes - 1] ^= Byte::from(0x80);
        state = Keccak::f1600(state);
        // squeezing phase
        (0..N)
            .step_by(rate_in_bytes)
            .fold(Vec::new(), |mut acc, pos| {
                let block_size = (N - pos).min(rate_in_bytes);
                acc.append(&mut state[0..block_size].to_vec());
                if acc.len() < N {
                    state = Keccak::f1600(state);
                }
                acc
            })
            .try_into()
            .unwrap_or_else(|v: Vec<Byte<Self>>| {
                panic!("Expected a Vec of length {N} (found {})", v.len())
            })
    }
}

pub trait WithBooleanBounds {
    fn with_boolean_bounds(&self) -> Self;
}

pub trait ToMontgomery {
    type Output: F25519;

    fn to_montgomery(self, is_expected_non_identity: bool) -> (Self::Output, Self::Output);
}

pub trait MxeX25519PrivateKey {
    fn mxe_x25519_private_key() -> Self;
}

pub trait MxeRescueKey {
    fn mxe_rescue_key(i: usize) -> Self;
}

/// Trait used to convert the ECDH output to the target field.
/// The implementor must make sure that the conversion is injective!
pub trait FromF25519<T: F25519> {
    #[allow(non_snake_case)]
    fn from_F25519(value: T) -> Vec<Self>
    where
        Self: Sized;
}
