use std::{collections::{BTreeMap, BTreeSet}, vec};

use crate::{ast::types::{BlockStmt, BreakStmt, CompileUnit, ContinueStmt, Expr, ExprValue, FuncDeclStmt, GlobalDeclStmt, IfElseBranch, IfStmt, ReturnStmt, Statement, VarDeclStmt, WhileStmt}, diagnostic::Diagnositics, ir::{err::IRError, types::{BinaryOp, CmpOp, Function, IRInstr, IRType, MoveRValue, UnaryOp, Variable, VariableOrIntLit, VariableType}}};
use crate::ast::types::BinaryOp as AstBinaryOp;
use crate::ast::types::UnaryOp as AstUnaryOp;
pub struct Generator {
    var_manager: VariableManager,
    function_map: BTreeMap<String, Function>,
    current_func_return_type: Option<IRType>,
    diagnostic: Diagnositics,
    current_exit_label: Vec<Option<(usize, usize)>>, // true exit, false exit
    // About exit label passing:
    // for non-logical parent expr, current_exit_label is None
    // this means we need to create label for logical child expr, and calculate the value
    // for logical parent expr, we can directly use the label for logical child expr, and no need to calculate
    // if child expr isn't logical, we need to do cmp to decide which label to goto
    while_exit_label: Vec<(usize, usize)>, // continue exit, break exit
    func_exit: Option<(usize, Option<Variable>)>, // (label, return_var)
    label_counter: usize
}

impl Generator {
    pub fn new() -> Self {
        let mut function_map = BTreeMap::new();
        function_map.insert("putint".to_string(), Function {
            name: "putint".to_string(),
            parameter_types: vec![IRType::I32],
            return_type: IRType::Void,
        });
        function_map.insert("getint".to_string(), Function {
            name: "getint".to_string(),
            parameter_types: vec![],
            return_type: IRType::I32,
        });
        Self {
            var_manager: VariableManager::new(),
            current_func_return_type: None,
            diagnostic: Diagnositics::new(),
            function_map,
            current_exit_label: vec![],
            while_exit_label: vec![],
            func_exit: None,
            label_counter: 0,
        }
    }
    fn request_label(&mut self) -> usize {
        let label = self.label_counter;
        self.label_counter += 1;
        label
    }
    pub fn emit(&mut self, compile_unit: CompileUnit) -> Vec<IRInstr> {
        self.generate_compile_unit(compile_unit)
    }
    pub fn get_diagnostics(&self) -> &Diagnositics {
        &self.diagnostic
    }
    fn generate_compile_unit(&mut self, compile_unit: CompileUnit) -> Vec<IRInstr> {
        let mut instrs = vec![];
        use GlobalDeclStmt::*;
        for decl in compile_unit.global_decls {
            match decl {
                VarDecl(var_decl) => {
                    instrs.extend(self.generate_var_decl(var_decl, true));
                }
                FuncDecl(func_decl) => {
                    instrs.extend(self.generate_func_decl(func_decl));
                }
            }
        }
        instrs
    }

    fn generate_var_decl(&mut self, var_decl: VarDeclStmt, is_global: bool) -> Vec<IRInstr> {
        let mut instrs = vec![];
        let var_type = if is_global { VariableType::Global } else { VariableType::Local };
        for value in var_decl.values {
            match self.var_manager.declare_variable(&value.name, var_type, var_decl.data_type.into()) {
                Ok(var) => {
                    if is_global { instrs.push(IRInstr::Declare(var)); }
                }
                Err(e) => {
                    self.diagnostic.add_from_ir_error(e, value.name_span);
                }
            }
        }

        instrs
    }

