use crate::{
    error::Error,
    interpreter::{
        data_store::SimpleAddress,
        fn_data::FnData,
        memory::Memory,
        pointer::Pointer,
        supported_type::{ArrayLen, CurveType, RustType, SupportedType, TypedStruct},
    },
};
use arcis_compiler::{
    core::{
        circuits::boolean::boolean_value::BooleanValue,
        expressions::field_expr::FieldExpr,
        global_value::{curve_value::CurveValue, value::FieldValue},
    },
    traits::{Reveal, Select},
    utils::field::BaseField,
};
use std::{
    collections::HashMap,
    hash::{Hash, Hasher},
    rc::Rc,
};

pub type ArcisField = BaseField;
pub type ArcisValue = FieldValue<ArcisField>;

pub type FieldId = usize;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StructKind {
    Named,
    Tuple,
    Unit,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StructOriginInner {
    User(usize),
}
/// Used for equality between `ArcisStruct`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct StructOrigin(StructOriginInner);

impl StructOrigin {
    pub fn from_user_struct(s: &syn::ItemStruct) -> StructOrigin {
        let ptr = s as *const syn::ItemStruct;
        StructOrigin(StructOriginInner::User(ptr.addr()))
    }
}

#[derive(Debug, Clone)]
pub struct ArcisStruct {
    kind: StructKind,
    /// Used for equality determination.
    origin: StructOrigin,
    name: String,
    generic_id: Option<usize>,
    n_fields: usize,
    named_fields: HashMap<syn::Ident, FieldId>,
    field_types: Vec<SupportedType>,
    field_names: Vec<syn::Ident>,
}
impl ArcisStruct {
    pub fn new(
        kind: StructKind,
        origin: StructOrigin,
        name: String,
        generic_id: Option<usize>,
        field_names: Vec<syn::Ident>,
        field_types: Vec<SupportedType>,
    ) -> Self {
        let named_fields = field_names
            .iter()
            .enumerate()
            .map(|(i, ident)| (ident.clone(), i))
            .collect();
        let n_fields = field_types.len();

        ArcisStruct {
            kind,
            origin,
            name,
            generic_id,
            n_fields,
            named_fields,
            field_types,
            field_names,
        }
    }
    pub fn get_kind(&self) -> StructKind {
        self.kind
    }
    pub fn get_name(&self) -> &str {
        &self.name
    }
    pub fn get_field_id(&self, name: &syn::Ident) -> Option<&FieldId> {
        self.named_fields.get(name)
    }
    pub fn get_n_fields(&self) -> usize {
        self.n_fields
    }
    pub fn get_field_types(&self) -> &[SupportedType] {
        &self.field_types
    }
    pub fn get_field_name(&self, id: FieldId) -> String {
        match self.kind {
            StructKind::Named => self.field_names[id].to_string(),
            StructKind::Tuple => id.to_string(),
            StructKind::Unit => "".to_string(),
        }
    }
    pub fn get_generic_id(&self) -> Option<usize> {
        self.generic_id
    }
}

impl PartialEq for ArcisStruct {
    fn eq(&self, other: &Self) -> bool {
        self.origin == other.origin
            && self.name == other.name
            && self.generic_id == other.generic_id
    }
}

impl Eq for ArcisStruct {}

impl Hash for ArcisStruct {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.name.hash(state);
    }
}
pub type StructId = Rc<ArcisStruct>;

#[derive(Debug, Clone, PartialEq)]
pub enum CollectionKind {
    Array(SupportedType),
    Struct(TypedStruct),
    Tuple,
}

impl CollectionKind {
    pub fn is_compatible(&self, other: &Self) -> bool {
        match (self, other) {
            (CollectionKind::Array(s), CollectionKind::Array(o)) => s.compatible_data_types(o),
            (CollectionKind::Struct(s1), CollectionKind::Struct(s2)) => s1 == s2,
            (CollectionKind::Tuple, CollectionKind::Tuple) => true,
            _ => false,
        }
    }
}

/// A float or integer.
#[derive(Debug, Clone, Copy)]
pub struct ValData(pub ArcisValue, pub RustType);

impl ValData {
    pub fn reveal(&self) -> ValData {
        ValData(ArcisValue::new(FieldExpr::Reveal(self.0)), self.1)
    }

    pub fn merge_data(&self, other: ValData, cond: BooleanValue) -> ValData {
        let rust_type = self
            .1
            .intersect(other.1)
            .expect("Incompatible numerical types are being merged.");
        ValData(cond.select(self.val(), other.val()), rust_type)
    }

    pub fn val(&self) -> ArcisValue {
        self.0
    }

    pub fn info(&self) -> RustType {
        self.1
    }

    fn push_expr_ids(&self, acc: &mut Vec<usize>) {
        acc.push(self.val().get_id());
    }

    fn get_type(&self) -> SupportedType {
        SupportedType::FieldValue(self.info())
    }
}

impl PartialEq for ValData {
    fn eq(&self, other: &Self) -> bool {
        self.val().get_id() == other.val().get_id() && self.info() == other.info()
    }
}

/// A boolean.
#[derive(Debug, Clone, Copy)]
pub struct BoolData(pub BooleanValue);

