use std::collections::BTreeMap;

use crate::{backend::{arm_instr::{ARMInstr, AddInstr, BInstr, BlInstr, CmpInstr, ConditionCode, LoadInstr, MoveInstr, MovtInstr, MovwInstr, MulInstr, PopInstr, PushInstr, RegisterOrImm, RsbInstr, SDivInstr, StoreInstr, SubInstr}, register_allocator::{REG_FP, REG_R0, REG_R1, REG_R2, REG_R3, REG_SP, Register, RegisterAlloc, RegisterAllocator}, types::ARMAsmVar}, ir::types::{Function, IRInstr, MoveRValue, Variable, VariableOrIntLit, VariableType}};
use crate::ir::types::BinaryOp as IRBinaryOp;
use crate::ir::types::CmpOp as IRCmpOp;
use crate::ir::types::UnaryOp as IRUnaryOp;

pub const ARM_STACK_ALIGNMENT: usize = 8;
pub struct Generator {
    instrs: Vec<ARMInstr>,
    var_inited: Vec<ARMAsmVar>,
    var_uninited: Vec<ARMAsmVar>,
    register_allocator: RegisterAllocator,
}

const DEFAULT_VAR_ALIGN: usize = 4;

const ARG_REGS: [Register; 4] = [REG_R0, REG_R1, REG_R2, REG_R3];
const ARM_DATA_IMM_CHUNK: i32 = 255;
const ARM_LOAD_STORE_MAX_OFFSET: i32 = 4095;

fn emit_load_imm32(dest: Register, value: i32, instrs: &mut Vec<ARMInstr>) {
    let bits = value as u32;
    let lower = (bits & 0xffff) as u16;
    let upper = (bits >> 16) as u16;

    instrs.push(MovwInstr::new_imm(dest, lower));
    if upper != 0 {
        instrs.push(MovtInstr::new_imm(dest, upper));
    }
}

fn emit_load_symbol_addr(dest: Register, name: String, instrs: &mut Vec<ARMInstr>) {
    instrs.push(MovwInstr::new_symbol(dest, name.clone()));
    instrs.push(MovtInstr::new_symbol(dest, name));
}

fn emit_sub_imm(dest: Register, left: Register, imm: i32, instrs: &mut Vec<ARMInstr>) {
    if imm == 0 {
        if dest != left {
            instrs.push(MoveInstr::new_uncond(dest, RegisterOrImm::Reg(left)));
        }
        return;
    }
    let mut remaining = imm;
    let first_chunk = remaining.min(ARM_DATA_IMM_CHUNK);
    instrs.push(SubInstr::new(dest, left, RegisterOrImm::Imm(first_chunk)));
    remaining -= first_chunk;
    while remaining > 0 {
        let chunk = remaining.min(ARM_DATA_IMM_CHUNK);
        instrs.push(SubInstr::new(dest, dest, RegisterOrImm::Imm(chunk)));
        remaining -= chunk;
    }
}

fn emit_add_sp_imm(imm: i32, instrs: &mut Vec<ARMInstr>) {
    let mut remaining = imm;
    while remaining > 0 {
        let chunk = remaining.min(ARM_DATA_IMM_CHUNK);
        instrs.push(AddInstr::new_sp(chunk));
        remaining -= chunk;
    }
}

fn emit_sub_sp_imm(imm: i32, instrs: &mut Vec<ARMInstr>) {
    let mut remaining = imm;
    while remaining > 0 {
        let chunk = remaining.min(ARM_DATA_IMM_CHUNK);
        instrs.push(SubInstr::new_sp(chunk));
        remaining -= chunk;
    }
}

fn emit_load_stack(dest: Register, offset: i32, reg_allocator: &mut RegisterAllocator, instrs: &mut Vec<ARMInstr>) {
    if offset <= ARM_LOAD_STORE_MAX_OFFSET {
        instrs.push(LoadInstr::new_stack(dest, offset));
    } else {
        let addr_alloc = reg_allocator.alloc_any().expect("Ran out of registers");
        emit_sub_imm(addr_alloc.reg, REG_FP, offset, instrs);
        instrs.push(LoadInstr::new(dest, addr_alloc.reg, None));
    }
}

