use crate::{
    error::{Error, SpanRange},
    interpreter::{
        data::{ArcisValue, BoolData, CollectionKind, CurveData, Data, ValData},
        data_store::{DataStore, MemoryLayer, SimpleAddress},
        fn_data::{ConditionalFn, FnData, IndexFn},
        pointer::{ConditionalPointer, ConstantIndexPointer, IndexPointer, Pointer},
        supported_type::{ArrayLen, SupportedType},
    },
};
use arcis_compiler::{
    core::{
        circuits::{arithmetic::index::one_hot_encode, boolean::boolean_value::BooleanValue},
        global_value::{curve_value::CurveValue, value::FieldValue},
    },
    utils::{field::ScalarField, used_field::UsedField},
};
use num_traits::ToPrimitive;
use std::{collections::HashSet, ops::Index};

#[derive(Debug)]
pub struct Memory {
    backend: Box<dyn DataStore>,
}

impl Default for Memory {
    fn default() -> Self {
        Self::new()
    }
}

impl Index<SimpleAddress> for Memory {
    type Output = Data;
    fn index(&self, index: SimpleAddress) -> &Data {
        self.backend.access(index)
    }
}

impl Memory {
    pub const SIMPLE_ADDRESS_OF_UNIT: SimpleAddress = SimpleAddress::zero();

    pub fn new() -> Memory {
        Memory {
            backend: Box::new(vec![Data::Unit]),
        }
    }

    pub fn pointed_type(&self, pointer: Pointer) -> SupportedType {
        self[pointer.similar_address()].get_type(self)
    }

