use crate::r#type::{Type, TypeError, Type::*, TypeError::*, PatFail::*, util::*};
use crate::sexp::{SExp, SExp::*, SLeaf::*, util::*};
use std::collections::HashMap;

impl SExp {

    /// Returns the type of valid expressions.
    /// Invalid expressions result in an error.
    ///
    /// Examples of simple expressions and their simple types:
    /// ```rust
    /// use myslip::{
    ///     r#type::{*, Type::*, TypeError::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// assert_eq!(Atom(Int(1)).type_check(), Ok(Integer));
    /// assert_eq!(Atom(False).type_check(),  Ok(Boolean));
    /// ```
    ///
    /// Quotes are given list types:
    /// ```rust
    /// use myslip::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// assert_eq!(
    ///     scons(Quote, scons(1, scons(False, Nil))).type_check(),
    ///     Ok(List(vec![QuoteTy, List(vec![Integer, Boolean])]))
    /// );
    /// ```
    ///
    /// Even though Nil is formatted as an empty list,
    /// it has its own type.
    /// ```rust
    /// use myslip::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// assert_eq!(
    ///     Atom(Nil).type_check(),
    ///     Ok(NilType)
    /// );
    /// ```
    ///
    /// Vectors are kind of special...
    /// ```rust
    /// use myslip::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// assert_eq!(
    ///     scons(Vector, scons(1, scons(3, Nil))).type_check(),
    ///     Ok(List(vec![VecType, vecof(Integer)]))
    /// );
    /// ```
    /// ...so please don't ask what their type is.
    ///
    /// Some common operators get arrow types:
    /// ```rust
    /// use myslip::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// assert_eq!(Atom(Sub).type_check(), Ok(arr(List(vec![Integer, Integer]), Integer)));
    /// assert_eq!(Atom(Div).type_check(), Ok(arr(List(vec![Integer, Integer]), Integer)));
    /// assert_eq!(Atom(Eq) .type_check(), Ok(arr(List(vec![Integer, Integer]), Boolean)));
    /// assert_eq!(Atom(Neq).type_check(), Ok(arr(List(vec![Integer, Integer]), Boolean)));
    /// assert_eq!(Atom(Gt) .type_check(), Ok(arr(List(vec![Integer, Integer]), Boolean)));
    /// assert_eq!(Atom(Lt) .type_check(), Ok(arr(List(vec![Integer, Integer]), Boolean)));
    /// assert_eq!(Atom(Ge) .type_check(), Ok(arr(List(vec![Integer, Integer]), Boolean)));
    /// assert_eq!(Atom(Le) .type_check(), Ok(arr(List(vec![Integer, Integer]), Boolean)));
    /// assert_eq!(Atom(And).type_check(), Ok(arr(List(vec![Boolean, Boolean]), Boolean)));
    /// assert_eq!(Atom(Or) .type_check(), Ok(arr(List(vec![Boolean, Boolean]), Boolean)));
    /// assert_eq!(Atom(Xor).type_check(), Ok(arr(List(vec![Boolean, Boolean]), Boolean)));
    /// assert_eq!(Atom(Not).type_check(), Ok(arr(Boolean, Boolean)));
    /// ```
    ///
    ///
    /// **Let-binding types**
    ///
    /// Let-bindings are quite interesting. Implementation is still a bit uncertain,
    /// but at least we know that
    /// ```rust
    /// use myslip::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// assert_eq!(
    ///     scons(
    ///         scons(Let, scons(var("x"), scons(False, Nil))),
    ///         scons(var("x"), Nil)
    ///     ).type_check(),
    ///     Ok(Boolean)
    /// );
    /// ```
    ///
    /// **Functions**
    ///
    /// One variable:
    /// ```rust
    /// use myslip::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// let varlist = scons(var("a"), Nil);
    /// let typelist = Atom(Ty(Integer));
    /// let ret = Atom(Ty(Integer));
    /// let body = scons(Add, scons(var("a"), scons(1, Nil)));
    /// assert_eq!(
    ///     scons(Fun, scons(varlist, scons(typelist, scons(ret, scons(body, Nil))))).type_check(),
    ///     Ok(arr(Integer, Integer))
    /// );
    /// ```
    /// Two-variable:
    /// ```rust
    /// use myslip::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// let varlist = scons(var("a"), scons(var("b"), Nil));
    /// let typelist = Atom(Ty(List(vec![Integer, Integer])));
    /// let ret = Atom(Ty(Boolean));
    /// let body = scons(Eq, varlist.clone());
    /// assert_eq!(
    ///     scons(Fun, scons(varlist, scons(typelist, scons(ret, scons(body, Nil))))).type_check(),
    ///     Ok(arr(List(vec![Integer, Integer]), Boolean))
    /// );
    /// ```
    /// Only keyword shouldnt panic:
    /// ```rust
    /// use myslip::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    /// Atom(Fun).type_check();
    /// ```
    ///
    /// **Pattern matching**
    ///
    /// If this works, I guess it seems to work...
    /// ```rust
    /// use myslip::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    ///     parse::parsetree::parse_to_ast
    /// };
    /// 
    /// assert_eq!(
    ///     parse_to_ast(
    ///         "+ (case (quote 1 2 3 true false) ((x 2 3 false true) (+ x 1)) ((0 0 0 true true) 0) (_ 1)) 0"
    ///     ).unwrap().type_check(),
    ///     Ok(Integer)
    /// );
    /// ```
    /// 
    ///
    /// Though perhaps the most important task of the type system
    /// is to increase safety by being able to warn about errors
    /// before evaluation. Here are some failing examples:
    /// ```rust
    /// use myslip::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// match scons(Sub, scons(1, scons(2, scons(3, Nil)))).type_check() {
    ///     Err(InvalidArgList { .. }) => (),
    ///     _ => panic!(
    ///         "passing over 2 arguments to '-' should result in InvalidArgList error"
    ///     ),
    /// };
    ///
    /// match scons(1, scons(2, Nil)).type_check() {
    ///     Err(InvalidOperator { .. }) => (),
    ///     _ => panic!(
    ///         "'1' as an operator should result in InvalidOperator error"
    ///     ),
    /// };
    ///
    /// match scons(Add, scons(Sub, scons(1, Nil))).type_check() {
    ///     Err(InvalidArgList { .. }) => (),
    ///     _ => panic!(
    ///         "passing '-' as an argument to '+' should return in InvalidArgList error"
    ///     ),
    /// };
    ///
    /// assert!(scons(And, scons(1, scons(Atom(True), Nil))).type_check().is_err());
    /// assert!(scons(Mul, scons(1, scons(Atom(True), Nil))).type_check().is_err());
    /// assert!(scons(Not, scons(1, Nil)).type_check().is_err());
    ///
    /// assert!(scons(Vector, scons(1, scons(True, Nil))).type_check().is_err());
    /// ```
    ///
    /// Also, free variables should result in an error
    /// as their type can't be known by the type checker.
    /// ```rust
    /// use myslip::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// match scons(Quote, scons(var("a"), Nil)).type_check() {
    ///     Err(UndefinedVariable(a)) if &a == "a" => (),
    ///     _ => panic!(
    ///         "passing a free variable in type check should result in UndefinedVariable error"
    ///     ),
    /// };
    /// ```
    pub fn type_check(&self) -> Result<Type, TypeError> {
	let res = self.infer_type(HashMap::new());
	match res {
	    Ok(res) => match res.is_concrete() {
		Ok(()) => Ok(res),
		Err(name) => Err(UnboundGeneric(name, res)),
	    },
	    e => e,
	}
    }