fn emit_store_stack(src: Register, offset: i32, reg_allocator: &mut RegisterAllocator, instrs: &mut Vec<ARMInstr>) {
    if offset <= ARM_LOAD_STORE_MAX_OFFSET {
        instrs.push(StoreInstr::new_stack(src, offset));
    } else {
        let addr_alloc = reg_allocator.alloc_any().expect("Ran out of registers");
        emit_sub_imm(addr_alloc.reg, REG_FP, offset, instrs);
        instrs.push(StoreInstr::new(src, addr_alloc.reg, None));
    }
}

fn load_variable(variable: Variable, reg_allocator: &mut RegisterAllocator, var_index_to_stack_offset: &BTreeMap<usize, usize>, instrs: &mut Vec<ARMInstr>) -> RegisterAlloc {
    if variable.data_type.is_array() {
        let var_alloc = reg_allocator.alloc(variable.clone()).expect("Ran out of registers");
        if variable.data_type.is_array_param() {
            match variable.var_type {
                VariableType::ParamTemp(param_index) => {
                    if param_index < ARG_REGS.len() {
                        instrs.push(MoveInstr::new_uncond(var_alloc.reg, RegisterOrImm::Reg(ARG_REGS[param_index])));
                    } else {
                        let stack_arg_offset = 8 + ((param_index - ARG_REGS.len()) * 4);
                        instrs.push(LoadInstr::new(var_alloc.reg, REG_FP, Some(RegisterOrImm::Imm(stack_arg_offset as i32))));
                    }
                }
                _ => {
                    let stack_offset = var_index_to_stack_offset.get(&variable.index).expect("Variable not declared");
                    emit_load_stack(var_alloc.reg, *stack_offset as i32, reg_allocator, instrs);
                }
            }
            return var_alloc;
        }
        match variable.var_type {
            VariableType::Global => {
                emit_load_symbol_addr(var_alloc.reg, format!("global_var_{}", variable.index), instrs);
            }
            _ => {
                let stack_offset = var_index_to_stack_offset.get(&variable.index).expect("Variable not declared");
                emit_sub_imm(var_alloc.reg, REG_FP, *stack_offset as i32, instrs);
            }
        }
        return var_alloc;
    }
    match variable.var_type {
        VariableType::Global => {
            let var_alloc = reg_allocator.alloc(variable.clone()).expect("Ran out of registers");
            let var_reg = var_alloc.reg;
            // if !var_alloc.is_reused {
                let address_alloc = reg_allocator.alloc_any().expect("Ran out of registers");
                emit_load_symbol_addr(address_alloc.reg, format!("global_var_{}", variable.index), instrs);
                instrs.push(LoadInstr::new(var_reg, address_alloc.reg, None));
            // }
            var_alloc
        },
        VariableType::ParamTemp(param_index) => {
            if param_index < ARG_REGS.len() {
                let reg = ARG_REGS[param_index];
                let var_alloc = reg_allocator.alloc_reg(reg).expect("Ran out of registers");
                var_alloc
            } else {
                let var_alloc = reg_allocator.alloc(variable.clone()).expect("Ran out of registers");
                let stack_arg_offset = 8 + ((param_index - ARG_REGS.len()) * 4);
                instrs.push(LoadInstr::new(var_alloc.reg, REG_FP, Some(RegisterOrImm::Imm(stack_arg_offset as i32))));
                var_alloc
            }
        },
        _ => {
            let stack_offset = var_index_to_stack_offset.get(&variable.index).expect("Variable not declared");
            let var_alloc = reg_allocator.alloc(variable.clone()).expect("Ran out of registers");
            // if !var_alloc.is_reused {
                emit_load_stack(var_alloc.reg, *stack_offset as i32, reg_allocator, instrs);
            // }
            var_alloc
        }
    }
}

