use crate::{
    error::SpanRange,
    interpreter::{
        data::{BoolData, CollectionKind, Data},
        data_store::SimpleAddress,
        memory::Memory,
        supported_type::SupportedType,
    },
};
use arcis_compiler::core::circuits::boolean::boolean_value::BooleanValue;

#[derive(Debug, Clone, Copy)]
pub struct ConditionalPointer {
    pub cond: BooleanValue,
    /// if cond is true, then it is equal to the pointer at that address
    pub then_ptr_address: SimpleAddress,
    /// if cond is false, then it is equal to the pointer at that address
    pub else_ptr_address: SimpleAddress,
    /// an address containing an data of the same type as the pointer
    pub same_type_address: SimpleAddress,
}

impl PartialEq for ConditionalPointer {
    fn eq(&self, other: &Self) -> bool {
        self.cond.get_id() == other.cond.get_id()
            && self.then_ptr_address == other.then_ptr_address
            && self.else_ptr_address == other.else_ptr_address
    }
}

#[derive(Debug, Clone, Copy)]
pub struct IndexPointer {
    /// address of the array containing one-hot-encoding
    pub ohe_address: SimpleAddress,
    /// address of the pointer to the array of data (same size as ohe_address)
    pub arr_ptr_address: SimpleAddress,
    /// an address containing an data of the same type as the pointer
    pub same_type_address: SimpleAddress,
}

impl PartialEq for IndexPointer {
    fn eq(&self, other: &Self) -> bool {
        self.ohe_address == other.ohe_address && self.arr_ptr_address == other.arr_ptr_address
    }
}
#[derive(Debug, Clone, Copy)]
pub struct ConstantIndexPointer {
    pub index: usize,
    /// address of the pointer to the array of data (same size as ohe_address)
    pub arr_ptr_address: SimpleAddress,
    /// an address containing an data of the same type as the pointer
    pub same_type_address: SimpleAddress,
    /// The Span of the indexing. Useful for bounds checking failure handling.
    pub span: SpanRange,
}

impl PartialEq for ConstantIndexPointer {
    fn eq(&self, other: &Self) -> bool {
        self.index == other.index && self.arr_ptr_address == other.arr_ptr_address
    }
}

/// Reference. This type is used to represent references.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Pointer {
    /// A simple reference that points to an address.
    Address(SimpleAddress),
    /// A reference that points to an address or another depending on a condition.
    /// Used for instance in `if c {&a} else {&b}`.
    Conditional(ConditionalPointer),
    /// A reference that points to a specific position within an array, a struct or a tuple.
    /// Used for instance in `(if c {&mut a} else {&mut b}).0 = 3;`.
    ConstantIndex(ConstantIndexPointer),
    /// A reference that points to a non-compile-time-known position within an array.
    /// Used for instance in `a[i] = 5;`.
    Index(IndexPointer),
}

impl Pointer {
    pub fn as_data(self) -> Data {
        Data::Pointer(self)
    }
    pub fn similar_address(&self) -> SimpleAddress {
        match self {
            Pointer::Address(a) => *a,
            Pointer::Conditional(ptr) => ptr.same_type_address,
            Pointer::Index(ptr) => ptr.same_type_address,
            Pointer::ConstantIndex(ptr) => ptr.same_type_address,
        }
    }

    pub fn ohe_and_address(
        self,
        memory: &mut Memory,
        cond: BooleanValue,
    ) -> Vec<(BooleanValue, SimpleAddress)> {
        let mut v = Vec::new();
        self.for_each(memory, cond, &mut |_, c, addr| {
            v.push((c, addr));
        });
        v
    }
    pub fn ohe_and_address_plural(
        addresses: Vec<(BooleanValue, Pointer)>,
        memory: &mut Memory,
    ) -> Vec<(BooleanValue, SimpleAddress)> {
        let mut v = Vec::new();
        for (cond, ptr) in addresses {
            ptr.for_each(memory, cond, &mut |_, c, addr| {
                v.push((c, addr));
            })
        }
        v
    }

    pub fn for_each(
        self,
        memory: &mut Memory,
        cond: BooleanValue,
        f: &mut impl FnMut(&mut Memory, BooleanValue, SimpleAddress),
    ) {
        match self {
            Pointer::Address(addr) => {
                f(memory, cond, addr);
            }
            Pointer::Conditional(s) => {
                let Data::Pointer(then_ptr) = memory[s.then_ptr_address] else {
                    panic!("non-pointer conditional pointer");
                };
                then_ptr.for_each(memory, cond & s.cond, f);
                let Data::Pointer(else_ptr) = memory[s.else_ptr_address] else {
                    panic!("non-pointer conditional pointer");
                };
                else_ptr.for_each(memory, cond & !s.cond, f);
            }
            Pointer::ConstantIndex(i) => {
                let Data::Pointer(arr_ptr) = memory[i.arr_ptr_address] else {
                    panic!("non-pointer index array pointer");
                };
                let v = arr_ptr.ohe_and_address(memory, cond);
                for (c, addr) in v {
                    let Data::Collection(v, _) = &memory[addr] else {
                        panic!("array is not an array")
                    };
                    let addr = v[i.index];
                    f(memory, c, addr);
                }
            }
            Pointer::Index(i) => {
                let Data::Pointer(arr_ptr) = memory[i.arr_ptr_address] else {
                    panic!("non-pointer index array pointer");
                };
                let Data::Collection(arr, CollectionKind::Array(_)) = &memory[i.ohe_address] else {
                    panic!("ohe is not an array")
                };
                let ohe: Vec<_> = arr
                    .iter()
                    .map(|x| {
                        let Data::Bool(BoolData(val)) = memory[*x] else {
                            panic!("ohe is not an array of Bools.")
                        };
                        val
                    })
                    .collect();
                let v = arr_ptr.ohe_and_address(memory, cond);
                for (c, addr) in v {
                    let Data::Collection(v, _) = &memory[addr] else {
                        panic!("array is not an array")
                    };
                    let n = v.len().min(ohe.len());
                    let v = v[..n].to_vec();
                    for (a, b) in ohe[..n].iter().zip(v) {
                        f(memory, c & *a, b);
                    }
                }
            }
        }
    }
    pub fn map<E>(
        self,
        memory: &mut Memory,
        f: &mut impl FnMut(&mut Memory, SimpleAddress) -> Result<SimpleAddress, E>,
    ) -> Result<Pointer, E> {
        let (ohe, addresses): (Vec<_>, Vec<_>) = self
            .ohe_and_address(memory, true.into())
            .into_iter()
            .map(|(c, addr)| (c, f(memory, addr)))
            .collect();
        let addresses = addresses.into_iter().collect::<Result<Vec<_>, _>>()?;
        let ohe_addresses = ohe
            .iter()
            .map(|b| memory.push(Data::Bool(BoolData(*b))))
            .collect::<Vec<_>>();
        let ohe_address = memory.push(Data::Collection(
            ohe_addresses,
            CollectionKind::Array(SupportedType::boolean()),
        ));
        let same_type_address = addresses[0];
        let ty = memory.pointed_type(same_type_address.as_pointer());
        let arr_addr = memory.push(Data::Collection(addresses, CollectionKind::Array(ty)));
        let arr_ptr_address = memory.push(arr_addr.as_data());

        Ok(Pointer::Index(IndexPointer {
            ohe_address,
            arr_ptr_address,
            same_type_address,
        }))
    }
}
