use std::collections::HashMap;

use crate::ir::*;

/// Runs the constant folding pass over the entire module, modifying it in place.
pub fn fold_constants(module: &mut Module) {
    for func in &mut module.functions {
        fold_function_constants(func);
    }
}

fn fold_function_constants(func: &mut Function) {
    for block in &mut func.blocks {
        let mut known_constants: HashMap<Register, Operand> = HashMap::new();

        for inst in &mut block.instructions {
            // 1. Substitute any known constants into the current instruction
            substitute_constants_in_inst(inst, &known_constants);

            // 2. Evaluate and rewrite instructions where possible
            match inst {
                Instruction::Alloc { .. } => {}
                Instruction::Copy { dest, src } => {
                    if is_constant(src) {
                        known_constants.insert(*dest, *src);
                    }
                }
                Instruction::Binary {
                    dest,
                    op,
                    src1,
                    src2,
                    ..
                } => {
                    if let Some(folded) = evaluate_binary(*op, src1, src2) {
                        known_constants.insert(*dest, folded);

                        // Rewrite the evaluated Binary instruction into a clean Assign
                        *inst = Instruction::Copy {
                            dest: *dest,
                            src: folded,
                        };
                    }
                }
                Instruction::Unary { dest, op, src, .. } => {
                    if let Some(folded) = evaluate_unary(*op, src) {
                        known_constants.insert(*dest, folded);

                        // Rewrite the evaluated Unary instruction into a clean Assign
                        *inst = Instruction::Copy {
                            dest: *dest,
                            src: folded,
                        };
                    }
                }
                // Memory and control flow boundaries cannot be statically folded here
                Instruction::Load { .. }
                | Instruction::Store { .. }
                | Instruction::Call { .. }
                | Instruction::Phi { .. } => {}
            }
        }

        // 3. Evaluate terminators
        substitute_constants_in_terminator(&mut block.terminator, &known_constants);

        if let Terminator::Branch {
            cond: Operand::Boolean(b),
            then_block,
            else_block,
        } = block.terminator
        {
            block.terminator = Terminator::Jump(if b { then_block } else { else_block });
        }
    }
}

// --- Helper Functions ---

fn is_constant(op: &Operand) -> bool {
    matches!(op, Operand::Integer(_) | Operand::Boolean(_))
}

fn substitute_constants_in_inst(inst: &mut Instruction, constants: &HashMap<Register, Operand>) {
    let replace = |op: &mut Operand| {
        if let Operand::Register(r) = op
            && let Some(c) = constants.get(r)
        {
            *op = *c;
        }
    };

    match inst {
        Instruction::Alloc { .. } => {}
        Instruction::Copy { src, .. } => replace(src),
        Instruction::Load { src, .. } => replace(src),
        Instruction::Store { dest, src, .. } => {
            replace(dest);
            replace(src);
        }
        Instruction::Binary { src1, src2, .. } => {
            replace(src1);
            replace(src2);
        }
        Instruction::Unary { src, .. } => replace(src),
        Instruction::Call { args, .. } => {
            for (_, arg_op) in args {
                replace(arg_op);
            }
        }
        Instruction::Phi { sources, .. } => {
            for (op, _) in sources.iter_mut() {
                replace(op);
            }
        }
    }
}

fn substitute_constants_in_terminator(
    term: &mut Terminator,
    constants: &HashMap<Register, Operand>,
) {
    let replace = |op: &mut Operand| {
        if let Operand::Register(r) = op
            && let Some(c) = constants.get(r)
        {
            *op = *c;
        }
    };

    match term {
        Terminator::Branch { cond, .. } => replace(cond),
        Terminator::Return {
            value: Some(val), ..
        } => replace(val),
        _ => {}
    }
}

fn evaluate_binary(op: BinaryOp, src1: &Operand, src2: &Operand) -> Option<Operand> {
    match (src1, src2) {
        (Operand::Integer(a), Operand::Integer(b)) => {
            let a = *a;
            let b = *b;

            match op {
                BinaryOp::Add => Some(Operand::Integer(a.wrapping_add(b))),
                BinaryOp::Sub => Some(Operand::Integer(a.wrapping_sub(b))),
                BinaryOp::SMul => Some(Operand::Integer((a as i64).wrapping_mul(b as i64) as u64)),
                BinaryOp::UMul => Some(Operand::Integer(a.wrapping_mul(b))),
                BinaryOp::UDiv => {
                    if b != 0 {
                        Some(Operand::Integer(a.wrapping_div(b)))
                    } else {
                        None
                    }
                }
                BinaryOp::SDiv => {
                    if b != 0 {
                        Some(Operand::Integer((a as i64).wrapping_div(b as i64) as u64))
                    } else {
                        None
                    }
                }
                BinaryOp::URem => {
                    if b != 0 {
                        Some(Operand::Integer(a.wrapping_rem(b)))
                    } else {
                        None
                    }
                }
                BinaryOp::SRem => {
                    if b != 0 {
                        Some(Operand::Integer((a as i64).wrapping_rem(b as i64) as u64))
                    } else {
                        None
                    }
                }
                BinaryOp::ICmp(cmp) => {
                    let res = match cmp {
                        ICmpOp::Eq => a == b,
                        ICmpOp::Ne => a != b,
                        ICmpOp::Ult => a < b,
                        ICmpOp::Ule => a <= b,
                        ICmpOp::Ugt => a > b,
                        ICmpOp::Uge => a >= b,
                        ICmpOp::Slt => (a as i64) < (b as i64),
                        ICmpOp::Sle => (a as i64) <= (b as i64),
                        ICmpOp::Sgt => (a as i64) > (b as i64),
                        ICmpOp::Sge => (a as i64) >= (b as i64),
                    };
                    Some(Operand::Boolean(res))
                }
            }
        }
        (Operand::Boolean(a), Operand::Boolean(b)) => match op {
            BinaryOp::ICmp(ICmpOp::Eq) => Some(Operand::Boolean(a == b)),
            BinaryOp::ICmp(ICmpOp::Ne) => Some(Operand::Boolean(a != b)),
            _ => None,
        },
        _ => None,
    }
}

fn evaluate_unary(op: UnaryOp, src: &Operand) -> Option<Operand> {
    match (op, src) {
        (UnaryOp::INeg, Operand::Integer(a)) => Some(Operand::Integer(a.wrapping_neg())),
        _ => None,
    }
}