fn save_variable(variable: Variable, reg: Register, reg_allocator: &mut RegisterAllocator, var_index_to_stack_offset: &BTreeMap<usize, usize>, instrs: &mut Vec<ARMInstr>) {
    match variable.var_type {
        VariableType::Global => {
            let address_alloc = reg_allocator.alloc_any().expect("Ran out of registers");
            emit_load_symbol_addr(address_alloc.reg, format!("global_var_{}", variable.index), instrs);
            instrs.push(StoreInstr::new(reg, address_alloc.reg, None));
        },
        VariableType::ParamTemp(_) => {
            todo!()
        },
        _ => {
            let stack_offset = var_index_to_stack_offset.get(&variable.index).expect("Variable not declared");
            emit_store_stack(reg, *stack_offset as i32, reg_allocator, instrs);
        }
    }
}
impl Generator {
    pub fn new() -> Self {
        Self {
            instrs: Vec::new(),
            var_inited: Vec::new(),
            var_uninited: Vec::new(),
            register_allocator: RegisterAllocator::new(),
        }
    }
    pub fn emit(&mut self, ir_instrs: Vec<IRInstr>) {
        for ir_instr in ir_instrs {
            match ir_instr {
                IRInstr::DefineFunc(func, args,  body) => self.emit_func_def(func, args, body),
                IRInstr::DeclareGlobal(var, init) => self.emit_global_decl(var, init),
                _ => unreachable!(),
            }
        }
    }
    pub fn to_text(&self) -> String {
        let mut text = String::new();
        text.push_str(".arch armv7ve\n.arm\n.fpu vfpv4\n");
        for var in &self.var_uninited {
            text.push_str(&format!(".comm {}, {}, {}\n", var.name, var.size, var.align));
        }
        for var in &self.var_inited {
            text.push_str(&format!(".data\n.align {}\n.global {}\n.type {}, %object\n{}:\n", var.align, var.name, var.name, var.name));
            text.push_str(&format!(".word {}\n", var.init.unwrap_or(0)));
        }
        text.push_str(".text\n");
        for instr in &self.instrs {
            text.push_str(&format!("{}\n", instr));
        }
        text
    }

    fn emit_global_decl(&mut self, var: Variable, init: Option<i32>) {
        let asm_var = ARMAsmVar {
            name: format!("global_var_{}", var.index),
            size: var.data_type.size_in_bytes(),
            align: DEFAULT_VAR_ALIGN,
            init,
        };
        if asm_var.init.is_some() {
            self.var_inited.push(asm_var);
        } else {
            self.var_uninited.push(asm_var);
        }
    }

    fn emit_func_def(&mut self, func: Function, _args: Vec<Variable>, body: Vec<IRInstr>) {
        self.instrs.push(ARMInstr::FunctionHead(func.name.clone(), 4)); // Assuming 4-byte alignment for simplicity
        self.instrs.push(PushInstr::new_push_fp_lr());
        self.instrs.push(MoveInstr::new_sp_to_fp());
        self.emit_func(&func.name, body);
    }

