mod basic_block;

use crate::{
    flatten::FlatFn,
    libqjs::parse::{AtomRef, ClosureVar, ImmOperand, Instr, JsValue, opcodes::OpCode},
};

pub use basic_block::BasicBlock;
use basic_block::find_basic_blocks;

#[derive(Clone)]
pub struct RegInstr {
    pub dsts: Vec<usize>,
    pub op: OpCode,
    pub imm: ImmOperand,
    pub args: Vec<usize>,
}

impl RegInstr {
    fn nop() -> Self {
        Self {
            dsts: vec![],
            op: OpCode::Nop,
            imm: ImmOperand::None,
            args: vec![],
        }
    }
}

impl std::fmt::Debug for RegInstr {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{:?} = {:?} {:?} {:?}",
            self.dsts, self.op, self.imm, self.args
        )
    }
}

#[derive(Debug)]
pub struct RegFn {
    pub name: AtomRef,
    pub arg_count: usize,
    pub val_count: usize,
    pub cpool: Vec<JsValue>,
    pub captures: Vec<ClosureVar>,
    pub code: Vec<RegInstr>,
    pub blocks: Vec<BasicBlock>,
}

impl RegFn {
    fn from_flat(flat_fn: FlatFn) -> Self {
        let mut blocks = find_basic_blocks(&flat_fn.code);
        let mut out = vec![RegInstr::nop(); flat_fn.code.len()];

        let mut next_val = flat_fn.var_count; // Reserve values for locals

        let mut entry_stacks: Vec<Option<Vec<usize>>> = vec![None; blocks.len()];
        let mut worklist = vec![];

        entry_stacks[0] = Some(vec![]);
        worklist.push(0);

        while let Some(block_idx) = worklist.pop() {
            let block = &mut blocks[block_idx];
            let mut stack: Vec<usize> = entry_stacks[block_idx].clone().unwrap_or_default();

            walk_block(block, &flat_fn.code, &mut stack, &mut next_val, &mut out);

            for succ_idx in block.succ.clone() {
                let succ_stack = &mut entry_stacks[succ_idx];

                match succ_stack {
                    Some(succ_stack) => add_block_postfix(block, &stack, succ_stack),
                    None => {
                        *succ_stack = Some(stack.clone());
                        worklist.push(succ_idx);
                    }
                }
            }
        }

        Self {
            name: flat_fn.name,
            arg_count: flat_fn.arg_count,
            val_count: next_val,
            cpool: flat_fn.cpool,
            captures: flat_fn.captures,
            code: out,
            blocks,
        }
    }
}

fn walk_block(
    block: &BasicBlock,
    code: &[Instr],
    stack: &mut Vec<usize>,
    next_val: &mut usize,
    out: &mut [RegInstr],
) {
    let mut fresh_val = || {
        let v = *next_val;
        *next_val += 1;
        v
    };

    for instr_idx in block.start..block.end {
        let instr = &code[instr_idx];

        let mut out_instr = RegInstr {
            dsts: vec![],
            op: instr.op,
            imm: instr.imm_operands,
            args: vec![],
        };

        let info = instr.op.info();
        let n_pop = info.n_pop as usize + extra_pops(instr);
        let n_push = info.n_push as usize;

        // Handle this instr's pops
        for _ in 0..n_pop {
            let v = stack.pop().unwrap();
            out_instr.args.push(v);
        }

        // Locals are not renamed, local read/writes need special handling
        // Handle read of a local value
        match is_local_read(instr) {
            Some(loc) => {
                stack.push(loc);
                out_instr.op = OpCode::Nop;
                out_instr.imm = ImmOperand::None;
                out[instr_idx] = out_instr;
                continue;
            }
            None if instr.op == OpCode::GetLoc0Loc1 => {
                stack.push(0);
                stack.push(1);
                out_instr.op = OpCode::Nop;
                out_instr.imm = ImmOperand::None;
                out[instr_idx] = out_instr;
                continue;
            }
            None => (),
        }

        // Handle write of a local value
        match is_local_write(instr) {
            Some(loc) => {
                out_instr.op = OpCode::SetLoc;
                out_instr.imm = ImmOperand::None;
                out_instr.dsts.push(loc);
                out[instr_idx] = out_instr;

                if n_push == 1 {
                    stack.push(loc);
                }

                continue;
            }
            None if instr.op == OpCode::SetLocUninitialized => {
                out_instr.op = OpCode::Nop;
                out_instr.imm = ImmOperand::None;
                out[instr_idx] = out_instr;
                continue;
            }
            None => (),
        }

        // Handle this instr's pushes
        for _ in 0..n_push {
            let v = fresh_val();
            stack.push(v);
            out_instr.dsts.push(v);
        }

        out[instr_idx] = out_instr;
    }
}

fn extra_pops(instr: &Instr) -> usize {
    match instr.op {
        OpCode::Call0 => 0,
        OpCode::Call1 => 1,
        OpCode::Call2 => 2,
        OpCode::Call3 => 3,
        _ => match instr.imm_operands {
            ImmOperand::Npop(npop) => npop as usize,
            ImmOperand::NpopU16 { npop, .. } => npop as usize,
            _ => 0,
        },
    }
}

fn is_local_read(instr: &Instr) -> Option<usize> {
    match instr.op {
        OpCode::GetLoc0 => Some(0),
        OpCode::GetLoc1 => Some(1),
        OpCode::GetLoc2 => Some(2),
        OpCode::GetLoc3 => Some(3),

        OpCode::GetLoc | OpCode::GetLoc8 | OpCode::GetLocCheck => match instr.imm_operands {
            ImmOperand::Loc(k) => Some(k as usize),
            _ => unreachable!(),
        },

        _ => None,
    }
}

fn is_local_write(instr: &Instr) -> Option<usize> {
    match instr.op {
        OpCode::PutLoc0 | OpCode::SetLoc0 => Some(0),
        OpCode::PutLoc1 | OpCode::SetLoc1 => Some(1),
        OpCode::PutLoc2 | OpCode::SetLoc2 => Some(2),
        OpCode::PutLoc3 | OpCode::SetLoc3 => Some(3),

        OpCode::PutLoc
        | OpCode::PutLoc8
        | OpCode::PutLocCheck
        | OpCode::PutLocCheckInit
        | OpCode::SetLoc
        | OpCode::SetLoc8 => match instr.imm_operands {
            ImmOperand::Loc(k) => Some(k as usize),
            _ => unreachable!(),
        },

        _ => None,
    }
}

fn add_block_postfix(block: &mut BasicBlock, stack: &[usize], succ_stack: &[usize]) {
    let overlap = stack.iter().rev().zip(succ_stack.iter().rev());

    for (val, succ_val) in overlap {
        block.postfix.push(RegInstr {
            dsts: vec![*succ_val],
            op: OpCode::SetLoc,
            imm: ImmOperand::None,
            args: vec![*val],
        });
    }
}

pub fn regify(flat_fns: Vec<FlatFn>) -> Vec<RegFn> {
    flat_fns.into_iter().map(RegFn::from_flat).collect()
}