//! Interpreter for Z.

pub mod types;

use types::*;

use crate::ast;

pub fn run(program: ast::Program) -> Result<ast::Value, Error> {
    let env_name = ast::EnvName("global".to_string());
    let mut state = State::new("global".into());

    defs_to_env(program.0, &env_name, &mut state.envs)?;

    let main = ast::Expr::Value(state.envs.get(&env_name)?.get(&"main".into())?.clone());

    let env: Env = state.envs.get(&env_name)?.clone();
    let result = eval_expr(&env, &mut state, &main)?;
    println!("{:?}", result);

    Ok(result)
}

fn eval_statement(
    expr_env: &mut Env,
    state: &mut State,
    statement: &ast::Statement,
) -> Result<ast::Value, Error> {
    match statement {
        ast::Statement::Expr(expr) => eval_expr(expr_env, state, expr),
        ast::Statement::Assign { name, expr } => {
            let result = eval_expr(expr_env, state, expr)?;
            expr_env.insert(name.clone(), result.clone());
            Ok(result)
        }
    }
}

fn eval_expr(expr_env: &Env, state: &mut State, expr: &ast::Expr) -> Result<ast::Value, Error> {
    match expr {
        ast::Expr::If {
            condition,
            then,
            r#else,
        } => match eval_expr(expr_env, state, condition)? {
            ast::Value::Boolean(condition) => {
                if condition {
                    eval_expr(expr_env, state, then)
                } else {
                    eval_expr(expr_env, state, r#else)
                }
            }
            v => Err(Error::NotABoolean(v)),
        },
        ast::Expr::Func(func) => {
            let env_name = state.generate_env(expr_env.clone())?;
            Ok(ast::Value::Closure {
                func: func.clone(),
                env: env_name,
            })
        }
        ast::Expr::Block(statements) => {
            let mut block_env = expr_env.clone();
            if let Some(last) = statements.iter().try_fold(None, |_, statement| {
                let result = eval_statement(&mut block_env, state, statement)?;
                Ok::<Option<ast::Value>, Error>(Some(result))
            })? {
                Ok(last)
            } else {
                Err(Error::LastStatementNotAnExpr)
            }
        }
        ast::Expr::Access { expr, field } => match eval_expr(expr_env, state, expr)? {
            ast::Value::Record(record) => eval_expr(
                expr_env,
                state,
                &ast::Expr::Value(record.get(field).map_err(Error::from)?.clone()),
            ),
            v => Err(Error::NotARecord(v)),
        },
        ast::Expr::Var(var) => {
            let value = expr_env.get(var)?.clone();
            eval_expr(expr_env, state, &ast::Expr::Value(value))
        }
        ast::Expr::MethodCall { expr, method, args } => {
            let r#type = expr.get_type();
            let method = state.method_env.get(r#type, method)?;
            match eval_expr(expr_env, state, &ast::helpers::func(method.clone()))? {
                ast::Value::Closure { func, env } => {
                    let mut closure_env: Env = state.envs.get(&env)?.clone();

                    if func.args.len() != args.len() + 1 {
                        Err(Error::ArgumentsMismatch)?;
                    }

                    let mut args = args.clone();
                    args.push(*expr.clone());

                    for (arg, e) in func.args.into_iter().zip(args.iter()) {
                        let evalled = eval_expr(&closure_env, state, e)?;
                        closure_env.0.insert(arg.name, evalled);
                    }
                    eval_expr(&closure_env, state, &func.body)
                }
                v => Err(Error::NotAFunction(v)),
            }
        }

        ast::Expr::FunCall { func, args } => match eval_expr(expr_env, state, func)? {
            ast::Value::Closure { func, env } => {
                let mut closure_env: Env = state.envs.get(&env)?.clone();

                if func.args.len() != args.len() {
                    Err(Error::ArgumentsMismatch)?;
                }

                for (arg, e) in func.args.into_iter().zip(args.iter()) {
                    let evalled = eval_expr(&closure_env, state, e)?;
                    closure_env.0.insert(arg.name, evalled);
                }
                eval_expr(&closure_env, state, &func.body)
            }
            v => Err(Error::NotAFunction(v)),
        },
        ast::Expr::Value(v) => match v {
            ast::Value::Deferred { expr, env } => {
                let closure_env: Env = state.envs.get(env)?.clone();
                eval_expr(&closure_env, state, expr)
            }
            _ => Ok(v.clone()),
        },
    }
}

fn defs_to_env(
    defs: Vec<ast::Definition>,
    env_name: &ast::EnvName,
    envs: &mut Envs,
) -> Result<(), Error> {
    let mut env = Env::new();

    for def in defs {
        let (name, closure) = match def {
            ast::Definition::Function { func, name } => (
                name,
                ast::Value::Closure {
                    func: Box::new(func),
                    env: env_name.clone(),
                },
            ),
            ast::Definition::Expr { expr, name } => (
                name,
                ast::Value::Deferred {
                    expr: Box::new(expr.clone()),
                    env: env_name.clone(),
                },
            ),
        };
        env.insert_nodup(&name, closure)?;
    }

    envs.0.insert(env_name.clone(), env);
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::helpers;

    #[test]
    fn main_0() {
        let program = vec![helpers::define_expr("main", 0.into())].into();
        let result = run(program);
        assert_eq!(result, Ok(0.into()));
    }
    #[test]
    fn var_lookup() {
        let program = vec![
            helpers::define_expr("main", "lit".into()),
            helpers::define_expr("lit", 0.into()),
        ]
        .into();
        let result = run(program);
        assert_eq!(result, Ok(0.into()));
    }
    #[test]
    fn var_assign_and_lookup() {
        let program = vec![helpers::define_expr(
            "main",
            vec![
                helpers::assign("zero", 0.into()),
                helpers::stmt_expr("zero".into()),
            ]
            .into(),
        )]
        .into();
        let result = run(program);
        assert_eq!(result, Ok(0.into()));
    }

    #[test]
    fn field_access() {
        let program = vec![
            helpers::define_expr("main", "record".into()),
            helpers::define_expr(
                "record",
                ast::Expr::from(vec![("my_field", 0.into())]).field("my_field"),
            ),
        ]
        .into();
        let result = run(program);
        assert_eq!(result, Ok(0.into()));
    }

    #[test]
    fn fun_call() {
        let program = vec![
            helpers::define_expr("main", "zero".into()),
            helpers::define_expr(
                "zero",
                helpers::func(ast::Fn {
                    args: vec![],
                    body: 0.into(),
                })
                .call(vec![]),
            ),
        ]
        .into();
        let result = run(program);
        assert_eq!(result, Ok(0.into()));
    }

    #[test]
    fn if_then_else() {
        let program = vec![helpers::define_expr(
            "main",
            ast::Expr::If {
                condition: Box::new(false.into()),
                then: Box::new(0.into()),
                r#else: Box::new(ast::Expr::If {
                    condition: Box::new(true.into()),
                    then: Box::new(1.into()),
                    r#else: Box::new(2.into()),
                }),
            },
        )]
        .into();
        let result = run(program);
        assert_eq!(result, Ok(1.into()));
    }

    #[test]
    fn fun_call_args() {
        let program = vec![
            helpers::define_expr("main", "zero".into()),
            helpers::define_expr(
                "zero",
                helpers::func(ast::Fn {
                    args: vec![
                        ast::Arg {
                            name: "a".into(),
                            r#type: ast::Type::Named("i64".into()),
                        },
                        ast::Arg {
                            name: "b".into(),
                            r#type: ast::Type::Named("i64".into()),
                        },
                    ],
                    body: "b".into(),
                })
                .call(vec![1.into(), 0.into()]),
            ),
        ]
        .into();
        let result = run(program);
        assert_eq!(result, Ok(0.into()));
    }

    // Errors

    #[test]
    fn duplicate_toplevel_defs() {
        let program = vec![
            helpers::define_expr("main", "record".into()),
            helpers::define_expr("main", 0.into()),
        ]
        .into();
        let result = run(program);
        assert_eq!(result, Err(Error::DuplicateNames("main".into())));
    }

    #[test]
    fn field_access_not_a_record() {
        let program = vec![
            helpers::define_expr("main", "record".into()),
            helpers::define_expr("record", ast::Expr::from(0).field("my_field")),
        ]
        .into();
        let result = run(program);
        assert_eq!(result, Err(Error::NotARecord(0.into())));
    }
    #[test]
    fn fun_call_not_a_function() {
        let program = vec![
            helpers::define_expr("main", "zero".into()),
            helpers::define_expr("zero", ast::Expr::from(0).call(vec![1.into()])),
        ]
        .into();
        let result = run(program);
        assert_eq!(result, Err(Error::NotAFunction(0.into())));
    }
}