    fn emit_func(&mut self, func_name: &str, instrs: Vec<IRInstr>) {
        let mut encounter_entry = false;
        let mut stack_size_needed = 0;
        let mut var_index_to_stack_offset = BTreeMap::new();
        for ir_instr in instrs {
            match ir_instr {
                IRInstr::Binary(dest, left, op, right) => self.emit_binary(dest, left, op, right, &var_index_to_stack_offset),
                IRInstr::Exit(v) => {
                    if let Some(v) = v {
                        let ret_alloc = self.register_allocator.alloc_reg(REG_R0).expect("Ran out of registers");
                        let ret_reg = ret_alloc.reg;
                        let v_alloc = load_variable(v, &mut self.register_allocator, &var_index_to_stack_offset, &mut self.instrs);
                        self.instrs.push(MoveInstr::new_uncond(ret_reg, RegisterOrImm::Reg(v_alloc.reg)));
                    }
                    self.instrs.push(MoveInstr::new_fp_to_sp());
                    self.instrs.push(PopInstr::new_pop_fp_pc());
                },
                IRInstr::FuncCall(func, args, ret) => self.emit_func_call(func, args, ret, &var_index_to_stack_offset),
                IRInstr::Move(dest, src) => self.emit_move(dest, src, &var_index_to_stack_offset),
                IRInstr::Load(dest, addr) => self.emit_load(dest, addr, &var_index_to_stack_offset),
                IRInstr::Store(addr, value) => self.emit_store(addr, value, &var_index_to_stack_offset),
                IRInstr::Declare(variable) => {
                    assert!(!encounter_entry, "Variable declarations must come before entry instruction");
                    let size = if variable.data_type.is_array_param() { 4 } else { variable.data_type.size_in_bytes() };
                    stack_size_needed = (stack_size_needed + size).next_multiple_of(ARM_STACK_ALIGNMENT);
                    var_index_to_stack_offset.insert(variable.index, stack_size_needed);
                },
                IRInstr::Entry => {
                    assert!(!encounter_entry, "Multiple entry instructions are not allowed");
                    encounter_entry = true;
                    emit_sub_sp_imm(stack_size_needed as i32, &mut self.instrs);
                },
                IRInstr::DefineFunc(_, _, _) | IRInstr::DeclareGlobal(_, _) => unreachable!(),
                IRInstr::Cmp(variable, left, cmp_op, right) => self.emit_cmp(variable, left, cmp_op, right, &var_index_to_stack_offset),
                IRInstr::Unary(variable, unary_op, variable1) => self.emit_unary(variable, unary_op, variable1, &var_index_to_stack_offset),
                IRInstr::Goto(index) => self.instrs.push(BInstr::new(format!("label_{}_{}", func_name, index))),
                IRInstr::CondGoto(variable, true_label_index, false_label_index) => {
                    let variable_alloc = load_variable(variable, &mut self.register_allocator, &var_index_to_stack_offset, &mut self.instrs);
                    let variable_reg = variable_alloc.reg;
                    self.instrs.push(CmpInstr::new(variable_reg, RegisterOrImm::Imm(0)));
                    self.instrs.push(BInstr::new_cond(ConditionCode::Ne, format!("label_{}_{}", func_name, true_label_index)));
                    self.instrs.push(BInstr::new(format!("label_{}_{}", func_name, false_label_index)));
                    
                },
                IRInstr::Label(index) => self.instrs.push(ARMInstr::Label(format!("label_{}_{}", func_name, index))),
            }
        }
    }
    fn emit_unary(&mut self, dest: Variable, unary_op: IRUnaryOp, variable: Variable, var_index_to_stack_offset: &BTreeMap<usize, usize>) {
        let variable_alloc = load_variable(variable, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs);
        let dest_alloc = self.register_allocator.alloc(dest.clone()).expect("Ran out of registers");
        match unary_op {
            IRUnaryOp::Neg => {
                self.instrs.push(RsbInstr::new(dest_alloc.reg, variable_alloc.reg, RegisterOrImm::Imm(0)));
            },
        }
        save_variable(dest, dest_alloc.reg, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs);
    }
    fn emit_cmp(&mut self, variable: Variable, left: VariableOrIntLit, cmp_op: IRCmpOp, right: VariableOrIntLit, var_index_to_stack_offset: &BTreeMap<usize, usize>) {
        let left_alloc = match left {
            VariableOrIntLit::Var(var) => {
                load_variable(var, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs)
            },
            VariableOrIntLit::IntLit(lit) => {
                let alloc = self.register_allocator.alloc_any().expect("Ran out of registers");
                emit_load_imm32(alloc.reg, lit, &mut self.instrs);
                alloc
            },
        };
        // TODO: check left == right
        let right_alloc = match right {
            VariableOrIntLit::Var(var) => {
                load_variable(var, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs)
            },
            VariableOrIntLit::IntLit(lit) => {
                let alloc = self.register_allocator.alloc_any().expect("Ran out of registers");
                emit_load_imm32(alloc.reg, lit, &mut self.instrs);
                alloc
            },
        };
        let variable_alloc = self.register_allocator.alloc(variable.clone()).expect("Ran out of registers");
        let variable_reg = variable_alloc.reg;
        self.instrs.push(CmpInstr::new(left_alloc.reg, RegisterOrImm::Reg(right_alloc.reg)));
        self.instrs.push(MoveInstr::new_uncond(variable_reg, RegisterOrImm::Imm(0)));
        self.instrs.push(MoveInstr::new_cond(cmp_op.into(), variable_reg, RegisterOrImm::Imm(1)));
        save_variable(variable, variable_alloc.reg, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs);
    }
    fn emit_func_call(&mut self, func: Function, args: Vec<Variable>, ret: Option<Variable>, var_index_to_stack_offset: &BTreeMap<usize, usize>) {

        self.instrs.push(PushInstr::new_push_caller_save());
        let stack_arg_count = args.len().saturating_sub(ARG_REGS.len());
        let stack_arg_size = stack_arg_count * 4;
        let caller_save_size = 5 * 4;
        let stack_padding = if (caller_save_size + stack_arg_size) % ARM_STACK_ALIGNMENT == 0 { 0 } else { 4 };
        let extra_stack_size = stack_arg_size + stack_padding;
        if extra_stack_size > 0 {
            emit_sub_sp_imm(extra_stack_size as i32, &mut self.instrs);
        }
        let mut arg_reg_allocs = Vec::new();
        for (i, arg) in args.into_iter().enumerate() {
            let arg_alloc = load_variable(arg, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs);
            if i < ARG_REGS.len() {
                arg_reg_allocs.push(self.register_allocator.alloc_reg(ARG_REGS[i]).expect("Ran out of registers"));
                self.instrs.push(MoveInstr::new_uncond(ARG_REGS[i], RegisterOrImm::Reg(arg_alloc.reg)));
            } else {
                let stack_offset = ((i - ARG_REGS.len()) * 4) as i32;
                self.instrs.push(StoreInstr::new(arg_alloc.reg, REG_SP, Some(RegisterOrImm::Imm(stack_offset))));
            }
        }
        self.instrs.push(BlInstr::new(func.name.clone()));
        if let Some(ret) = ret {
            save_variable(ret, REG_R0, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs);
        }

        if extra_stack_size > 0 {
            emit_add_sp_imm(extra_stack_size as i32, &mut self.instrs);
        }

        self.instrs.push(PopInstr::new_pop_caller_save());
    }
    fn emit_move(&mut self, dest: Variable, src: MoveRValue, var_index_to_stack_offset: &BTreeMap<usize, usize>) {
        let dest_alloc = self.register_allocator.alloc(dest.clone()).expect("Ran out of registers");
        let dest_register = dest_alloc.reg;
        match src {
            MoveRValue::Var(variable) => {
                let src_alloc = load_variable(variable, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs);
                self.instrs.push(MoveInstr::new_uncond(dest_register, RegisterOrImm::Reg(src_alloc.reg)));
            },
            MoveRValue::ConstInt(literal_int) => emit_load_imm32(dest_register, literal_int, &mut self.instrs),
        };
        save_variable(dest, dest_alloc.reg, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs);
    }

