use std::collections::HashMap;

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

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

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

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

#[derive(PartialEq, Eq, Clone)]
pub enum Type {
    BotType,
    SelfType,
    PrimitiveType(PrimitiveId),
    ClassType(ClassId),
    VirtualClassType(ClassId),
    ParametricType(ParamId, Vec<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)).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) -> Option<&mut TypeDef<'a>> {
        self.primitive_defs.get_mut(id.0)
    }

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

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

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

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

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

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

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

    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: &Option<HashMap<VariableId, Type>>) -> Type {
        if let Some(m) = map {
            match self {
                Type::TypeVariable(id) => m.get(id).unwrap_or(self).clone(),
                Type::ParametricType(id, params) => {
                    Type::ParametricType(*id, params.iter().map(|v| v.subst(map)).collect())
                }
                _ => self.clone(),
            }
        } else {
            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.inv_subst(map)).collect())
            },
            _ => self.clone()
        }
    }

    pub fn get_subst(&self, ctx: &GlobalContext) -> Option<HashMap<VariableId, Type>> {
        match self {
            Type::ParametricType(id, params) => {
                let vars = &ctx.get_parametric(*id).unwrap().params;
                Some(vars.iter().zip(params).map(|(v, p)| (*v, p.clone())).collect())
            },
            _ => None
        }
    }
}