use std::fmt::Debug;

use crate::{
    flatten::{Capture, FlatFn},
    parse::{AtomRef, ImmOperand, JsValue, opcodes::OpCode},
};

pub struct RegInstr {
    pub dsts: Vec<usize>,
    pub op: OpCode,
    pub imm: ImmOperand,
    pub args: Vec<usize>,
}

impl RegInstr {
    pub fn new(op: OpCode, imm: ImmOperand, dsts: Vec<usize>, args: Vec<usize>) -> Self {
        Self {
            dsts,
            op,
            imm,
            args,
        }
    }
}

impl 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: String,
    pub arg_count: usize,
    pub var_count: usize,
    pub value_names: Vec<AtomRef>,
    pub captures: Vec<Capture>,
    pub cpool: Vec<JsValue>,
    pub code: Vec<RegInstr>,
    next_val: usize,
}

impl RegFn {
    fn from_flat(flat_fn: FlatFn) -> Self {
        let mut stack = vec![];
        let mut output = Vec::with_capacity(flat_fn.code.len());
        let mut next_val = flat_fn.var_count; // Reserve values for locals

        let mut fresh_val = || {
            let v = next_val;
            next_val += 1;
            v
        };

        for instr in flat_fn.code {
            let mut reg_instr = RegInstr {
                dsts: vec![],
                op: instr.op,
                imm: instr.imm_operands,
                args: vec![],
            };

            let extra_pop = 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,
                },
            };

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

            for _ in 0..n_pop {
                let v = stack.pop().unwrap();
                reg_instr.args.push(v);
            }

            for _ in 0..n_push {
                let v = fresh_val();
                stack.push(v);
                reg_instr.dsts.push(v);
            }

            output.push(reg_instr);
        }

        Self {
            name: flat_fn.name,
            arg_count: flat_fn.arg_count,
            var_count: flat_fn.var_count,
            value_names: flat_fn.value_names,
            captures: flat_fn.captures,
            cpool: flat_fn.cpool,
            code: output,
            next_val,
        }
    }

    pub fn fresh_val(&mut self) -> usize {
        let v = self.next_val;
        self.next_val += 1;
        v
    }
}

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