#[cfg(debug_assertions)]
use crate::core::expressions::domain::Domain;
use crate::{
    core::{
        actually_used_field::ActuallyUsedField,
        bounds::{Bounds, IsBounds},
        expressions::{
            bit_expr::{BitExpr, RandomBitId},
            expr::Expr,
            other_expr::OtherExpr,
        },
        global_value::{
            curve_value::CurveValue,
            global_expr_store::with_global_expr_store_as_local,
            value::FieldValue,
        },
        ir_builder::{ExprStore, IRBuilder},
        mxe_input::{MxeBitInput, MxeInput},
    },
    traits::{GetBit, Keccak, RandomBit, Reveal, Select},
    utils::{
        elliptic_curve::ProjectiveEdwardsPoint,
        field::{BaseField, ScalarField},
    },
};
use core::panic;
use rand::Rng;
use std::ops::{BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Not};

pub trait Boolean:
    BitAnd<Output = Self>
    + BitAndAssign
    + BitXor<Output = Self>
    + BitXorAssign
    + Not<Output = Self>
    + RandomBit
    + Reveal
    + Keccak
    + Select<Self, Self, Self>
    + Select<Vec<Self>, Vec<Self>, Vec<Self>>
    + Copy
    + From<bool>
{
}
#[derive(Debug, Clone, Copy)]
pub struct BooleanValue(usize);

impl BooleanValue {
    pub fn new(id: usize) -> Self {
        #[cfg(debug_assertions)]
        with_global_expr_store_as_local(|expr_store| {
            let _ = bool::unwrap(*expr_store.get_bounds(id));
        });

        Self(id)
    }

    pub fn from_expr(expr: Expr<usize>) -> Self {
        Self(with_global_expr_store_as_local(|expr_store| {
            let id = expr_store.new_expr(expr);
            #[cfg(debug_assertions)]
            let _ = bool::unwrap(*expr_store.get_bounds(id));
            id
        }))
    }

    pub fn get_id(&self) -> usize {
        self.0
    }

    pub fn expr(&self) -> Expr<usize> {
        with_global_expr_store_as_local(|expr_store| expr_store.get_expr(self.0).clone())
    }

    pub fn is_plaintext(&self) -> bool {
        with_global_expr_store_as_local(|expr_store| expr_store.get_is_plaintext(self.get_id()))
    }

    pub fn ed25519_secret_key(i: usize) -> Self {
        Self::from_expr(Expr::Other(OtherExpr::MxeKey(MxeInput::Bit(
            MxeBitInput::Ed25519SecretKey(i),
        ))))
    }

    pub fn ed25519_signing_key_hash_prefix(i: usize) -> Self {
        Self::from_expr(Expr::Other(OtherExpr::MxeKey(MxeInput::Bit(
            MxeBitInput::Ed25519SigningKeyHashPrefix(i),
        ))))
    }

    pub fn as_constant(&self) -> Option<bool> {
        let Bounds::Bit(b) =
            with_global_expr_store_as_local(|expr_store| *expr_store.get_bounds(self.get_id()))
        else {
            panic!("BooleanValue has non-boolean bounds.")
        };
        b.as_constant()
    }
}

impl BitAnd for BooleanValue {
    type Output = BooleanValue;

    fn bitand(self, rhs: Self) -> Self::Output {
        Self(with_global_expr_store_as_local(|expr_store| {
            <IRBuilder as ExprStore<BaseField>>::push_bit(
                expr_store,
                BitExpr::And(self.0, rhs.get_id()),
            )
        }))
    }
}

impl BitAndAssign for BooleanValue {
    fn bitand_assign(&mut self, rhs: Self) {
        *self = *self & rhs;
    }
}

impl BitXor for BooleanValue {
    type Output = BooleanValue;

    fn bitxor(self, rhs: Self) -> Self::Output {
        Self(with_global_expr_store_as_local(|expr_store| {
            <IRBuilder as ExprStore<BaseField>>::push_bit(
                expr_store,
                BitExpr::Xor(self.0, rhs.get_id()),
            )
        }))
    }
}

impl BitXorAssign for BooleanValue {
    fn bitxor_assign(&mut self, rhs: Self) {
        *self = *self ^ rhs;
    }
}

impl BitOr for BooleanValue {
    type Output = BooleanValue;

    fn bitor(self, rhs: Self) -> Self::Output {
        !(!self & !rhs)
    }
}

