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(List(vec![])), // TODO make it all identity functions

            SCons(op, l) if **op == Atom(Quote) => (*l).infer_list_type(ctx),

            SCons(op, l) => {
                let opertype = (*op).infer_type(ctx.clone())?;
                let argstype =  (*l).infer_list_type(ctx)?;
                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(UndefinedType, UndefinedType),
                        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))
    }

}