    fn same_type(&self, pointer: Pointer, data: &Data) -> bool {
        self.pointed_type(pointer)
            .compatible_data_types(&data.get_type(self))
    }
    pub fn merge_ptr(
        &mut self,
        cond: BooleanValue,
        then_data: Pointer,
        else_data: Pointer,
    ) -> Pointer {
        if then_data == else_data {
            return then_data;
        }
        let type_addr = then_data.similar_address();
        let then_addr = self.push(Data::Pointer(then_data));
        let else_addr = self.push(Data::Pointer(else_data));
        Pointer::Conditional(ConditionalPointer {
            cond,
            then_ptr_address: then_addr,
            else_ptr_address: else_addr,
            same_type_address: type_addr,
        })
    }
    fn merge_data(
        &mut self,
        cond: BooleanValue,
        then_data: &Data,
        else_data: &Data,
    ) -> Result<Data, String> {
        if *then_data == *else_data {
            return Ok(then_data.clone());
        }
        if let Some(c) = cond.as_constant() {
            return Ok((if c { then_data } else { else_data }).clone());
        }
        let res = match (then_data, else_data) {
            (Data::Unit, Data::Unit) => Data::Unit,
            (Data::Val(t), Data::Val(e)) => Data::Val(t.merge_data(*e, cond)),
            (Data::Bool(t), Data::Bool(e)) => Data::Bool(t.merge_data(*e, cond)),
            (Data::Curve(t), Data::Curve(e)) => Data::Curve(t.merge_data(*e, cond)),
            (Data::Pointer(t), Data::Pointer(e)) => Data::Pointer(self.merge_ptr(cond, *t, *e)),
            (Data::Collection(t, k_t), Data::Collection(e, k_e)) => {
                assert_eq!(t.len(), e.len());
                assert!(k_t.is_compatible(k_e));
                let new_addresses: Result<Vec<_>, _> = t
                    .iter()
                    .zip(e)
                    .map(|(t_addr, e_addr)| {
                        let t_data = self[*t_addr].clone();
                        let e_data = self[*e_addr].clone();
                        let new_data = self.merge_data(cond, &t_data, &e_data)?;
                        Result::<_, String>::Ok(self.push(new_data))
                    })
                    .collect();
                let new_addresses = new_addresses?;
                Data::Collection(new_addresses, k_t.clone())
            }
            (Data::Fn(_), Data::Fn(_)) => {
                let t_addr = self.push(then_data.clone());
                let e_addr = self.push(else_data.clone());
                Data::Fn(FnData::Cond(ConditionalFn::new(cond, t_addr, e_addr)))
            }
            _ => {
                let then_ty = then_data.get_type(self);
                let else_ty = else_data.get_type(self);
                return if then_ty != else_ty {
                    Err(format!(
                        "Incompatible types `{then_ty}` and `{else_ty}` in merge_data."
                    ))
                } else {
                    Err(format!(
                        "Unimplemented conditional merge for type `{then_ty}`."
                    ))
                };
            }
        };
        Ok(res)
    }
    fn merge_array_of_data(
        &mut self,
        ohe: &[BooleanValue],
        data_addresses: &[SimpleAddress],
    ) -> Data {
        assert_eq!(ohe.len(), data_addresses.len());
        assert!(!data_addresses.is_empty());
        let first_data = &self[data_addresses[0]];
        assert!(data_addresses.iter().all(|addr| {
            let res = self.same_type(addr.as_pointer(), first_data);
            if !res {
                panic!(
                    "Different Types in Arrays. item:{}, first item:{}.",
                    self.pointed_type(addr.as_pointer()),
                    first_data.get_type(self)
                );
            }
            res
        }));
        match first_data {
            Data::Unit => Data::Unit,
            Data::Val(_) => {
                let mut rust_type = None;
                let vals: Vec<_> = data_addresses
                    .iter()
                    .map(|addr| {
                        let Data::Val(ValData(val, r)) = self[*addr] else {
                            unreachable!("incompatible types")
                        };
                        rust_type = Some(r);
                        val
                    })
                    .collect();
                let bounds = vals
                    .iter()
                    .map(|x| x.bounds())
                    .fold(ArcisValue::from(0).bounds(), |a, b| a.union(b));
                let result = ohe
                    .iter()
                    .zip(vals)
                    .fold(ArcisValue::from(0), |acc, (ohe, val)| {
                        acc + ArcisValue::from(*ohe) * val
                    });

                Data::Val(ValData(
                    result.with_bounds(bounds),
                    rust_type.expect("Empty array"),
                ))
            }
            Data::Bool(_) => {
                let vals: Vec<_> = data_addresses
                    .iter()
                    .map(|addr| {
                        let Data::Bool(BoolData(val)) = self[*addr] else {
                            unreachable!("incompatible types")
                        };
                        val
                    })
                    .collect();
                let result = ohe
                    .iter()
                    .zip(vals)
                    .fold(BooleanValue::from(false), |acc, (ohe, val)| {
                        acc ^ (*ohe & val)
                    });
                Data::Bool(BoolData(result))
            }
            Data::Curve(_) => {
                let mut rust_type = None;
                let vals: Vec<_> = data_addresses
                    .iter()
                    .map(|addr| {
                        let Data::Curve(CurveData(val, r)) = self[*addr] else {
                            unreachable!("incompatible types")
                        };
                        rust_type = Some(r);
                        val
                    })
                    .collect();
                Data::Curve(CurveData(
                    CurveValue::multiscalar_mul(
                        ohe.iter()
                            .map(|val| FieldValue::from(*val))
                            .collect::<Vec<FieldValue<ScalarField>>>(),
                        vals,
                    ),
                    rust_type.expect("Empty array"),
                ))
            }
            Data::Pointer(p) => {
                let pointers = data_addresses
                    .iter()
                    .map(|addr| {
                        let Data::Pointer(ptr) = self[*addr] else {
                            panic!(
                                "Address is not a pointer, it is a {}.",
                                self[*addr].get_type(self)
                            );
                        };
                        ptr
                    })
                    .collect::<Vec<_>>();
                Data::Pointer(self.merge_array_of_pointers(*p, ohe, &pointers))
            }
            Data::Collection(arr, kind) => {
                let kind = kind.clone();
                let len = arr.len();
                let v = (0..len)
                    .map(|i| {
                        let field_addresses: Vec<_> = data_addresses
                            .iter()
                            .map(|x| {
                                let Data::Collection(a, k) = &self[*x] else {
                                    panic!("incompatible types")
                                };
                                assert_eq!(*k, kind);
                                a[i]
                            })
                            .collect();
                        let data = self.merge_array_of_data(ohe, &field_addresses);
                        self.push(data)
                    })
                    .collect();
                Data::Collection(v, kind)
            }
            Data::Fn(_) => {
                let ohe_addresses = ohe
                    .iter()
                    .map(|x1| self.push(Data::Bool(BoolData(*x1))))
                    .collect::<Vec<_>>();
                let ohe_address = self.push(Data::Collection(
                    ohe_addresses,
                    CollectionKind::Array(SupportedType::boolean()),
                ));
                let fn_address = self.push(Data::Collection(
                    data_addresses.to_vec(),
                    CollectionKind::Array(SupportedType::Fn),
                ));
                Data::Fn(FnData::Index(IndexFn::new(ohe_address, fn_address)))
            }
        }
    }
    pub fn merge_array_of_pointers(
        &mut self,
        p: Pointer,
        ohe: &[BooleanValue],
        pointers: &[Pointer],
    ) -> Pointer {
        let same_type_address = p.similar_address();

        let ptrs = ohe
            .iter()
            .zip(pointers)
            .map(|(a, b)| (*a, *b))
            .collect::<Vec<_>>();
        let v = Pointer::ohe_and_address_plural(ptrs, self);
        let ohe_addresses = v
            .iter()
            .map(|x1| self.push(Data::Bool(BoolData(x1.0))))
            .collect::<Vec<_>>();
        let addresses = v.into_iter().map(|x1| x1.1).collect::<Vec<_>>();
        let ty = self.pointed_type(addresses[0].as_pointer());
        let arr_addr = self.push(Data::Collection(addresses, CollectionKind::Array(ty)));
        let arr_ptr_address = self.push(arr_addr.as_data());
        let ohe_address = self.push(Data::Collection(
            ohe_addresses,
            CollectionKind::Array(SupportedType::boolean()),
        ));
        Pointer::Index(IndexPointer {
            ohe_address,
            arr_ptr_address,
            same_type_address,
        })
    }
    pub fn compute(&mut self, pointer: Pointer) -> Result<Data, Error> {
        let res = match pointer {
            Pointer::Address(u) => self.clone_pointed(u),
            Pointer::Conditional(s) => {
                let then_ptr = &self[s.then_ptr_address];
                let Data::Pointer(then_ptr) = then_ptr else {
                    panic!("Tried to compute non-pointer pointer")
                };
                let then_data = self.compute(*then_ptr)?;
                let else_ptr = &self[s.else_ptr_address];
                let Data::Pointer(else_ptr) = else_ptr else {
                    panic!("Tried to compute non-pointer pointer")
                };
                let else_data = self.compute(*else_ptr)?;
                self.merge_data(s.cond, &then_data, &else_data)
                    .map_err(|err| Error::new(SpanRange::call_site(), err))?
            }
            Pointer::ConstantIndex(i) => {
                let Data::Pointer(ptr) = self[i.arr_ptr_address] else {
                    panic!("Tried to compute non-pointer pointer")
                };
                let Data::Collection(array_data_addresses, _) = self.compute(ptr)? else {
                    panic!("array is not an array")
                };
                let idx = i.index;
                if array_data_addresses.len() <= idx {
                    return Err(Error::new(
                        i.span,
                        format!(
                            "index {idx} out of bounds. It should have been below {}.",
                            array_data_addresses.len()
                        ),
                    ));
                }
                self[array_data_addresses[idx]].clone()
            }
            _ => {
                let v = pointer.ohe_and_address(self, true.into());
                let ohe: Vec<_> = v.iter().map(|x| x.0).collect();
                let data_addresses: Vec<_> = v.into_iter().map(|x| x.1).collect();
                self.merge_array_of_data(&ohe, &data_addresses)
            }
        };
        Ok(res)
    }
    pub fn conditional_assign(
        &mut self,
        pointer: Pointer,
        data: &Data,
        cond: BooleanValue,
    ) -> Result<(), String> {
        if !self.same_type(pointer, data) {
            return Err("Incompatible types in assign.".to_string());
        }
        let mut err = None;
        pointer.for_each(self, cond, &mut |memory, c, addr| {
            let new_data = memory.merge_data(c, data, &memory[addr].clone());
            let new_data = match new_data {
                Ok(new_data) => new_data,
                Err(e) => {
                    err = Some(e);
                    return;
                }
            };
            let res = memory.assign(addr.as_pointer(), new_data);
            res.unwrap_or_else(|e| err = Some(e.to_string()))
        });
        match err {
            None => Ok(()),
            Some(s) => Err(s),
        }
    }
    pub fn assign(&mut self, pointer: Pointer, data: Data) -> Result<(), String> {
        if !self.same_type(pointer, &data) {
            return Err(format!(
                "Incompatible types in assign: {} into {}.",
                data.get_type(self),
                self.pointed_type(pointer)
            ));
        }
        match pointer {
            Pointer::Address(u) => {
                match data {
                    Data::Unit => {
                        // nothing to do
                    }
                    Data::Collection(new_arr, _) => {
                        let Data::Collection(old_arr, _) = self[u].clone() else {
                            unreachable!("not same type in assignment");
                        };
                        assert_eq!(new_arr.len(), old_arr.len());
                        let n = new_arr.len();
                        for i in 0..n {
                            let data = self[new_arr[i]].clone();
                            self.assign(Pointer::Address(old_arr[i]), data)?;
                        }
                    }
                    _ => {
                        self.backend.set(u, data);
                    }
                }
            }
            _ => {
                self.conditional_assign(pointer, &data, true.into())?;
            }
        }
        Ok(())
    }
    pub fn clone_and_push(&mut self, pointer: Pointer) -> Result<SimpleAddress, Error> {
        let data = self.compute(pointer)?;
        let res = self.push(data);
        Ok(res)
    }
    pub fn push(&mut self, data: Data) -> SimpleAddress {
        self.backend.push(data)
    }
    fn clone_and_push_pointed(&mut self, address: SimpleAddress) -> SimpleAddress {
        match self[address].clone() {
            Data::Unit => Self::SIMPLE_ADDRESS_OF_UNIT,
            Data::Collection(arr, k) => {
                let v = arr
                    .iter()
                    .map(|x| self.clone_and_push_pointed(*x))
                    .collect();
                self.push(Data::Collection(v, k))
            }
            y => self.push(y),
        }
    }
    fn clone_pointed(&mut self, address: SimpleAddress) -> Data {
        match self[address].clone() {
            Data::Collection(arr, k) => {
                let v = arr
                    .iter()
                    .map(|x| self.clone_and_push_pointed(*x))
                    .collect();
                Data::Collection(v, k)
            }
            y => y,
        }
    }
    pub fn conditional_merge(
        &mut self,
        cond: BooleanValue,
        then_layer: MemoryLayer,
        else_layer: MemoryLayer,
    ) -> Result<(), String> {
        let changed_addresses: HashSet<_> = then_layer
            .keys()
            .chain(else_layer.keys())
            .cloned()
            .collect();
        for address in changed_addresses {
            let then_data = then_layer
                .get(&address)
                .unwrap_or_else(|| self.backend.access(address))
                .clone();
            let else_data = else_layer
                .get(&address)
                .unwrap_or_else(|| self.backend.access(address))
                .clone();

            assert!(
                then_data.can_be_modified(),
                "Modified unmodifiable type {}",
                then_data.get_type(self)
            );
            assert!(
                else_data.can_be_modified(),
                "Modified unmodifiable type {}",
                else_data.get_type(self)
            );

            let merged_data = self.merge_data(cond, &then_data, &else_data)?;
            self.backend.set(address, merged_data);
        }
        Ok(())
    }
    fn find_underlying_arr_pointer(
        &mut self,
        arr_ptr: Pointer,
        index: ArcisValue,
        span: SpanRange,
    ) -> Result<(Pointer, SimpleAddress, usize), Error> {
        let Data::Collection(arr, _) = &self[arr_ptr.similar_address()] else {
            return if matches!(&self[arr_ptr.similar_address()], Data::Pointer(_)) {
                let ptr = self.compute(arr_ptr)?;
                let Data::Pointer(ptr) = ptr else {
                    panic!(
                        "Pointer wasn't a pointer, when it should have been. It was {}.",
                        ptr.get_type(self)
                    );
                };
                self.find_underlying_arr_pointer(ptr, index, span)
            } else {
                Err(Error::new(span, "Trying to access non-array by index."))
            };
        };
        let Some(idx) = index.bounds().to_unsigned_number_pair().0.to_usize() else {
            return Err(Error::new(span, "index way out of bounds"));
        };
        if idx >= arr.len() {
            return Err(Error::new(span, "index out of bounds"));
        }
        Ok((arr_ptr, arr[0], arr.len()))
    }
    pub fn create_index(
        &mut self,
        arr_ptr: Pointer,
        index: ArcisValue,
        span: SpanRange,
    ) -> Result<Pointer, Error> {
        let (arr_ptr, same_type_address, same_type_len) =
            self.find_underlying_arr_pointer(arr_ptr, index, span)?;
        let arr_ptr_address = self.push(Data::Pointer(arr_ptr));

        if let Some(index) = index.as_constant() {
            let Some(index) = index.to_unsigned_number().to_usize() else {
                return Err(Error::new(
                    span,
                    format!("index out of range, way too high: {:?}", index),
                ));
            };
            if index >= same_type_len {
                return Err(Error::new(
                    span,
                    format!(
                        "Index ({index}) out of range. It should be below {}, a possible array length.",
                        same_type_len
                    ),
                ));
            }
            return if let Pointer::Address(arr_address) = arr_ptr {
                let Data::Collection(arr, _) = &self[arr_address] else {
                    return Err(Error::new(span, "Trying to access non-array by index"));
                };
                if index < arr.len() {
                    Ok(arr[index].as_pointer())
                } else {
                    Err(Error::new(
                        span,
                        format!(
                            "Index ({index}) out of range. It should be below {}, the array length.",
                            arr.len()
                        ),
                    ))
                }
            } else {
                Ok(Pointer::ConstantIndex(ConstantIndexPointer {
                    index,
                    arr_ptr_address,
                    same_type_address,
                    span,
                }))
            };
        }
        let SupportedType::Array(_, ArrayLen::Computed(n)) = self.pointed_type(arr_ptr) else {
            return Err(Error::new(span, "Indexed object is not an array."));
        };

        let ohe_vals = one_hot_encode(index, n);
        let v: Vec<_> = ohe_vals
            .into_iter()
            .map(|ohe_val| self.push(Data::Bool(BoolData(ohe_val))))
            .collect();
        let ohe_address = self.push(Data::Collection(
            v,
            CollectionKind::Array(SupportedType::boolean()),
        ));
        Ok(Pointer::Index(IndexPointer {
            ohe_address,
            arr_ptr_address,
            same_type_address,
        }))
    }
    pub fn create_slice(
        &mut self,
        arr_ptr: Pointer,
        begin: usize,
        end: usize, // inclusive
        span: SpanRange,
    ) -> Result<Pointer, Error> {
        if end < begin {
            return Err(Error::new(span, "Slice is empty."));
        }
        let (arr_ptr, _, same_type_len) =
            self.find_underlying_arr_pointer(arr_ptr, begin.into(), span)?;
        if end >= same_type_len {
            return Err(Error::new(span, "End of range is too big."));
        }
        arr_ptr
            .map(self, &mut |memory, addr| {
                let Data::Collection(arr, ty) = &memory[addr] else {
                    return Err("Not an array.");
                };
                let new_addr = memory.push(Data::Collection(arr[begin..=end].to_vec(), ty.clone()));
                Ok(new_addr)
            })
            .map_err(|_| Error::new(span, "Trying to slice non-slice by range."))
    }
    pub fn swap_backend(&mut self, new_backend: &mut Box<dyn DataStore>) {
        std::mem::swap(&mut self.backend, new_backend);
    }
    pub fn get_underlying_data(&mut self, pointer: Pointer) -> Result<Data, Error> {
        let mut data = self.compute(pointer)?;
        while let Data::Pointer(ptr) = data {
            data = self.compute(ptr)?;
        }
        Ok(data)
    }
    pub fn in_conditional_execution(&self) -> bool {
        self.backend.has_layer()
    }
    pub fn deep_clone(&mut self, simple_address: SimpleAddress) -> Result<SimpleAddress, Error> {
        let data = self[simple_address].clone();
        let res = match data {
            Data::Unit => Memory::SIMPLE_ADDRESS_OF_UNIT,
            Data::Pointer(d) => {
                let addr = self.clone_and_push(d)?;
                let new_addr = self.deep_clone(addr)?;
                self.push(new_addr.as_data())
            }
            Data::Collection(s, k) => {
                let v = s
                    .into_iter()
                    .map(|addr| self.deep_clone(addr))
                    .collect::<Result<Vec<_>, _>>()?;
                self.push(Data::Collection(v, k))
            }
            _ => self.push(data),
        };
        Ok(res)
    }
}
