use super::typed_ir;
use super::types;
use super::Context;

impl Context {
    pub fn infer_types(&mut self, mut block: typed_ir::Body) -> typed_ir::Body {
        let new_root = self.infer_expr_types(&block.expr);
        block.expr = new_root;
        block
    }

    pub fn infer_expr_types(&mut self, expr: &typed_ir::Expr) -> typed_ir::Expr {
        if self.types.is_fully_typed(expr.typ) {
            // there is nothing more to infer
            return expr.clone();
        }
        match &expr.kind {
            typed_ir::ExprKind::Literal(lit) => expr.clone().with_type(lit.typ),
            // we can not see beyond this expression right now
            typed_ir::ExprKind::Path(_) => expr.clone(),
            typed_ir::ExprKind::Call(call) => {
                let args_typed: Vec<_> = call
                    .args
                    .iter()
                    .map(|ex| self.infer_expr_types(ex))
                    .collect();
                let callee_def = self.callables.get(call.called);

                let param_types: Vec<_> = callee_def.args.iter().map(|param| param.1).collect();
                let inferred_args: Vec<_> = param_types
                    .iter()
                    .zip(&args_typed)
                    .map(|(param, arg)| self.types.infer_type(*param, arg.typ))
                    .collect();

                let mut genargs: Vec<_> = callee_def.genargs.iter().map(|a| a.1).collect();

                let mut new_type = callee_def.ret_type;

                if !genargs.is_empty() {
                    // need to infer generic arguments
                    for inf_res in inferred_args {
                        match inf_res {
                            types::InferenceResult::First(_) => todo!(),
                            types::InferenceResult::Second(_) => todo!(),
                            types::InferenceResult::TypeVar(dbi, tvar, typ) => {
                                assert_eq!(dbi, 0);
                                // TODO: type check argument instead of just using it
                                genargs[tvar as usize] = typ;
                            }
                            types::InferenceResult::Incompatible => todo!(),
                            types::InferenceResult::Ambigous => todo!(),
                        }
                    }

                    // TODO: HACKY HACKY HACK
                    new_type = genargs[0];
                }

                let mut new_expr = expr.clone();
                new_expr.typ = new_type;
                new_expr.kind = typed_ir::ExprKind::Call(typed_ir::Call {
                    called: call.called,
                    args: args_typed,
                    genargs,
                });
                new_expr
            }
            typed_ir::ExprKind::Match(match_) => {
                let new_arms: Vec<_> = match_
                    .arms
                    .iter()
                    .map(|(pat, val)| (self.infer_expr_types(pat), self.infer_expr_types(val)))
                    .collect();
                // TODO: hacky hacky hacky
                let res_type = new_arms.first().unwrap().1.typ;
                let new_match = typed_ir::Match {
                    expr: self.infer_expr_types(&match_.expr),
                    arms: new_arms,
                };
                let mut new_expr = expr.clone().with_type(res_type);
                new_expr.kind = typed_ir::ExprKind::Match(Box::new(new_match));
                new_expr
            }
        }
    }
}