    fn emit_load(&mut self, dest: Variable, addr: Variable, var_index_to_stack_offset: &BTreeMap<usize, usize>) {
        let addr_alloc = load_variable(addr, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs);
        let dest_alloc = self.register_allocator.alloc(dest.clone()).expect("Ran out of registers");
        self.instrs.push(LoadInstr::new(dest_alloc.reg, addr_alloc.reg, None));
        save_variable(dest, dest_alloc.reg, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs);
    }

    fn emit_store(&mut self, addr: Variable, value: Variable, var_index_to_stack_offset: &BTreeMap<usize, usize>) {
        let addr_alloc = load_variable(addr, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs);
        let value_alloc = load_variable(value, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs);
        self.instrs.push(StoreInstr::new(value_alloc.reg, addr_alloc.reg, None));
    }

    fn emit_binary(&mut self, dest: Variable, left: Variable, op: IRBinaryOp, right: Variable, var_index_to_stack_offset: &BTreeMap<usize, usize>) {
        let dest_alloc = self.register_allocator.alloc(dest.clone()).expect("Ran out of registers");
        let dest_reg = dest_alloc.reg;
        // should consider left == right
        let left_alloc = load_variable(left.clone(), &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs);
        let (_right_alloc, right_reg) = if left != right {
            let right_alloc = load_variable(right, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs);
            let right_reg = right_alloc.reg;
            (Some(right_alloc), right_reg)
        } else {
            (None, left_alloc.reg)
        };
        match op {
            IRBinaryOp::Add => {
                self.instrs.push(AddInstr::new(dest_reg, left_alloc.reg, RegisterOrImm::Reg(right_reg)));
            },
            IRBinaryOp::Sub => {
                self.instrs.push(SubInstr::new(dest_reg, left_alloc.reg, RegisterOrImm::Reg(right_reg)));
            },
            IRBinaryOp::Mul => {
                self.instrs.push(MulInstr::new(dest_reg, left_alloc.reg, RegisterOrImm::Reg(right_reg)));
            },
            IRBinaryOp::Div => {
                self.instrs.push(SDivInstr::new(dest_reg, left_alloc.reg, RegisterOrImm::Reg(right_reg)));
            },
            IRBinaryOp::Mod => {
                let temp_alloc = self.register_allocator.alloc_any().expect("Ran out of registers");
                let temp_reg = temp_alloc.reg;
                self.instrs.push(SDivInstr::new(temp_reg, left_alloc.reg, RegisterOrImm::Reg(right_reg)));
                self.instrs.push(MulInstr::new(temp_reg, temp_reg, RegisterOrImm::Reg(right_reg)));
                self.instrs.push(SubInstr::new(dest_reg, left_alloc.reg, RegisterOrImm::Reg(temp_reg)));
            },
        }
        save_variable(dest, dest_alloc.reg, &mut self.register_allocator, var_index_to_stack_offset, &mut self.instrs);
    }

}