    pub fn infer_type(&self, mut ctx: HashMap<String, Type>) -> Result<Type, TypeError> {

        match self {

            Atom(Int(_)) => Ok(Integer),
            Atom(True | False) => Ok(Boolean),
            Atom(Var(name)) => ctx.get(name)
                                  .ok_or(UndefinedVariable(name.to_string()))
                                  .cloned(),
            Atom(Add) => Ok(arr(vecof(Integer), Integer)),
            Atom(Mul) => Ok(arr(vecof(Integer), Integer)),
            Atom(Sub) => Ok(arr(List(vec![Integer, Integer]), Integer)),
            Atom(Div) => Ok(arr(List(vec![Integer, Integer]), Integer)),
            Atom(Eq | Neq | Lt | Gt | Le | Ge) =>
                Ok(arr(List(vec![Integer, Integer]), Boolean)),
            Atom(Or | And | Xor) => Ok(arr(List(vec![Boolean, Boolean]), Boolean)),
            Atom(Not)   => Ok(arr(Boolean, Boolean)),
            Atom(Nil)   => Ok(NilType),
            Atom(Quote) => Ok(arr(
                               vt("T"),
                               List(vec![QuoteTy, vt("T")])
                           )),
            Atom(Vector) => Ok(arr(
                                vecof(vt("T")),
                                List(vec![VecType, vecof(vt("T"))])
                            )),
            Atom(Let)    => Ok(LetType),
	    Atom(Print)  => Ok(arr(vt("_"), List(vec![]))),
	    Atom(Ty(_))  => Ok(TypeLit),
	    Atom(Fun)    => Err(FunAsAtom),
	    Atom(Case)   => Err(CaseAsAtom),
	    Atom(RestPat(_)) => Err(RestAsAtom),

            SCons(op, l) => {

                // Let-expressions
                if let Some((varname, val)) = (*op).clone().check_let() {
                    let valtype = val.infer_type(ctx.clone())?;
                    ctx.insert(varname, valtype);
                    return match (**l).clone() {
                        SCons(exp, nil) if *nil == Atom(Nil) => *exp,
                        t => t
                    }.infer_type(ctx);
                }
                if **op == Atom(Let) {
                    return Err(LetAsOperator(scons(op.clone(), l.clone())));
                }

		// Anonymous functions
		if scons(op.clone(), l.clone()).is_fun() {
		    return scons(op.clone(), l.clone()).get_fun_type(ctx);
		}

		// Nil vector
		if (**op).clone() == Atom(Nil) {
		    return match (**l).clone() {
			SCons(v, n) if *n == Atom(Nil) => match *v {
			    Atom(Ty(t)) => Ok(vecof(t)),
			    _ => Err(OtherError),
			},
			_ => Err(OtherError)
		    }
		}

		// Case expressions
		if let Some((scrutinee, patarms)) = scons(op.clone(), l.clone()).check_case() {
		    let scruty = scrutinee.infer_type(ctx.clone())?;
		    let scruty = match scruty {
			List(v) if (
			    v.get(0) == Some(&QuoteTy)
				|| v.get(0) == Some(&VecType)
			) && v.get(1).is_some() =>
			    v[1].clone(),
			t => t,
		    };
		    let mut ty: Option<Type> = None;
		    let mut has_wildcard = false;
		    for patandarm in patarms {
			let (pat, arm) = match patandarm {
			    SCons(pat, arm) => Ok((*pat, *arm)),
			    _ => Err(InvalidPattern(NoArm(patandarm))),
			}?;
			let arm = match arm {
			    SCons(x, n) if *n == Atom(Nil) => *x,
			    t => t,
			};
			if let Atom(Var(_)) = pat {
			    has_wildcard = true;
			}
			let mut newctx = ctx.clone();
			for (name, ty) in pat.matches_type(&scruty)? {
			    newctx.insert(name, ty);
			}
			match &ty {
			    None => ty = Some(arm.infer_type(newctx)?),
			    Some(t) => ty = Some(t.least_general_supertype(&arm.infer_type(newctx.clone())?)),
			}
		    }
		    if !has_wildcard {
			return Err(NoWildcardInCase(scrutinee));
		    }
		    match ty {
			Some(t) => return Ok(t),
			None => return Err(OtherError),
		    }
		}

                // Normal operation
                let opertype = (*op).infer_type(ctx.clone())?;
                let argstype =  (*l).infer_list_type(ctx)?;

		let conv_args = match (opertype.clone(), argstype.clone()) {
		    (Arrow(from, _), a) => match a.clone().into_type(&*from) {
			Ok(s) => Ok(s),
			Err(()) => Err(InvalidArgList {
			    arglist: (**l).clone(),
			    expected: *from,
			    found: a,
			})
		    },
		    (a, _) => {
			Err(InvalidOperator {
			    operator: *op.clone(),
			    expected: arr(vt("_"), vt("_")),
			    found:    a,
			})
		    }
		}?;

                let opertype = if opertype.is_concrete().is_ok() {
                    opertype
                } else {
                    opertype.infer_generics(&conv_args)?
                };

                match (opertype, argstype) {
                    (Arrow(a, b), c) => {
                        if c.aka(&*a) {
			    Ok(*b)
			} else {
                            Err(InvalidArgList {
                                    arglist:  (**l).clone(),
                                    expected: *a,
                                    found:    c,
                            })
                        }
                    },
                    (t, _) => Err(InvalidOperator {
                        operator: *op.clone(),
                        expected: arr(
                            VarType(String::from("_")),
                            VarType(String::from("_"))
                        ),
                        found:    t,
                    }),
                }
            },

        }

    }

