use crate::ParserOptions;
use alloc::vec::Vec;
use tinywasm_types::{Instruction, WasmFunctionData};

pub(crate) fn optimize_instructions(
    mut instructions: Vec<Instruction>,
    function_data: &mut WasmFunctionData,
    self_func_addr: u32,
    options: &ParserOptions,
) -> Vec<Instruction> {
    rewrite(&mut instructions, self_func_addr);
    if options.dce {
        dce(&mut instructions, function_data);
    }
    instructions
}

fn rewrite(instructions: &mut [Instruction], self_func_addr: u32) {
    for read in 0..instructions.len() {
        match instructions[read] {
            Instruction::LocalCopy32(a, b) if a == b => instructions[read] = Instruction::Nop,
            Instruction::LocalCopy64(a, b) if a == b => instructions[read] = Instruction::Nop,
            Instruction::LocalCopy128(a, b) if a == b => instructions[read] = Instruction::Nop,
            Instruction::Call(addr) if addr == self_func_addr => instructions[read] = Instruction::CallSelf,
            Instruction::ReturnCall(addr) if addr == self_func_addr => instructions[read] = Instruction::ReturnCallSelf,
            Instruction::I32Add => {
                if read > 1
                    && let (Instruction::LocalGet32(a), Instruction::LocalGet32(b)) =
                        (instructions[read - 2], instructions[read - 1])
                {
                    instructions[read - 2] = Instruction::Nop;
                    instructions[read - 1] = Instruction::Nop;
                    instructions[read] = Instruction::I32AddLocals(a, b);
                }

                if read > 0 {
                    match instructions[read - 1] {
                        Instruction::I32Const(c) if read > 1 => {
                            if let Instruction::LocalGet32(local) = instructions[read - 2] {
                                instructions[read - 2] = Instruction::Nop;
                                instructions[read - 1] = Instruction::LocalGet32(local);
                                instructions[read] = Instruction::I32AddConst(c);
                            } else {
                                instructions[read - 1] = Instruction::Nop;
                                instructions[read] = Instruction::I32AddConst(c);
                            }
                        }
                        Instruction::I32Const(c) => {
                            instructions[read - 1] = Instruction::Nop;
                            instructions[read] = Instruction::I32AddConst(c);
                        }
                        _ => {}
                    }
                }
            }
            Instruction::I64Add => {
                if read > 1
                    && let (Instruction::LocalGet64(a), Instruction::LocalGet64(b)) =
                        (instructions[read - 2], instructions[read - 1])
                {
                    instructions[read - 2] = Instruction::Nop;
                    instructions[read - 1] = Instruction::Nop;
                    instructions[read] = Instruction::I64AddLocals(a, b);
                }

                if read > 0 {
                    match instructions[read - 1] {
                        Instruction::I64Const(c) if read > 1 => {
                            if let Instruction::LocalGet64(local) = instructions[read - 2] {
                                instructions[read - 2] = Instruction::Nop;
                                instructions[read - 1] = Instruction::LocalGet64(local);
                                instructions[read] = Instruction::I64AddConst(c);
                            } else {
                                instructions[read - 1] = Instruction::Nop;
                                instructions[read] = Instruction::I64AddConst(c);
                            }
                        }
                        Instruction::I64Const(c) => {
                            instructions[read - 1] = Instruction::Nop;
                            instructions[read] = Instruction::I64AddConst(c);
                        }
                        _ => {}
                    }
                }
            }
            Instruction::I64Rotl => {
                if read > 1
                    && let (Instruction::I64Xor, Instruction::I64Const(c)) =
                        (instructions[read - 2], instructions[read - 1])
                {
                    instructions[read - 2] = Instruction::Nop;
                    instructions[read - 1] = Instruction::Nop;
                    instructions[read] = Instruction::I64XorRotlConst(c);
                }
            }
            Instruction::I32Store(memarg) => {
                if read > 1
                    && let (Instruction::LocalGet32(addr_local), Instruction::LocalGet32(value_local)) =
                        (instructions[read - 2], instructions[read - 1])
                    && let (Ok(addr_local), Ok(value_local)) = (u8::try_from(addr_local), u8::try_from(value_local))
                {
                    instructions[read - 2] = Instruction::Nop;
                    instructions[read - 1] = Instruction::Nop;
                    instructions[read] = Instruction::I32StoreLocalLocal(memarg, addr_local, value_local);
                }
            }
            Instruction::I64Store(memarg) => {
                if read > 1
                    && let (Instruction::LocalGet32(addr_local), Instruction::LocalGet64(value_local)) =
                        (instructions[read - 2], instructions[read - 1])
                    && let (Ok(addr_local), Ok(value_local)) = (u8::try_from(addr_local), u8::try_from(value_local))
                {
                    instructions[read - 2] = Instruction::Nop;
                    instructions[read - 1] = Instruction::Nop;
                    instructions[read] = Instruction::I64StoreLocalLocal(memarg, addr_local, value_local);
                }
            }
            Instruction::MemoryFill(mem) => {
                if read > 1
                    && let (Instruction::I32Const(val), Instruction::I32Const(size)) =
                        (instructions[read - 2], instructions[read - 1])
                {
                    instructions[read - 2] = Instruction::Nop;
                    instructions[read - 1] = Instruction::Nop;
                    instructions[read] = Instruction::MemoryFillImm(mem, val as u8, size);
                }
            }

            Instruction::LocalGet32(dst) => {
                if read > 0
                    && let Instruction::LocalSet32(src) = instructions[read - 1]
                    && src == dst
                {
                    instructions[read - 1] = Instruction::LocalTee32(src);
                    instructions[read] = Instruction::Nop;
                }
            }
            Instruction::LocalGet64(dst) => {
                if read > 0
                    && let Instruction::LocalSet64(src) = instructions[read - 1]
                    && src == dst
                {
                    instructions[read - 1] = Instruction::LocalTee64(src);
                    instructions[read] = Instruction::Nop;
                }
            }
            Instruction::LocalGet128(dst) => {
                if read > 0
                    && let Instruction::LocalSet128(src) = instructions[read - 1]
                    && src == dst
                {
                    instructions[read - 1] = Instruction::LocalTee128(src);
                    instructions[read] = Instruction::Nop;
                }
            }
            Instruction::LocalSet32(dst) => {
                if read > 0 {
                    match instructions[read - 1] {
                        Instruction::LocalGet32(src) => {
                            instructions[read - 1] = Instruction::Nop;
                            instructions[read] =
                                if src == dst { Instruction::Nop } else { Instruction::LocalCopy32(src, dst) };
                        }
                        Instruction::I32Const(c) => {
                            instructions[read - 1] = Instruction::Nop;
                            instructions[read] = Instruction::LocalSetConst32(dst, c);
                        }
                        Instruction::F32Const(c) => {
                            instructions[read - 1] = Instruction::Nop;
                            instructions[read] =
                                Instruction::LocalSetConst32(dst, i32::from_ne_bytes(c.to_bits().to_ne_bytes()));
                        }
                        _ => {}
                    }
                }

                if read > 1 {
                    match (instructions[read - 2], instructions[read - 1]) {
                        (Instruction::LocalGet32(src), Instruction::I32AddConst(c)) if src == dst => {
                            instructions[read - 2] = Instruction::Nop;
                            instructions[read - 1] = Instruction::Nop;
                            instructions[read] = Instruction::LocalAddConst32(dst, c);
                        }
                        (Instruction::LocalGet32(addr), Instruction::I32Load(memarg)) => {
                            if let (Ok(addr), Ok(dst)) = (u8::try_from(addr), u8::try_from(dst)) {
                                instructions[read - 2] = Instruction::Nop;
                                instructions[read - 1] = Instruction::Nop;
                                instructions[read] = Instruction::I32LoadLocalSet(memarg, addr, dst);
                            }
                        }
                        _ => {}
                    }
                }
            }
            Instruction::LocalSet64(dst) => {
                if read > 0 {
                    match instructions[read - 1] {
                        Instruction::LocalGet64(src) => {
                            instructions[read - 1] = Instruction::Nop;
                            instructions[read] =
                                if src == dst { Instruction::Nop } else { Instruction::LocalCopy64(src, dst) };
                        }
                        Instruction::I64Const(c) => {
                            instructions[read - 1] = Instruction::Nop;
                            instructions[read] = Instruction::LocalSetConst64(dst, c);
                        }
                        Instruction::F64Const(c) => {
                            instructions[read - 1] = Instruction::Nop;
                            instructions[read] =
                                Instruction::LocalSetConst64(dst, i64::from_ne_bytes(c.to_bits().to_ne_bytes()));
                        }
                        _ => {}
                    }
                }

                if read > 1
                    && let (Instruction::LocalGet64(src), Instruction::I64AddConst(c)) =
                        (instructions[read - 2], instructions[read - 1])
                    && src == dst
                {
                    instructions[read - 2] = Instruction::Nop;
                    instructions[read - 1] = Instruction::Nop;
                    instructions[read] = Instruction::LocalAddConst64(dst, c);
                }
            }
            Instruction::LocalSet128(dst) => {
                if read > 0
                    && let Instruction::LocalGet128(src) = instructions[read - 1]
                {
                    instructions[read - 1] = Instruction::Nop;
                    instructions[read] =
                        if src == dst { Instruction::Nop } else { Instruction::LocalCopy128(src, dst) };
                }
            }
            Instruction::LocalTee32(dst) => {
                if read > 0
                    && let Instruction::LocalGet32(src) = instructions[read - 1]
                    && src == dst
                {
                    instructions[read] = Instruction::Nop;
                }

                if read > 1
                    && let (Instruction::LocalGet32(addr), Instruction::I32Load(memarg)) =
                        (instructions[read - 2], instructions[read - 1])
                    && let (Ok(addr), Ok(dst)) = (u8::try_from(addr), u8::try_from(dst))
                {
                    instructions[read - 2] = Instruction::Nop;
                    instructions[read - 1] = Instruction::Nop;
                    instructions[read] = Instruction::I32LoadLocalTee(memarg, addr, dst);
                }
            }
            Instruction::LocalTee64(dst) if read > 0 => match instructions[read - 1] {
                Instruction::LocalGet64(src) if src == dst => {
                    instructions[read] = Instruction::Nop;
                }
                Instruction::I64XorRotlConst(c) => {
                    instructions[read - 1] = Instruction::Nop;
                    instructions[read] = Instruction::I64XorRotlConstTee(c, dst);
                }
                _ => {}
            },
            Instruction::LocalTee128(dst) => {
                if read > 0
                    && let Instruction::LocalGet128(src) = instructions[read - 1]
                    && src == dst
                {
                    instructions[read] = Instruction::Nop;
                }
            }
            Instruction::Drop32 => {
                if read > 0
                    && let Instruction::LocalTee32(local) = instructions[read - 1]
                {
                    instructions[read - 1] = Instruction::LocalSet32(local);
                    instructions[read] = Instruction::Nop;
                }
            }
            Instruction::Drop64 => {
                if read > 0
                    && let Instruction::LocalTee64(local) = instructions[read - 1]
                {
                    instructions[read - 1] = Instruction::LocalSet64(local);
                    instructions[read] = Instruction::Nop;
                }
            }
            Instruction::Drop128 => {
                if read > 0
                    && let Instruction::LocalTee128(local) = instructions[read - 1]
                {
                    instructions[read - 1] = Instruction::LocalSet128(local);
                    instructions[read] = Instruction::Nop;
                }
            }
            _ => {}
        }
    }
}