    fn generate_func_decl(&mut self, func_decl: FuncDeclStmt) -> Vec<IRInstr> {
        if self.function_map.contains_key(&func_decl.name) {
            self.diagnostic.add_from_ir_error(IRError::FunctionHasBeenDefined(func_decl.name.clone()), func_decl.name_span);
            return vec![];
        }
        self.current_func_return_type = Some(func_decl.return_type.into());
        self.var_manager.enter_scope();
        let parameters: Vec<Variable> =  match func_decl.params.iter().map(|param| {
            match self.var_manager.declare_variable(&param.name, VariableType::Local, param.param_type.into()) {
                Ok(var) => Ok(var),
                Err(e) => {
                    self.diagnostic.add_from_ir_error(e, param.name_span);
                    Err(())
                }
            }
        }).collect() {
            Ok(p) => p,
            Err(()) => return vec![],
        };
        let temp_parameters = parameters.iter().map(|param| self.var_manager.declare_param_temp(param.data_type)).collect::<Vec<_>>();
        let mut body_instrs = vec![];
        self.func_exit = Some((self.request_label(),  {
            let ret_type = func_decl.return_type.into();
            if ret_type != IRType::Void {
                Some(self.var_manager.declare_unamed_local(ret_type))
            } else {
                None
            }
        }));
        let block_instrs = self.generate_block_stmt(func_decl.body);
        for var in self.var_manager.get_cur_func_variables() {
            if matches!(var.var_type, VariableType::ParamTemp) {
                continue;
            }
            body_instrs.push(IRInstr::Declare(var));
        }
        body_instrs.push(IRInstr::Entry);
        parameters.iter().zip(temp_parameters.iter()).for_each(|(param, temp_param)| {
            body_instrs.push(IRInstr::Move(*temp_param, MoveRValue::Var(*param)));
        });
        body_instrs.extend(block_instrs);
        let func_exit = self.func_exit.take().unwrap();
        body_instrs.push(IRInstr::Label(func_exit.0));
        body_instrs.push(IRInstr::Exit(func_exit.1));
        self.var_manager.exit_scope();
        self.current_func_return_type = None;
        self.var_manager.clear_local_counter();
        let func  = Function {
            name: func_decl.name,
            parameter_types: temp_parameters.iter().map(|v| v.data_type).collect(),
            return_type: func_decl.return_type.into(),
        };
        self.function_map.insert(func.name.clone(), func.clone());
        vec![IRInstr::DefineFunc(func, parameters, body_instrs)]
    }
    fn generate_block_stmt(&mut self, block_stmt: BlockStmt) -> Vec<IRInstr> {
        let mut instrs = vec![];
        for stmt in block_stmt.statements {
            instrs.extend(self.generate_statement(stmt));
        }
        instrs
    }

    fn generate_statement(&mut self, stmt: Statement) -> Vec<IRInstr> {
        use Statement::*;
        let instrs = match stmt {
            Return(return_stmt) => self.generate_return_stmt(return_stmt),
            If(if_stmt) => self.generate_if_stmt(if_stmt),
            While(while_stmt) => self.generate_while_stmt(while_stmt),
            Break(break_stmt) => self.generate_break_stmt(break_stmt),
            Continue(continue_stmt) => self.generate_continue_stmt(continue_stmt),
            Block(block_stmt) => {
                self.var_manager.enter_scope();
                let block_instrs = self.generate_block_stmt(block_stmt);
                self.var_manager.exit_scope();
                block_instrs
            },
            Expr(expr) => {
                self.current_exit_label.push(None);
                let (instrs, _) = match self.generate_expr(expr) {
                    Some(res) => res,
                    None => {
                        self.current_exit_label.pop();
                        return vec![];
                    }
                };
                self.current_exit_label.pop();
                instrs
            },
            VarDecl(var_decl) => self.generate_var_decl(var_decl, false),
        };
        instrs
    }

