use super::typed_ir;
use super::typed_ir::{Expr, ExprId, ExprKind};
use super::types;
use super::{CompileError, CompileErrorKind, Context};
use std::collections::HashMap;

impl Context {
    pub fn infer_types(&mut self, block: typed_ir::Body) -> typed_ir::Body {
        // TODO: ugly ugly hack
        let try_1 = self
            .infer_body_types(&block)
            .expect("could not infer types");
        let try_2 = self
            .infer_body_types(&try_1)
            .expect("could not infer types");
        self.infer_body_types(&try_2)
            .expect("could not infer types")
    }

    pub fn infer_body_types(
        &mut self,
        body: &typed_ir::Body,
    ) -> Result<typed_ir::Body, CompileError> {
        let mut new_exprs = HashMap::new();
        for (expr_id, expr) in &body.exprs {
            if self.types.is_fully_typed(expr.typ) {
                new_exprs.insert(*expr_id, expr.clone());
                continue;
            }
            match &expr.kind {
                ExprKind::Literal(lit) => {
                    // TODO: don't try to overwrite the type of a literal
                    let infres = self.types.infer_type(expr.typ, lit.typ);
                    new_exprs.insert(*expr_id, expr.clone().with_type(lit.typ));
                }
                ExprKind::Path(_) => {
                    new_exprs.insert(*expr_id, expr.clone());
                }
                ExprKind::Call(call) => {
                    let called_def = self.callables.get(call.called);
                    let param_types = called_def.args.iter().map(|param| param.1);

                    let mut genargs: Vec<_> = called_def.genargs.iter().map(|a| a.1).collect();
                    let inferred_args: Vec<_> = param_types
                        .zip(&call.args)
                        .map(|(param, arg)| {
                            self.types
                                .infer_type(param, body.exprs.get(arg).unwrap().typ)
                        })
                        .collect();

                    if !genargs.is_empty() {
                        // need to infer generic arguments
                        for inf_res in inferred_args {
                            match inf_res {
                                types::InferenceResult::TypeVar(dbi, tvar, typ) => {
                                    assert_eq!(dbi, 0);
                                    // TODO: type check argument instead of just using it
                                    genargs[tvar as usize] = typ;
                                }
                                _ => todo!(),
                            }
                        }
                    }
                    let ret_type = match self.types.infer_type(expr.typ, called_def.ret_type) {
                        types::InferenceResult::TypeVar(dbi, tvar, typ) => {
                            assert_eq!(dbi, 0);
                            genargs[tvar as usize]
                        }
                        types::InferenceResult::First(typ) => typ,
                        x => todo!("{x:?}"),
                    };
                    let new_expr = typed_ir::Expr {
                        kind: typed_ir::ExprKind::Call(typed_ir::Call {
                            genargs,
                            ..call.clone()
                        }),
                        typ: ret_type,
                        ..expr.clone()
                    };
                    new_exprs.insert(*expr_id, new_expr);
                }
                ExprKind::Match(match_) => {
                    // TODO: hacky hacky hacky
                    let res_type = body.exprs.get(&match_.arms.first().unwrap().1).unwrap().typ;
                    let new_expr = expr.clone().with_type(res_type);
                    new_exprs.insert(*expr_id, new_expr);
                }
            }
        }
        /*
        for (expr_id, expr) in &new_exprs {
            if !self.types.is_fully_typed(expr.typ) {
               return Err(CompileError::new(CompileErrorKind::TodoError("fail".to_owned())))
            }
        }
        */
        let new_body = typed_ir::Body {
            exprs: new_exprs,
            ..body.clone()
        };
        Ok(new_body)
    }
}