use crate::{
    error::Error,
    interpreter::{
        data::{BoolData, Data, ValData},
        pointer::Pointer,
        Interpreter,
    },
};
use arcis_compiler::{
    core::circuits::boolean::{boolean_value::BooleanValue, utils::all_circuit},
    traits::Equal,
};
use quote::ToTokens;

fn are_equal_acc(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
    values: &mut Vec<BooleanValue>,
) -> Result<(), Error> {
    let &[arg0, arg1] = args else {
        return Err(Error::new_spanned(expr, "Comparisons have two arguments."));
    };
    let val0 = interpreter.memory.compute(arg0)?;
    let val1 = interpreter.memory.compute(arg1)?;
    let ty0 = val0.get_type(&interpreter.memory);
    let ty1 = val1.get_type(&interpreter.memory);
    match (val0, val1) {
        (Data::Val(ValData(i0, _)), Data::Val(ValData(i1, _))) => values.push(i0.eq(i1)),
        (Data::Bool(BoolData(b0)), Data::Bool(BoolData(b1))) => values.push(!b0 ^ b1),
        (Data::Pointer(p0), Data::Pointer(p1)) => {
            are_equal_acc(interpreter, &[p0, p1], expr, values)?
        }
        (Data::Collection(v0, _), Data::Collection(v1, _)) => {
            if v0.len() != v1.len() {
                return Err(Error::new_spanned(
                    expr,
                    format!("Comparison between different collection types `{ty0}` and `{ty1}`."),
                ));
            }
            for (a0, a1) in v0.into_iter().zip(v1) {
                are_equal_acc(
                    interpreter,
                    &[a0.as_pointer(), a1.as_pointer()],
                    expr,
                    values,
                )?;
            }
        }
        (Data::Unit, Data::Unit) => {}
        (Data::Curve(c0), Data::Curve(c1)) => values.push((c0.0 - c1.0).is_identity()),
        _ => {
            return if ty0 == ty1 {
                Err(Error::new_spanned(
                    expr,
                    format!("Comparison unimplemented for type `{ty0}`."),
                ))
            } else {
                Err(Error::new_spanned(
                    expr,
                    format!("Comparison between different types `{ty0}` and `{ty1}`."),
                ))
            };
        }
    }
    Ok(())
}

fn are_equal_impl(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<BooleanValue, Error> {
    let mut acc = Vec::new();
    are_equal_acc(interpreter, args, expr, &mut acc)?;
    // We want the lowest possible depth among secret-shared values.
    let mut plaintext = Vec::new();
    let mut secret = Vec::new();
    for b in acc {
        if b.is_plaintext() {
            plaintext.push(b);
        } else {
            secret.push(b);
        }
    }
    let plaintext_eq = all_circuit(plaintext, Default::default());
    let secret_eq = all_circuit(secret, Default::default());
    Ok(plaintext_eq & secret_eq)
}

pub fn are_equal(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    let res = are_equal_impl(interpreter, args, expr)?;
    Ok(interpreter
        .memory
        .push(Data::Bool(BoolData(res)))
        .as_pointer())
}

pub fn are_different(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    let res = are_equal_impl(interpreter, args, expr)?;
    Ok(interpreter
        .memory
        .push(Data::Bool(BoolData(!res)))
        .as_pointer())
}