impl BoolData {
    pub fn reveal(&self) -> BoolData {
        BoolData(self.0.reveal())
    }
    pub fn new_with_cap(v: ArcisValue) -> Self {
        BoolData(ArcisValue::new(FieldExpr::Cap(v, 1)).into())
    }
    pub fn merge_data(&self, other: BoolData, cond: BooleanValue) -> BoolData {
        BoolData(cond.select(self.0, other.0))
    }
}

impl PartialEq for BoolData {
    fn eq(&self, other: &Self) -> bool {
        self.0.get_id() == other.0.get_id()
    }
}

/// A point on the Curve. Not exposed, only used internally.
#[derive(Debug, Clone, Copy)]
pub struct CurveData(pub CurveValue, pub CurveType);

impl CurveData {
    pub fn reveal(&self) -> CurveData {
        CurveData(self.0.reveal(), self.1)
    }

    pub fn merge_data(&self, other: CurveData, cond: BooleanValue) -> CurveData {
        let rust_type = self
            .1
            .intersect(other.1)
            .expect("Incompatible curve types are being merged.");
        CurveData(cond.select(self.0, other.0), rust_type)
    }

    pub fn info(&self) -> CurveType {
        self.1
    }

    fn get_type(&self) -> SupportedType {
        SupportedType::Curve(self.info())
    }
}

impl PartialEq for CurveData {
    fn eq(&self, other: &Self) -> bool {
        self.0.get_id() == other.0.get_id()
    }
}

/// Data: smallest item of memory.
#[derive(Debug, Clone, PartialEq)]
pub enum Data {
    /// Unit Tuple.
    Unit,
    /// Integer or float.
    Val(ValData),
    /// Boolean.
    Bool(BoolData),
    /// Point on curve. Only used internally, not exposed.
    Curve(CurveData),
    /// Reference.
    Pointer(Pointer),
    /// Array, struct, or (non-unit) tuple.
    Collection(Vec<SimpleAddress>, CollectionKind),
    /// Function.
    Fn(FnData),
}

impl Data {
    pub fn get_type(&self, memory: &Memory) -> SupportedType {
        match self {
            Data::Val(s) => s.get_type(),
            Data::Bool(_) => SupportedType::Bool,
            Data::Curve(s) => s.get_type(),
            Data::Pointer(p) => memory.pointed_type(*p).reference(),
            Data::Unit => SupportedType::Unit,
            Data::Collection(a, k) => match k {
                CollectionKind::Array(ty) => {
                    SupportedType::Array(Box::new(ty.clone()), ArrayLen::Computed(a.len()))
                }
                CollectionKind::Struct(s) => SupportedType::Struct(s.clone()),
                CollectionKind::Tuple => {
                    let v = a
                        .iter()
                        .map(|addr| memory.pointed_type(addr.as_pointer()))
                        .collect();
                    SupportedType::Tuple(v)
                }
            },
            Data::Fn(_) => SupportedType::Fn,
        }
    }

    pub fn reveal(&self, memory: &mut Memory) -> Result<Data, Error> {
        let res = match self {
            Data::Unit => Data::Unit,
            Data::Val(v) => Data::Val(v.reveal()),
            Data::Bool(v) => Data::Bool(v.reveal()),
            Data::Curve(v) => Data::Curve(v.reveal()),
            Data::Pointer(p) => {
                let pointed_data = memory.compute(*p)?;
                let revealed_data = pointed_data.reveal(memory)?;
                memory.push(revealed_data).as_data()
            }
            Data::Collection(v, k) => {
                let new_v: Result<Vec<_>, Error> = v
                    .iter()
                    .map(|addr| {
                        let data = memory[*addr].clone();
                        let revealed_data = data.reveal(memory)?;
                        Ok(memory.push(revealed_data))
                    })
                    .collect();
                Data::Collection(new_v?, k.clone())
            }
            Data::Fn(_) => self.clone(),
        };
        Ok(res)
    }
    pub fn to_expr_ids_impl(&self, memory: &mut Memory, res: &mut Vec<usize>) -> Result<(), Error> {
        match self {
            Data::Val(v) => v.push_expr_ids(res),
            Data::Bool(BoolData(b)) => res.push(ArcisValue::from(*b).get_id()),
            Data::Curve(c) => res.push(c.0.get_id()),
            Data::Pointer(p) => {
                let data = memory.compute(*p)?;
                data.to_expr_ids_impl(memory, res)?
            }
            Data::Collection(arr, _) => {
                for addr in arr {
                    let data = memory[*addr].clone();
                    data.to_expr_ids_impl(memory, res)?;
                }
            }
            Data::Unit => {}  // No expressions.
            Data::Fn(_) => {} // No expressions.
        }
        Ok(())
    }
    pub fn to_expr_ids(&self, memory: &mut Memory) -> Result<Vec<usize>, Error> {
        let mut res = Vec::new();
        self.to_expr_ids_impl(memory, &mut res)?;
        Ok(res)
    }
    pub fn can_be_modified(&self) -> bool {
        matches!(
            self,
            Data::Unit | Data::Val(_) | Data::Bool(_) | Data::Curve(_) | Data::Pointer(_)
        )
    }
}