    fn generate_return_stmt(&mut self, return_stmt: ReturnStmt) -> Vec<IRInstr> {
        let mut instrs = vec![];
        let func_exit = self.func_exit.unwrap();
        match return_stmt.value {
            Some(expr) => {
                if func_exit.1.is_none() {
                    // shouldn't return value but return stmt has expr;
                    self.diagnostic.add_from_ir_error(IRError::ReturnExpressionOnVoidFunction, return_stmt.span);
                    return vec![];
                }
                self.current_exit_label.push(None);
                let (value_instrs, value_var) = match self.generate_expr(expr) {
                    Some(res) => res,
                    None => {
                        self.current_exit_label.pop();
                        return vec![];
                    }
                };
                self.current_exit_label.pop();
                if value_var.is_none() {
                    self.diagnostic.add_from_ir_error(IRError::InvalidOperand(IRType::Void), return_stmt.span);
                    self.current_exit_label.pop();
                    return vec![];
                }
                instrs.extend(value_instrs);
                instrs.push(IRInstr::Move(func_exit.1.unwrap(), MoveRValue::Var(value_var.unwrap())));
                // if value_var.is_none() {
                //     self.diagnostic.add_from_ir_error(IRError::InvalidOperand(IRType::Void), return_stmt.span);
                //     return vec![];
                // }
                // let value_var = value_var.unwrap();
            }
            None => {
                if func_exit.1.is_some() {
                    // should return value but return stmt has no expr;
                    self.diagnostic.add_from_ir_error(IRError::TypeMismatch(self.current_func_return_type.unwrap(), IRType::Void), return_stmt.span);
                }
            },
        }
        instrs.push(IRInstr::Goto(func_exit.0));
        instrs
    }
    fn generate_while_stmt(&mut self, while_stmt: WhileStmt) -> Vec<IRInstr> {
        let mut instrs = vec![];
        let cond_label = self.request_label();
        let body_label = self.request_label();
        let exit_label = self.request_label();
        instrs.push(IRInstr::Label(cond_label));
        self.current_exit_label.push(Some((body_label, exit_label)));
        let while_cond_span = while_stmt.condition.span;
        let (cond_instrs, cond_var) = match self.generate_expr(while_stmt.condition) {
            Some(res) => res,
            None => {
                self.current_exit_label.pop();
                return vec![];
            }
        };
        self.current_exit_label.pop();
        if cond_var.is_none() {
            self.diagnostic.add_from_ir_error(IRError::InvalidOperand(IRType::Void), while_cond_span);
            return vec![];
        }
        instrs.extend(cond_instrs);
        instrs.push(IRInstr::Label(body_label));
        self.while_exit_label.push((cond_label, exit_label));
        instrs.extend(self.generate_block_stmt(while_stmt.body));
        self.while_exit_label.pop();
        instrs.push(IRInstr::Goto(cond_label));
        instrs.push(IRInstr::Label(exit_label));
        instrs
    }
    fn generate_if_stmt(&mut self, if_stmt: IfStmt) -> Vec<IRInstr> {
        let mut instrs = vec![];
        let then_label = self.request_label();
        let exit_label = self.request_label();
        let mut labels = vec![then_label];
        for _ in &if_stmt.ifelse_branch {
            let else_if_label = self.request_label();
            let else_if_body_label = self.request_label();

            labels.push(else_if_label);
            labels.push(else_if_body_label);
        }
        if if_stmt.else_branch.is_some() {
            let else_label = self.request_label();
            labels.push(else_label);
        }
        labels.push(exit_label);
        // now generate if expr, true exit to labels[0], false exit to labels[1]
        self.current_exit_label.push(Some((labels[0], labels[1])));
        let cond_span = if_stmt.condition.span;
        let (cond_instrs, cond_var) = match self.generate_expr(if_stmt.condition) {
            Some(res) => res,
            None => {
                self.current_exit_label.pop();
                return vec![];
            }
        };
        self.current_exit_label.pop();
        if cond_var.is_none() {
            self.diagnostic.add_from_ir_error(IRError::InvalidOperand(IRType::Void), cond_span);
            return vec![];
        }
        instrs.extend(cond_instrs);
        instrs.push(IRInstr::Label(labels[0]));
        instrs.extend(self.generate_block_stmt(if_stmt.then_branch));
        instrs.push(IRInstr::Goto(exit_label));
        for (i, else_if_branch) in if_stmt.ifelse_branch.into_iter().enumerate() {
            let IfElseBranch { condition: else_if_cond, then_branch: else_if_block } = else_if_branch;
            instrs.push(IRInstr::Label(labels[i * 2 + 1]));
            self.current_exit_label.push(Some((labels[i * 2 + 2], labels[i * 2 + 3])));
            let else_if_cond_span = else_if_cond.span;
            let (else_if_cond_instrs, else_if_cond_var) = match self.generate_expr(else_if_cond) {
                Some(res) => res,
                None => {
                    self.current_exit_label.pop();
                    continue;
                }
            };
            self.current_exit_label.pop();
            if else_if_cond_var.is_none() {
                self.diagnostic.add_from_ir_error(IRError::InvalidOperand(IRType::Void), else_if_cond_span);
                return vec![];
            }
            instrs.extend(else_if_cond_instrs);
            instrs.push(IRInstr::Label(labels[i * 2 + 2]));
            instrs.extend(self.generate_block_stmt(else_if_block));
            instrs.push(IRInstr::Goto(exit_label));
        }
        if let Some(else_block) = if_stmt.else_branch {
            instrs.push(IRInstr::Label(labels[labels.len() - 2]));
            instrs.extend(self.generate_block_stmt(else_block));
            instrs.push(IRInstr::Goto(exit_label));
        }
        instrs.push(IRInstr::Label(exit_label));
        instrs
    }
    fn generate_continue_stmt(&mut self, stmt: ContinueStmt) -> Vec<IRInstr> {
        if let Some((continue_label, _)) = self.while_exit_label.last() {
            vec![IRInstr::Goto(*continue_label)]
        } else {
            self.diagnostic.add_from_ir_error(IRError::ContinueOutsideLoop, stmt.span);
            vec![]
        }
    }
    fn generate_break_stmt(&mut self, stmt: BreakStmt) -> Vec<IRInstr> {
        if let Some((_, break_label)) = self.while_exit_label.last() {
            vec![IRInstr::Goto(*break_label)]
        } else {
            self.diagnostic.add_from_ir_error(IRError::BreakOutsideLoop, stmt.span);
            vec![]
        }
    }
    fn generate_expr(&mut self, expr: Expr) -> Option<(Vec<IRInstr>, Option<Variable>)> {
        // there may be some expr that doesn't produce value, like void func call
        let (mut instrs, var) = match expr.value {
            ExprValue::IntLit(i) => {
                // TODO: convert check
                let var = self.var_manager.declare_temp(IRType::I32);
                (vec![IRInstr::Move(var, MoveRValue::ConstInt(i as i32))], Some(var))
            },
            ExprValue::Var(name) => {
                if let Some(var) = self.var_manager.get_variable(&name) {
                    (vec![], Some(var))
                } else {
                    self.diagnostic.add_from_ir_error(IRError::VariableNotFound(name.clone()), expr.span);
                    return None;
                }
            },
            ExprValue::Assign { lvalue, rvalue } => {
                // TODO: only support simple variable assignment now, need to support more complex lvalue in the future
                if let ExprValue::Var(name) = lvalue.value {
                    let var = match self.var_manager.get_variable(&name) {
                        Some(var) => var,
                        None => {
                            self.diagnostic.add_from_ir_error(IRError::VariableNotFound(name.clone()), lvalue.span);
                            return None;
                        }
                    };
                    self.current_exit_label.push(None);
                    let rvalue_span = rvalue.span;
                    let (mut instrs, rvalue_var) = self.generate_expr(*rvalue)?;
                    self.current_exit_label.pop();
                    // TODO: further check
                    // if var.data_type != rvalue_var.data_type {
                    //     self.diagnostic.add_from_ir_error(IRError::TypeMismatch(var.data_type, rvalue_var.data_type), lvalue.span);
                    //     return (vec![], None);
                    // }
                    if rvalue_var.is_none() {
                        self.diagnostic.add_from_ir_error(IRError::InvalidOperand(IRType::Void), rvalue_span);
                        return None;
                    }
                    instrs.push(IRInstr::Move(var, MoveRValue::Var(rvalue_var.unwrap())));
                    let temp_var = self.var_manager.declare_temp(var.data_type);
                    instrs.push(IRInstr::Move(temp_var, MoveRValue::Var(var)));
                    (instrs, Some(temp_var))
                } else {
                    self.diagnostic.add_from_ir_error(IRError::InvalidAssignmentTarget, lvalue.span);
                    return None;
                }
            },
            ExprValue::UnaryOp { op, operand } => {
                let mut parent_is_logical = true;
                let exit_passdown = if matches!(op, AstUnaryOp::Not) {
                    Some(self.current_exit_label.last().cloned().flatten().map_or_else(|| {
                        parent_is_logical = false;
                        (self.request_label(), self.request_label())
                    }, |(true_exit, false_exit)| (false_exit, true_exit)))
                } else {
                    None
                };
                // NOT:
                //   $x = cmp eq left_var, 0
                //   bc, $x, true_exit, false_exit
                //   false_exit: #only when parent isn't logical
                //   true_exit:
                //   $dest = $x

                // +--------+-------------------------+-------------------------+----------------------------+
                // | parent | (true_exit, false_exit) | (true_exit, false_exit) | None                       |
                // | self   | not                     | add/sub                 | not(true_exit, false_exit) |
                // | child  | (false_exit, true_exit) | None                    | (true_exit, false_exit)    |
                // +--------+-------------------------+-------------------------+----------------------------+


                self.current_exit_label.push(exit_passdown);
                let operand_span = operand.span;
                let (mut instrs, operand_var) = self.generate_expr(*operand)?;
                self.current_exit_label.pop();
                if operand_var.is_none() {
                    self.diagnostic.add_from_ir_error(IRError::InvalidOperand(IRType::Void), operand_span);
                    return None;
                }
                let operand_var = operand_var.unwrap();
                let dest_var = match op {
                    AstUnaryOp::Add => {
                        let dest_var = self.var_manager.declare_temp(operand_var.data_type);
                        instrs.push(IRInstr::Move(dest_var, MoveRValue::Var(operand_var)));
                        dest_var
                    },
                    AstUnaryOp::Sub => {
                        let dest_var = self.var_manager.declare_temp(operand_var.data_type);
                        instrs.push(IRInstr::Unary(dest_var, UnaryOp::Neg, operand_var));
                        dest_var
                    },
                    AstUnaryOp::Not => {
                        let dest_var = self.var_manager.declare_unamed_local(operand_var.data_type);
                        // child will do the cmp
                        if !parent_is_logical {
                            let exit = exit_passdown.unwrap(); // (false_exit, true_exit) (consider `not`)
                            let final_exit = self.request_label();
                            instrs.push(IRInstr::Label(exit.1));
                            instrs.push(IRInstr::Move(dest_var, MoveRValue::ConstInt(1)));
                            instrs.push(IRInstr::Goto(final_exit));
                            instrs.push(IRInstr::Label(exit.0));
                            instrs.push(IRInstr::Move(dest_var, MoveRValue::ConstInt(0)));
                            instrs.push(IRInstr::Goto(final_exit));
                            instrs.push(IRInstr::Label(final_exit));

                        }
                        // return directly since it's logical expr
                        return Some((instrs, Some(dest_var)));
                    }
                };
                (instrs, Some(dest_var))
            },
            ExprValue::BinaryOp { lhs, op, rhs } => {
                let lhs_span = lhs.span;
                let rhs_span = rhs.span;

                // +--------+-------------------------+-------------------------+-------------------------+
                // | parent | (true_exit, false_exit) | (true_exit, false_exit) | (true_exit, false_exit) |
                // | self   | and(next)               | or(next)                | others                  |
                // | lhs    | (next, false_exit)      | (true_exit, next)       | None                    |
                // | rhs    | (true_exit, false_exit) | (next, false_exit)      | None                    |
                // +--------+-------------------------+-------------------------+-------------------------+


                // +--------+----------------------------------+-------------------------------+--------+
                // | parent | None                             | None                          | None   |
                // | self   | and(true_exit, next, false_exit) | or(true_exit,next,false_exit) | others |
                // | lhs    | (next, false_exit)               | (true_exit, next)             | None   |
                // | rhs    | (true_exit, false_exit)          | (next, false_exit)            | None   |
                // +--------+----------------------------------+-------------------------------+--------+


                let mut parent_exit = None;
                let exit = if op.is_logical() {
                    Some(self.current_exit_label.last().cloned().flatten().map_or(
                        (self.request_label(), self.request_label(), self.request_label()),
                        |(true_exit, false_exit)| {
                            parent_exit = Some((true_exit, false_exit));
                            (true_exit, self.request_label(), false_exit)
                        }
                    ))
                } else {
                    None
                }; // (true_exit, next_label, false_exit)
                let lhs_exit_passdown = exit.map(|(true_exit, next_label, false_exit)| match op {
                    AstBinaryOp::And => (next_label, false_exit),
                    AstBinaryOp::Or => (true_exit, next_label),
                    _ => unreachable!(),
                });
                let rhs_exit_passdown = exit.map(|(true_exit, _, false_exit)| match op {
                    AstBinaryOp::And => (true_exit, false_exit),
                    AstBinaryOp::Or => (true_exit, false_exit),
                    _ => unreachable!(),
                });
                self.current_exit_label.push(lhs_exit_passdown);
                let (mut instrs, left_var) = self.generate_expr(*lhs)?;
                self.current_exit_label.pop();
                if left_var.is_none() {
                    self.diagnostic.add_from_ir_error(IRError::InvalidOperand(IRType::Void), lhs_span);
                    return None;
                }
                let mut left_var = left_var.unwrap();
                self.current_exit_label.push(rhs_exit_passdown);
                let (right_instrs, right_var) = self.generate_expr(*rhs)?;
                self.current_exit_label.pop();
                if right_var.is_none() {
                    self.diagnostic.add_from_ir_error(IRError::InvalidOperand(IRType::Void), rhs_span);
                    return None;
                }
                let mut right_var = right_var.unwrap();
                // check implicit convert
                let convert_to;
                if op.is_logical() {
                    // we dont really care since we use cmp
                    convert_to = left_var.data_type;
                } else {
                    if let Some(ty) = IRType::get_elevate_result(left_var.data_type, right_var.data_type) {
                        convert_to = ty;
                    } else {
                        self.diagnostic.add_from_ir_error(IRError::IncompatiableOperand(left_var.data_type, right_var.data_type), lhs_span);
                        self.diagnostic.add_from_ir_error(IRError::IncompatiableOperand(left_var.data_type, right_var.data_type), rhs_span);
                        return None;
                    } 
                }
                // do implicit convert if needed
                // TODO: further check
                if !op.is_logical() && convert_to != left_var.data_type {
                    let temp_var = self.var_manager.declare_temp(convert_to);
                    instrs.push(IRInstr::Move(temp_var, MoveRValue::Var(left_var)));
                    left_var = temp_var;
                }
                let result_type;
                match op {
                    AstBinaryOp::Add | AstBinaryOp::Sub | AstBinaryOp::Mul | AstBinaryOp::Div | AstBinaryOp::Mod => {
                        result_type = left_var.data_type;
                    },
                    AstBinaryOp::Equal | AstBinaryOp::NotEqual | AstBinaryOp::Less | AstBinaryOp::Greater | AstBinaryOp::LessEqual | AstBinaryOp::GreaterEqual 
                    | AstBinaryOp::And | AstBinaryOp::Or => {
                        result_type = IRType::I1;
                    }
                }
                let dest_var = self.var_manager.declare_temp(result_type);
                match op {
                    AstBinaryOp::And | AstBinaryOp::Or => {
                        instrs.push(IRInstr::Label(exit.unwrap().1));
                    },
                    _ => {}
                }
                // for logical operator, considering short-circuit
                // AND: 
                //   $x = cmp ne left_var, 0
                //   bc, $x, next, false_exit
                //   next:
                //   $x = cmp eq right_var, 0
                //   bc, $x, true_exit, false_exit
                //   false_exit: #only when parent isn't logical
                //   true_exit:
                // OR:
                //   $x = cmp ne left_var, 0
                //   bc, $x, true_exit, next
                //   next:
                //   $x = cmp ne right_var, 0
                //   bc, $x, true_exit, false_exit
                //   false_exit: #only when parent isn't logical
                //   true_exit:
                instrs.extend(right_instrs);
                if !op.is_logical() && convert_to != right_var.data_type {
                    let temp_var = self.var_manager.declare_temp(convert_to);
                    instrs.push(IRInstr::Move(temp_var, MoveRValue::Var(right_var)));
                    right_var = temp_var;
                }
                if !op.is_logical() && let Some((true_exit, false_exit)) = parent_exit {
                    let final_exit = self.request_label();
                    instrs.push(IRInstr::Label(true_exit));
                    instrs.push(IRInstr::Move(dest_var, MoveRValue::ConstInt(1)));
                    instrs.push(IRInstr::Goto(final_exit));
                    instrs.push(IRInstr::Label(false_exit));
                    instrs.push(IRInstr::Move(dest_var, MoveRValue::ConstInt(0)));
                    instrs.push(IRInstr::Label(final_exit));


                }
                if op.is_logical() {
                    return Some((instrs, Some(dest_var)));
                } else {
                    if op.is_cmp() {
                        instrs.push(IRInstr::Cmp(dest_var, VariableOrIntLit::Var(left_var), op.into(), VariableOrIntLit::Var(right_var)));
                    } else {
                        instrs.push(IRInstr::Binary(dest_var, left_var, op.into(), right_var));
                    }
                }
                // if !op.is_logical() {
                //     if let Some((true_exit, false_exit))  = self.current_exit_label.last().cloned().flatten() {
                //         // parent is logical
                //         instrs.push(IRInstr::Cmp(dest_var, VariableOrIntLit::Var(dest_var), CmpOp::Ne, VariableOrIntLit::IntLit(0)));
                //         instrs.push(IRInstr::CondGoto(dest_var, true_exit, false_exit));
                //     } else {
                //         // parent isn't logical
                //         if op.is_cmp() {
                //             instrs.push(IRInstr::Cmp(dest_var, VariableOrIntLit::Var(left_var), op.into(), VariableOrIntLit::Var(right_var)));
                //         } else {
                //             instrs.push(IRInstr::Binary(dest_var, left_var, op.into(), right_var));
                //         }
                //     }
                // }
                (instrs, Some(dest_var))
            },
            ExprValue::FuncCall(func_name, args) => {
                let mut instrs = vec![];
                let mut arg_vars = vec![];
                let func_def = if let Some(func_def) = self.function_map.get(&func_name) {
                    func_def
                } else {
                    self.diagnostic.add_from_ir_error(IRError::FunctionNotFound(func_name.clone()), expr.span);
                    return None;
                }.clone();

                if args.len() < func_def.parameter_types.len() {
                    self.diagnostic.add_from_ir_error(IRError::TooFewArguments(func_def.parameter_types.len(), args.len()), expr.span);
                    return None;
                }
                if args.len() > func_def.parameter_types.len() {
                    self.diagnostic.add_from_ir_error(IRError::TooManyArguments(func_def.parameter_types.len(), args.len()), expr.span);
                    return None;
                }
                let mut has_error = false;
                for parameter_type in &func_def.parameter_types {
                    if matches!(parameter_type, IRType::Void) {
                        self.diagnostic.add_from_ir_error(IRError::InvalidParameterType(IRType::Void), expr.span);
                        has_error = true;
                    }
                }
                if has_error {
                    return None;
                }
                for (i, arg) in args.into_iter().enumerate() {
                    let (arg_instrs, arg_var) = self.generate_expr(arg)?;
                    let parameter_type = func_def.parameter_types.get(i).unwrap();
                    if *parameter_type != arg_var.map_or(IRType::Void, |v| v.data_type) {
                        self.diagnostic.add_from_ir_error(IRError::TypeMismatch(*parameter_type, arg_var.map_or(IRType::Void, |v| v.data_type)), expr.span);
                        has_error = true;
                        continue;
                    }
                    instrs.extend(arg_instrs);
                    arg_vars.push(arg_var.unwrap());
                }
                if has_error {
                    return None;
                }
                let ret_variable = if matches!(func_def.return_type, IRType::Void) {
                    None
                } else {
                    Some(self.var_manager.declare_temp(func_def.return_type))
                };
                instrs.push(IRInstr::FuncCall(func_def.clone(), arg_vars, ret_variable));
                (instrs, ret_variable)
            }
        };
        if let Some((true_exit, false_exit)) = self.current_exit_label.last().cloned().flatten() {
            let cmp_var = self.var_manager.declare_temp(IRType::I1);
            instrs.push(IRInstr::Cmp(cmp_var, VariableOrIntLit::Var(var.unwrap()), CmpOp::Ne, VariableOrIntLit::IntLit(0)));
            instrs.push(IRInstr::CondGoto(cmp_var, true_exit, false_exit));
            Some((instrs, var))
        } else {
            Some((instrs, var))
        }
    }
}

