use std::collections::HashMap;
use std::rc::Rc;

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

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

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

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

#[derive(PartialEq, Eq, Clone, Hash)]
pub enum Type {
    BotType,
    SelfType,
    PrimitiveType(PrimitiveId),
    ClassType(ClassId),
    VirtualClassType(ClassId),
    ParametricType(ParamId, Vec<Rc<Type>>),
    TypeVariable(VariableId),
}

pub struct FnDef {
    pub args: Vec<Type>,
    pub result: Option<Type>,
}

pub struct TypeDef<'a> {
    pub name: &'a str,
    pub fields: HashMap<&'a str, Type>,
    pub methods: HashMap<&'a str, FnDef>,
}

pub struct ClassDef<'a> {
    pub base: TypeDef<'a>,
    pub parents: Vec<ClassId>,
}

pub struct ParametricDef<'a> {
    pub base: TypeDef<'a>,
    pub params: Vec<VariableId>,
}

pub struct VarDef<'a> {
    pub name: &'a str,
    pub bound: Vec<Type>,
}

pub struct GlobalContext<'a> {
    primitive_defs: Vec<TypeDef<'a>>,
    class_defs: Vec<ClassDef<'a>>,
    parametric_defs: Vec<ParametricDef<'a>>,
    var_defs: Vec<VarDef<'a>>,
    sym_table: HashMap<&'a str, Type>,
}

impl<'a> GlobalContext<'a> {
    pub fn new(primitives: Vec<TypeDef<'a>>) -> GlobalContext {
        let mut sym_table = HashMap::new();
        for (i, t) in primitives.iter().enumerate() {
            sym_table.insert(t.name, Type::PrimitiveType(PrimitiveId(i)));
        }
        return GlobalContext {
            primitive_defs: primitives,
            class_defs: Vec::new(),
            parametric_defs: Vec::new(),
            var_defs: Vec::new(),
            sym_table,
        };
    }

    pub fn add_class(&mut self, def: ClassDef<'a>) {
        self.sym_table.insert(
            def.base.name,
            Type::ClassType(ClassId(self.class_defs.len())),
        );
        self.class_defs.push(def);
    }
    pub fn add_parametric(&mut self, def: ParametricDef<'a>) {
        let params = def
            .params
            .iter()
            .map(|&v| Type::TypeVariable(v).into())
            .collect();
        self.sym_table.insert(
            def.base.name,
            Type::ParametricType(ParamId(self.parametric_defs.len()), params),
        );
        self.parametric_defs.push(def);
    }

    pub fn add_variable(&mut self, def: VarDef<'a>) {
        self.sym_table.insert(
            def.name,
            Type::TypeVariable(VariableId(self.var_defs.len())),
        );
        self.var_defs.push(def);
    }

    pub fn add_variable_private(&mut self, def: VarDef<'a>) {
        self.var_defs.push(def);
    }

    pub fn get_primitive_mut(&mut self, id: PrimitiveId) -> &mut TypeDef<'a> {
        self.primitive_defs.get_mut(id.0).unwrap()
    }

    pub fn get_primitive(&self, id: PrimitiveId) -> &TypeDef {
        self.primitive_defs.get(id.0).unwrap()
    }

    pub fn get_class_mut(&mut self, id: ClassId) -> &mut ClassDef<'a> {
        self.class_defs.get_mut(id.0).unwrap()
    }

    pub fn get_class(&self, id: ClassId) -> &ClassDef {
        self.class_defs.get(id.0).unwrap()
    }

    pub fn get_parametric_mut(&mut self, id: ParamId) -> &mut ParametricDef<'a> {
        self.parametric_defs.get_mut(id.0).unwrap()
    }

    pub fn get_parametric(&self, id: ParamId) -> &ParametricDef {
        self.parametric_defs.get(id.0).unwrap()
    }

    pub fn get_variable_mut(&mut self, id: VariableId) -> &mut VarDef<'a> {
        self.var_defs.get_mut(id.0).unwrap()
    }

    pub fn get_variable(&self, id: VariableId) -> &VarDef {
        self.var_defs.get(id.0).unwrap()
    }

    pub fn get_type(&self, name: &str) -> Option<Type> {
        // TODO: change this to handle import
        self.sym_table.get(name).map(|v| v.clone())
    }
}

impl Type {
    pub fn subst(&self, map: &HashMap<VariableId, Type>) -> Type {
        match self {
            Type::TypeVariable(id) => map.get(id).unwrap_or(self).clone(),
            Type::ParametricType(id, params) => Type::ParametricType(
                *id,
                params
                    .iter()
                    .map(|v| v.as_ref().subst(map).into())
                    .collect(),
            ),
            _ => self.clone(),
        }
    }

    pub fn inv_subst(&self, map: &Vec<(Type, Type)>) -> Type {
        for (from, to) in map.iter() {
            if self == from {
                return to.clone();
            }
        }
        match self {
            Type::ParametricType(id, params) => Type::ParametricType(
                *id,
                params
                    .iter()
                    .map(|v| v.as_ref().inv_subst(map).into())
                    .collect(),
            ),
            _ => self.clone(),
        }
    }

    pub fn get_subst(&self, ctx: &GlobalContext) -> HashMap<VariableId, Type> {
        match self {
            Type::ParametricType(id, params) => {
                let vars = &ctx.get_parametric(*id).params;
                vars.iter()
                    .zip(params)
                    .map(|(v, p)| (*v, p.as_ref().clone()))
                    .collect()
            }
            // if this proves to be slow, we can use option type
            _ => HashMap::new(),
        }
    }
}