use crate::{
    error::{Error, SpanRange},
    interpreter::{
        data::{ArcisField, ArcisValue, BoolData, Data, ValData},
        memory::Memory,
        pointer::Pointer,
        supported_type::{ArrayLen, RustType, SupportedType},
        Interpreter,
    },
};
use arcis_compiler::{
    core::circuits::{
        arithmetic::sort::{batcher_odd_even_mergesort, MPCSortable, MPCSwappable, Shuffle},
        boolean::boolean_value::BooleanValue,
    },
    traits::{FromLeBits, RandomBit, Reveal},
    utils::used_field::UsedField,
};
use num_traits::ToPrimitive;
use quote::ToTokens;

struct InterpreterSort<'a, 'ast> {
    interpreter: &'a mut Interpreter<'ast>,
    arr_pointer: Pointer,
    len: usize,
}

impl<'a, 'ast> InterpreterSort<'a, 'ast> {
    fn new(interpreter: &'a mut Interpreter<'ast>, arr_pointer: Pointer) -> Result<Self, Error> {
        let pointed_type = interpreter.memory.pointed_type(arr_pointer);
        let &SupportedType::Array(_, ArrayLen::Computed(len)) = pointed_type.underlying_type()
        else {
            return Err(Error::new(
                SpanRange::call_site(),
                "Sort called on a non-array",
            ));
        };
        Ok(Self {
            interpreter,
            arr_pointer,
            len,
        })
    }
}

impl MPCSwappable<ArcisField> for InterpreterSort<'_, '_> {
    fn conditional_swap(
        &mut self,
        condition: BooleanValue,
        first_index: usize,
        second_index: usize,
    ) {
        let first_ptr = self
            .interpreter
            .memory
            .create_index(self.arr_pointer, first_index.into(), SpanRange::call_site())
            .unwrap();
        let second_ptr = self
            .interpreter
            .memory
            .create_index(
                self.arr_pointer,
                second_index.into(),
                SpanRange::call_site(),
            )
            .unwrap();
        let first_val = self.interpreter.memory.compute(first_ptr).unwrap();
        let second_val = self.interpreter.memory.compute(second_ptr).unwrap();
        // could be optimized to use one less Beaver in the simple array of value case.
        self.interpreter
            .memory
            .conditional_assign(first_ptr, &second_val, condition)
            .expect("Typing error when shuffling array ??");
        self.interpreter
            .memory
            .conditional_assign(second_ptr, &first_val, condition)
            .expect("Typing error when shuffling array ???");
    }

    fn len(&self) -> usize {
        self.len
    }
}

impl MPCSortable<ArcisField> for InterpreterSort<'_, '_> {
    fn get_val(&mut self, index: usize) -> ArcisValue {
        let ptr = self
            .interpreter
            .memory
            .create_index(self.arr_pointer, index.into(), SpanRange::call_site())
            .unwrap();
        let Data::Val(ValData(v, _)) = self.interpreter.memory.compute(ptr).unwrap() else {
            panic!("Sort called on a non-integer array");
        };
        v
    }
}

pub fn sort_int(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    _expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    let mut sort = InterpreterSort::new(interpreter, args[0])?;
    batcher_odd_even_mergesort(&mut sort, true);
    Ok(Memory::SIMPLE_ADDRESS_OF_UNIT.as_pointer())
}

pub fn shuffle(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    _expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    let sort = InterpreterSort::new(interpreter, args[0])?;
    let mut shuffle = Shuffle::new(sort);
    shuffle.mpc_shuffle();
    Ok(Memory::SIMPLE_ADDRESS_OF_UNIT.as_pointer())
}

pub fn gen_bool(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    if !args.is_empty() {
        return Err(Error::new_spanned(expr, "This should have 0 arguments."));
    }
    let val = BooleanValue::random();
    Ok(interpreter
        .memory
        .push(Data::Bool(BoolData(val)))
        .as_pointer())
}
pub fn gen_integer_from_width(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    if args.len() != 1 {
        return Err(Error::new_spanned(expr, "This should have 1 argument."));
    }
    let arg = args[0];
    let Data::Val(ValData(v, _)) = interpreter.memory.compute(arg)? else {
        return Err(Error::new_spanned(
            expr,
            "This function expects an integer argument.",
        ));
    };
    let Some(width) = v.as_constant() else {
        return Err(Error::new_spanned(
            expr,
            "This function expects a compile-time known integer argument.",
        ));
    };
    let Some(width) = width.to_unsigned_number().to_usize() else {
        return Err(Error::new_spanned(
            expr,
            "This function's argument is way too big. Only 0 to 128 (included) are supported",
        ));
    };
    if width > 128 {
        return Err(Error::new_spanned(
            expr,
            "This function expects an argument that's not too big (<= 128).",
        ));
    }
    let val = ArcisValue::from_le_bits(
        (0..width)
            .map(|_| BooleanValue::random())
            .collect::<Vec<BooleanValue>>(),
        false,
    );
    Ok(interpreter
        .memory
        .push(Data::Val(ValData(val, RustType::Integer(None))))
        .as_pointer())
}

pub fn gen_public_integer_from_width(
    interpreter: &mut Interpreter,
    args: &[Pointer],
    expr: &dyn ToTokens,
) -> Result<Pointer, Error> {
    let Pointer::Address(addr) = gen_integer_from_width(interpreter, args, expr)? else {
        unreachable!()
    };
    let Data::Val(ValData(val, _)) = interpreter.memory[addr] else {
        unreachable!()
    };
    Ok(interpreter
        .memory
        .push(Data::Val(ValData(val.reveal(), RustType::Integer(None))))
        .as_pointer())
}