struct VariableManager {
    variable_map: BTreeMap<String, Vec<Variable>>,
    scopes: Vec<BTreeSet<String>>,
    global_counter: usize,
    local_counter: usize,
    local_var_type: Vec<Variable>,
}

impl VariableManager {
    pub fn new() -> Self {
        Self {
            variable_map: BTreeMap::new(),
            scopes: vec![BTreeSet::new()],
            global_counter: 0,
            local_counter: 0,
            local_var_type: vec![],
        }
    }
    pub fn enter_scope(&mut self) {
        self.scopes.push(BTreeSet::new());
    }

    pub fn exit_scope(&mut self) {
        let variables = self.scopes.pop().unwrap();
        for var in variables {
            self.variable_map.get_mut(&var).unwrap().pop();
        }
    }

    fn declare_variable(&mut self, name: &str, var_type: VariableType, var_data_type: IRType) -> Result<Variable, IRError> {
        if self.scopes.last().unwrap().contains(name) {
            return Err(IRError::VariableHasBeenDefined(name.to_string()));
        }
        let variable = match var_type {
            VariableType::Global => {
                let var = Variable { index: self.global_counter, var_type, data_type: var_data_type };
                self.global_counter += 1;
                var
            }
            VariableType::Local => {
                let var = Variable { index: self.local_counter, var_type, data_type: var_data_type };
                self.local_counter += 1;
                var
            }
            _ => unreachable!(),
        };
        self.variable_map.entry(name.to_string()).or_default().push(variable);
        self.scopes.last_mut().unwrap().insert(name.to_string());
        if matches!(var_type, VariableType::Local) {
            self.local_var_type.push(variable);
        }
        Ok(variable)
    }
    pub fn declare_gloabal(&mut self, name: &str, var_data_type: IRType) -> Result<Variable, IRError> {
        self.declare_variable(name, VariableType::Global, var_data_type)
    }
    pub fn declare_local(&mut self, name: &str, var_data_type: IRType) -> Result<Variable, IRError> {
        self.declare_variable(name, VariableType::Local, var_data_type)
    }
    pub fn declare_temp(&mut self, var_data_type: IRType) -> Variable {
        let var = Variable { index: self.local_counter, var_type: VariableType::Temp, data_type: var_data_type };
        self.local_counter += 1;
        self.local_var_type.push(var);
        var
    }
    pub fn declare_unamed_local(&mut self, var_data_type: IRType) -> Variable {
        let var = Variable { index: self.local_counter, var_type: VariableType::Local, data_type: var_data_type };
        self.local_counter += 1;
        self.local_var_type.push(var);
        var
    }
    pub fn declare_param_temp(&mut self, var_data_type: IRType) -> Variable {
        let var = Variable { index: self.local_counter, var_type: VariableType::ParamTemp, data_type: var_data_type };
        self.local_counter += 1;
        self.local_var_type.push(var);
        var
    }
    pub fn clear_local_counter(&mut self) {
        self.local_counter = 0;
        self.local_var_type.clear();
    }
    pub fn get_cur_func_variables(&self) -> Vec<Variable> {
        self.local_var_type.iter().cloned().collect()
    }
    pub fn get_variable(&self, name: &str) -> Option<Variable> {
        self.variable_map.get(name).and_then(|vars| vars.last()).cloned()
    }

}