    pub fn infer_list_type(&self, ctx: HashMap<String, Type>) -> Result<Type, TypeError> {
        let mut res = vec![];
        for exp in self.clone().parts() {
            res.push(exp.infer_type(ctx.clone())?);
        }
        Ok(List(res))
    }

}


impl Type {

    /// Infers the type of generic 'VarType's using type of arguments.
    ///
    /// In case there are generic variables that can't be inferred,
    /// returns an TypeError::UnboundVariable.
    fn infer_generics(&self, argtype: &Type) -> Result<Type, TypeError> {
        match self {
            Arrow(from, to) => {

                let generics = match (*from).infer_generics_ctx(argtype, Vec::new()) {
                    Ok(x) => Ok(x),
                    Err(None) => Err(ArgumentsDontMatchGeneric {
                        argtype: argtype.clone(),
                        generictype: self.clone(),
                    }),
                    Err(Some(e)) => Err(e),
                }?;
                let mut restype = (**to).clone();
                for (name, ty) in generics {
                    restype = restype.subst(&name, &ty);
                }

                match restype.is_concrete() {
                    Ok(()) => Ok(arr(argtype.clone(), restype)),
                    Err(unbound) => Err(UnboundGeneric(unbound, self.clone())),
                }
            },
            _ => Err(OtherError)
        }
    }


