use crate::interpreter::{
    consts::HandledConst,
    data::{Data, ValData},
    data_store::SimpleAddress,
    memory::Memory,
    supported_type::{ArrayLen, RustType, SupportedType},
};
use std::collections::HashMap;

#[derive(Debug)]
pub enum MatchKind<'a> {
    NoDeref,
    CountDeref(&'a mut usize),
}

impl MatchKind<'_> {
    pub fn may_deref(&self) -> bool {
        matches!(self, MatchKind::CountDeref(_))
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct VarId(usize);

#[derive(Debug, Clone, PartialEq)]
struct VarContainer<T: Clone> {
    map: HashMap<VarId, T>,
    n_vars: usize,
}

impl<T: Clone> Default for VarContainer<T> {
    fn default() -> Self {
        let map = Default::default();
        let n_vars = 0;
        Self { map, n_vars }
    }
}

impl<T: Clone> VarContainer<T> {
    fn new_var(&mut self) -> VarId {
        let result = VarId(self.n_vars);
        self.n_vars += 1;
        result
    }
    fn get(&self, id: VarId) -> Option<T> {
        self.map.get(&id).cloned()
    }
    fn set(&mut self, id: VarId, value: T) -> Result<(), ()> {
        let opt = self.map.insert(id, value);
        if opt.is_none() {
            Ok(())
        } else {
            Err(())
        }
    }
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct GenericContext {
    const_vars: VarContainer<HandledConst>,
    type_vars: VarContainer<SupportedType>,
}

impl GenericContext {
    fn matching_array_lens(
        &mut self,
        memory: &mut Memory,
        s_len: ArrayLen,
        t_len: ArrayLen,
    ) -> bool {
        match (s_len, t_len) {
            (ArrayLen::Computed(s), ArrayLen::Computed(t)) => s == t,
            (ArrayLen::Computed(_), ArrayLen::ConstVar(_)) => {
                // weird case: TODO emit ERR
                false
            }
            (ArrayLen::ConstVar(s), ArrayLen::Computed(t)) => {
                match self.const_vars.get(s) {
                    None => {
                        let addr = memory.push(Data::Val(ValData(t.into(), RustType::usize())));
                        let _ = self.const_vars.set(s, HandledConst::Computed(addr));
                        true
                    }
                    Some(_) => {
                        // weird case: TODO emit ERR
                        false
                    }
                }
            }
            (ArrayLen::ConstVar(s), ArrayLen::ConstVar(t)) => {
                match self.const_vars.get(s) {
                    None => {
                        let _ = self.const_vars.set(s, HandledConst::Generic(t));
                        true
                    }
                    Some(_) => {
                        // weird case: TODO emit ERR
                        false
                    }
                }
            }
        }
    }

    #[must_use]
    pub fn includes_type(
        &mut self,
        memory: &mut Memory,
        ty: &SupportedType,
        other: &SupportedType,
        mut kind: MatchKind<'_>,
    ) -> bool {
        match (ty, other) {
            (SupportedType::ImplAny, _) => true,
            (SupportedType::Struct(s), SupportedType::Struct(t)) => {
                s.id() == t.id()
                    && s.types()
                        .iter()
                        .zip(t.types())
                        .all(|(s, t)| self.includes_type(memory, s, t, MatchKind::NoDeref))
            }
            (SupportedType::FieldValue(s), SupportedType::FieldValue(t)) => s.intersect(*t).is_ok(),
            (SupportedType::Bool, SupportedType::Bool) => true,
            (SupportedType::Curve(s), SupportedType::Curve(t)) => s.intersect(*t).is_ok(),
            (SupportedType::Ref(s), SupportedType::Ref(t)) => {
                self.includes_type(memory, s, t, kind)
            }
            (SupportedType::Array(s, s_len), SupportedType::Array(t, t_len)) => {
                self.includes_type(memory, s, t, MatchKind::NoDeref)
                    && self.matching_array_lens(memory, *s_len, *t_len)
            }
            (SupportedType::Slice(s, _), SupportedType::Array(t, _)) => {
                self.includes_type(memory, s, t, MatchKind::NoDeref)
            }
            (SupportedType::Slice(s, _), SupportedType::Slice(t, _)) => {
                self.includes_type(memory, s, t, MatchKind::NoDeref)
            }
            (SupportedType::Unit, SupportedType::Unit) => true,
            (SupportedType::Fn, SupportedType::Fn) => true,
            (SupportedType::Tuple(s), SupportedType::Tuple(t)) => s
                .iter()
                .zip(t)
                .all(|(s, t)| self.includes_type(memory, s, t, MatchKind::NoDeref)),
            (SupportedType::Var(idx), SupportedType::Var(idx_2)) if idx == idx_2 => true,
            (SupportedType::Var(idx), _) => {
                match self.type_vars.get(*idx) {
                    None => {
                        let _ = self.type_vars.set(*idx, other.clone());
                        true
                    } // Maybe ?
                    Some(s) => self.includes_type(memory, &s, other, kind),
                }
            }
            (_, SupportedType::Var(idx)) => {
                match self.type_vars.get(*idx) {
                    None => false, // Maybe ?
                    Some(s) => self.includes_type(memory, ty, &s, kind),
                }
            }
            (_, SupportedType::Ref(s)) if kind.may_deref() => {
                let MatchKind::CountDeref(ref mut count) = kind else {
                    unreachable!("We just checked it may deref.")
                };
                **count += 1;
                self.includes_type(memory, ty, s, kind)
            }
            _ => {
                if ty == other {
                    // returning true would work, but probably lead to bad cases being forgotten.
                    panic!("Unimplemented `includes_type` for `{}`.", ty);
                } else {
                    false
                }
            }
        }
    }
    pub fn set_type_var(&mut self, idx: VarId, ty: SupportedType) -> Result<(), &'static str> {
        self.type_vars
            .set(idx, ty)
            .map_err(|_| "Type generic already set.")
    }
    pub fn set_const_var(&mut self, idx: VarId, val: HandledConst) -> Result<(), &'static str> {
        self.const_vars
            .set(idx, val)
            .map_err(|_| "Constant generic already set.")
    }
    pub fn get_type_var(&self, idx: VarId) -> Option<SupportedType> {
        self.type_vars.get(idx)
    }
    pub fn get_const_var(&self, idx: VarId) -> Option<HandledConst> {
        self.const_vars.get(idx)
    }
    pub fn get_const_var_addr(&self, idx: VarId) -> Option<SimpleAddress> {
        match self.const_vars.get(idx)? {
            HandledConst::Computed(s) => Some(s),
            _ => None,
        }
    }

    pub fn new_type_var(&mut self) -> VarId {
        self.type_vars.new_var()
    }
    pub fn new_const_var(&mut self) -> VarId {
        self.const_vars.new_var()
    }
}