#[cfg(test)]
mod tests {
    use std::io::BufRead;
    use std::path::Path;
    use std::fs::File;
    use crate::ast::graph::AstGraphExt;
    use std::io::Write;
    use crate::frontend::lexer::Lexer;
    use crate::frontend::parser::Parser;
    use crate::utils::case_list::CaseList;
    use crate::utils::num_sequence::NumberSequence;

    pub use super::*;
    fn test_case(case_str: &str) {
        let case_sequence = NumberSequence::from_str(case_str).unwrap();
        let case_list = CaseList::from_dir(&Path::new("./testcases")).unwrap();
        let mut error_case_cnt = 0;
        for case_no in case_sequence {
            let case_path = case_list.get_case_path(case_no).unwrap();
            println!("{}", case_path.display());
            let file = File::open(&case_path).unwrap();
            let mut buf_reader = std::io::BufReader::new(file);
            let mut lexer = Lexer::new();
            let mut full_text = String::new();
            loop {
                let mut line = String::new();
                let bytes_read = buf_reader.read_line(&mut line).unwrap();
                if bytes_read == 0 {
                    break;
                }
                full_text.push_str(&line);
                lexer.parse_next_str(&line);
            }
            let (tokens, diagnostics) = lexer.finish();
            let mut is_error = false;
            if !diagnostics.is_empty() {
                diagnostics.print(&format!("{}", case_path.display()), &full_text);
                is_error = true;
            }
            let mut parser = Parser::new(tokens, diagnostics);
            let compile_unit = parser.parse();
            let case_name = case_list.get_case_name(case_no).unwrap().strip_suffix(".c").unwrap();
            if !parser.diagnostics.is_empty() {
                parser.diagnostics.print(&format!("{}", case_path.display()), &full_text);
                is_error = true;
            }
            let mut generator = Generator::new();
            let ir = generator.emit(compile_unit);
            if !generator.diagnostic.is_empty() {
                generator.diagnostic.print(&format!("{}", case_path.display()), &full_text);
                is_error = true;
            }
            let mut ir_file = File::create(format!("output/{}.ir", case_name)).unwrap();
            for instr in ir {
                writeln!(ir_file, "{}", instr).unwrap();
            }
            if is_error {
                error_case_cnt += 1;
            }
        }
        if error_case_cnt > 0 {
            panic!("Found {} cases with errors", error_case_cnt);
        }

    }
    #[test]
    fn test_expr() {
        test_case("0-3,14-25");
        // test_case("0-3,14-25");
    }
    #[test]
    fn test_if_while() {
        test_case("26-32,34-41,46-51,57");
    }
}