fn dce(instructions: &mut Vec<Instruction>, function_data: &mut WasmFunctionData) {
    let old_len = instructions.len();
    if old_len == 0 {
        return;
    }

    let mut removed_before = Vec::with_capacity(old_len + 1);
    removed_before.push(0u32);
    instructions.iter().for_each(|instr| {
        let removed = removed_before.last().copied().unwrap_or(0) + u32::from(matches!(instr, Instruction::Nop));
        removed_before.push(removed);
    });

    let removed_total = removed_before[old_len];
    if removed_total == 0 {
        return;
    }

    let compacted_len = old_len as u32 - removed_total;

    function_data.branch_table_targets.iter_mut().for_each(|ip| {
        let old_target = *ip as usize;
        if old_target <= old_len {
            *ip -= removed_before[old_target];
            debug_assert!(*ip < compacted_len, "remapped jump target points past end of function");
        }
    });

    instructions.retain_mut(|instr| {
        let ip = match instr {
            Instruction::Jump(ip)
            | Instruction::JumpIfZero(ip)
            | Instruction::JumpIfNonZero(ip)
            | Instruction::BranchTable(ip, _, _) => ip,
            _ => return !matches!(instr, Instruction::Nop),
        };

        let old_target = *ip as usize;
        if old_target > old_len {
            return !matches!(instr, Instruction::Nop);
        }

        *ip -= removed_before[old_target];
        debug_assert!(*ip < compacted_len, "remapped jump target points past end of function");
        !matches!(instr, Instruction::Nop)
    });
}