    fn infer_generics_ctx(
        &self,
        argtype: &Type,
        ctx: Vec<(String, Type)>
    ) -> Result<Vec<(String, Type)>, Option<TypeError>> {
        match (self, argtype) {

            (a, b) if a == b => Ok(ctx),

            (Arrow(a1, a2), Arrow(b1, b2)) => {
                let mut r1 = a1.infer_generics_ctx(b1, ctx.clone())?;
                let r2 = a2.infer_generics_ctx(b2, ctx.clone())?;
                r1.extend_from_slice(&r2);
                r1.extend_from_slice(&ctx);
                Ok(r1)
            },

            (List(v1), List(v2)) => {
                let mut res = ctx.clone();
                for (t1, t2) in v1.into_iter().zip(v2.into_iter()) {
                    let newctx = t1.infer_generics_ctx(t2, ctx.clone())?;
                    res.extend_from_slice(&newctx);
                }
                Ok(res)
            },

            (VecOf(t1), VecOf(t2)) => {
                let mut res = ctx.clone();
                let newctx = t1.infer_generics_ctx(t2, ctx.clone())?;
                res.extend_from_slice(&newctx);
                Ok(res)
            },

            (VarType(name), ty) => {
                let mut res = ctx.clone();
                res.push((name.clone(), ty.clone()));
                Ok(res)
            },

            (_a, _b) => Err(None),

        }
    }

}


#[cfg(test)]
mod tests {
    use super::{*, TypeError};

    #[test]
    fn test_failing_infer_generics() {
        assert!(
            arr(Integer, VarType("X".to_string())).infer_generics(&Integer).is_err()
        );
    }

    #[test]
    fn test_infer_generics() {

        // Simple identity function, really all we
        // care about (this language doesn't attempt
        // to have useful generic programming capabilities,
        // but some simple features are required for typing
        // quotes)...
        assert_eq!(
            arr(VarType("T".to_string()), VarType("T".to_string())).infer_generics(&Integer),
            Ok(arr(Integer, Integer))
        );

        // ...but let's make it work for other simple cases too,
        // maybe someone finds use for these.
        assert_eq!(
            arr(
                    List(vec![Integer, VarType("T".to_string())]),
                    List(vec![VarType("T".to_string()), Integer])
                ).infer_generics(
                    &List(vec![Integer, arr(Integer, Integer)])
                ),
            Ok(arr(
                List(vec![Integer, arr(Integer, Integer)]),
                List(vec![arr(Integer, Integer), Integer])
            ))
        );

        assert_eq!(
            arr(VarType("T".to_string()), Integer).infer_generics(&List(vec![Integer])),
            Ok(arr(List(vec![Integer]), Integer))
        );

        assert_eq!(
            arr(List(vec![VarType("A".to_string()), VarType("B".to_string())]), VarType("B".to_string()))
                .infer_generics(&List(vec![Integer, arr(Integer, Integer)])),
            Ok(arr(
                    List(vec![Integer, arr(Integer, Integer)]),
                    arr(Integer, Integer)
                ))
        );

        assert_eq!(
            arr(
                arr(VarType("A".to_string()), VarType("B".to_string())),
                arr(VarType("B".to_string()), VarType("A".to_string()))
            ).infer_generics(&arr(Integer, arr(Integer, Integer))),
            Ok(arr(
                arr(Integer, arr(Integer, Integer)),
                arr(arr(Integer, Integer), Integer)
            ))
        );

    }
}