impl BitOrAssign for BooleanValue {
    fn bitor_assign(&mut self, rhs: Self) {
        *self = *self | rhs;
    }
}

impl Not for BooleanValue {
    type Output = BooleanValue;

    fn not(self) -> Self::Output {
        Self(with_global_expr_store_as_local(|expr_store| {
            <IRBuilder as ExprStore<BaseField>>::push_bit(expr_store, BitExpr::Not(self.0))
        }))
    }
}

impl RandomBit for BooleanValue {
    fn random() -> Self {
        Self(with_global_expr_store_as_local(|expr_store| {
            let before_len = expr_store.len();
            let res = expr_store.new_expr(Expr::Bit(BitExpr::Random(RandomBitId::new())));
            let after_len = expr_store.len();
            assert_eq!(after_len, before_len + 1, "Randomness was reused.");
            assert_eq!(res, before_len, "Randomness was reused.");
            res
        }))
    }
}

impl Reveal for BooleanValue {
    fn reveal(self) -> Self {
        Self(with_global_expr_store_as_local(|expr_store| {
            expr_store.new_expr(Expr::Bit(BitExpr::Reveal(self.0)))
        }))
    }
}

impl Select<BooleanValue, BooleanValue, BooleanValue> for BooleanValue {
    fn select(self, x: BooleanValue, y: BooleanValue) -> BooleanValue {
        y ^ self & (x ^ y)
    }
}

impl Select<Vec<BooleanValue>, Vec<BooleanValue>, Vec<BooleanValue>> for BooleanValue {
    fn select(self, x: Vec<BooleanValue>, y: Vec<BooleanValue>) -> Vec<BooleanValue> {
        assert_eq!(x.len(), y.len());
        x.into_iter()
            .zip(y)
            .map(|(bit_x, bit_y)| self.select(bit_x, bit_y))
            .collect::<Vec<BooleanValue>>()
    }
}

impl<F: ActuallyUsedField> Select<FieldValue<F>, FieldValue<F>, FieldValue<F>> for BooleanValue {
    fn select(self, x: FieldValue<F>, y: FieldValue<F>) -> FieldValue<F> {
        FieldValue::<F>::from(self).select(x, y)
    }
}

impl Select<CurveValue, CurveValue, CurveValue> for BooleanValue {
    fn select(self, x: CurveValue, y: CurveValue) -> CurveValue {
        y + FieldValue::<ScalarField>::from(self) * (x - y)
    }
}

impl
    Select<
        ProjectiveEdwardsPoint<FieldValue<BaseField>>,
        ProjectiveEdwardsPoint<FieldValue<BaseField>>,
        ProjectiveEdwardsPoint<FieldValue<BaseField>>,
    > for BooleanValue
{
    fn select(
        self,
        x: ProjectiveEdwardsPoint<FieldValue<BaseField>>,
        y: ProjectiveEdwardsPoint<FieldValue<BaseField>>,
    ) -> ProjectiveEdwardsPoint<FieldValue<BaseField>> {
        ProjectiveEdwardsPoint::new(
            (
                self.select(x.X, y.X),
                self.select(x.Y, y.Y),
                self.select(x.Z, y.Z),
            ),
            x.is_on_curve && y.is_on_curve,
            x.is_ell_torsion && y.is_ell_torsion,
        )
    }
}

impl From<bool> for BooleanValue {
    fn from(val: bool) -> Self {
        Self(with_global_expr_store_as_local(|expr_store| {
            <IRBuilder as ExprStore<BaseField>>::push_bit(expr_store, BitExpr::Val(val))
        }))
    }
}

impl<F: ActuallyUsedField> From<FieldValue<F>> for BooleanValue {
    fn from(val: FieldValue<F>) -> Self {
        let bounds = val.bounds();
        if !bounds.is_arithmetic_boolean() {
            panic!("bounds must be boolean (found {:?})", bounds);
        }
        val.get_bit(0, false)
    }
}

impl Boolean for BooleanValue {}

impl RandomBit for bool {
    fn random() -> Self {
        let rng = &mut rand::thread_rng();
        rng.gen_bool(0.5)
    }
}

impl Reveal for bool {
    fn reveal(self) -> Self {
        self
    }
}

impl<T> Select<T, T, T> for bool {
    fn select(self, a: T, b: T) -> T {
        if self {
            a
        } else {
            b
        }
    }
}

impl Boolean for bool {}
