use crate::r#type::{Type, TypeError, Type::*, TypeError::*, util::*};
use crate::sexp::{SExp, SExp::*, SLeaf::*};
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 melisp::{
    ///     r#type::{*, Type::*, TypeError::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// assert_eq!(Atom(Int(1)).type_check(), Ok(Integer));
    /// ```
    ///
    /// Quotes are given list types:
    /// ```rust
    /// use melisp::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// assert_eq!(
    ///     scons(Quote, scons(1, scons(2, Nil))).type_check(),
    ///     Ok(List(vec![Integer, Integer]))
    /// );
    /// ```
    /// Though so is Nil given too:
    /// ```rust
    /// use melisp::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// assert_eq!(
    ///     Atom(Nil).type_check(),
    ///     Ok(List(vec![]))
    /// );
    /// ```
    ///
    /// Some common operators get arrow types:
    /// ```rust
    /// use melisp::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// assert_eq!(Atom(Add).type_check(), Ok(arr(List(vec![Integer, Integer]), Integer)));
    /// assert_eq!(Atom(Mul).type_check(), Ok(arr(List(vec![Integer, Integer]), Integer)));
    /// 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)));
    /// ```
    ///
    /// 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 melisp::{
    ///     r#type::{*, Type::*, TypeError::*, util::*},
    ///     sexp::{SExp::*, SLeaf::*, util::*},
    /// };
    ///
    /// let err = scons(Mul, scons(1, Nil)).type_check();
    /// match err {
    ///     Err(InvalidArgList { .. }) => (),
    ///     _ => panic!(
    ///         "passing only 1 argument to '*' should result in InvalidArgList error, found '{:?}'",
    ///         err
    ///     ),
    /// };
    ///
    /// 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"
    ///     ),
    /// };
    /// ```
    ///
    /// Also, free variables should result in an error
    /// as their type can't be known by the type checker.
    /// ```rust
    /// use melisp::{
    ///     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> {
        self.infer_type(HashMap::new())
    }


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

        match self {

            Atom(Int(_)) => Ok(Integer),
            Atom(Var(name)) => ctx.get(name)
                                  .ok_or(UndefinedVariable(name.to_string()))
                                  .cloned(),
            Atom(Add) => Ok(arr(List(vec![Integer, Integer]), Integer)), // TODO varlen
            Atom(Mul) => Ok(arr(List(vec![Integer, Integer]), Integer)), // TODO varlen
            Atom(Sub) => Ok(arr(List(vec![Integer, Integer]), Integer)),
            Atom(Div) => Ok(arr(List(vec![Integer, Integer]), Integer)),
            Atom(Nil)   => Ok(List(vec![])),
            Atom(Quote) => Ok(arr(
                               VarType("T".to_string()),
                               VarType("T".to_string())
                           )),

            SCons(op, l) => {
                let opertype = (*op).infer_type(ctx.clone())?;
                let argstype =  (*l).infer_list_type(ctx)?;

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

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

        }

    }

    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 = (*from).infer_generics_ctx(argtype, Vec::new())?;
                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)),
                }
            },
            _ => Err(OtherError)
        }
    }

    fn infer_generics_ctx(
        &self,
        argtype: &Type,
        ctx: Vec<(String, Type)>
    ) -> Result<Vec<(String, Type)>, 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)
            },

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

            (_a, _b) => Err(InvalidArgList {
                arglist: Atom(Var("undefined".to_string())), // TODO: hacky as heck
                expected: self.clone(),
                found: argtype.clone(),
            }),

        }
    }

}


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

    #[test]
    fn test_failing_infer_generics() {
        assert_eq!(
            arr(Integer, VarType("X".to_string())).infer_generics(&Integer),
            Err(TypeError::UnboundGeneric(String::from("X")))
        );
    }

    #[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)
            ))
        );

    }
}