use crate::Word;
use super::{ Program, Address };

impl Program {
    /* Intermediate code generation */
    pub(crate) fn collapse_instructions(&mut self) {
        // Fix instruction addresses
        let mut offsets = Vec::new();
        for instr in &self.instructions {
            offsets.push(instr.collapse_len() + offsets.last().unwrap_or(&0));
        }
        offsets.insert(0, 0);

        for instr in &mut self.instructions {
            for arg in &mut instr.args {
                match arg {
                    Address::Instruction(num)
                        => *arg = Address::Instruction(offsets[*num]),
                    Address::InstructionOffset(num, offset)
                        => *arg = Address::InstructionOffset(offsets[*num], *offset),
                    Address::InstructionConstant(num)
                        => *arg = Address::InstructionConstant(offsets[*num]),
                    _ => {},
                }
            }
        }

        // Collapse all instructions
        let mut collapsed = Vec::new();

        for instr in &mut self.instructions {
            collapsed.extend(instr.collapse(collapsed.len()));
        }

        self.instructions = collapsed;
    }

    /* Machine code generation */

    pub(crate) fn instructions_to_constants(&mut self) {
        let absolute_start = self.absolute_start + 1;

        let mut constlen = self.constants.len();
        for instr in &self.instructions {
            for arg in &instr.args {
                if let Address::InstructionConstant(_) = arg {
                    constlen += 1;
                }
            }
        }

        for instr in &mut self.instructions {
            for arg in &mut instr.args {
                if let Address::InstructionConstant(num) = arg {
                    self.constants.push((constlen + absolute_start + *num * 3) as Word);
                    *arg = Address::Constants(self.constants.len() - 1);
                }
            }
        }

        self.constants.insert(0, (absolute_start + constlen) as Word);

        for instr in &mut self.instructions {
            for arg in &mut instr.args {
                self.constants.push(match arg {
                    Address::Memory(num) => *num,
                    Address::Constants(num) => absolute_start + *num,
                    Address::Instruction(num)
                        => constlen + absolute_start + *num * 3,
                    Address::InstructionOffset(num, off)
                        => *off + constlen + absolute_start + *num * 3,
                    _ => unreachable!("Definition addresses should've already been removed!"),
                } as Word);
            }
        }
        self.instructions.clear();
    }

    pub(crate) fn to_bytes(self, pc: usize) -> Vec<u8> {
        let mut res = Vec::new();
        res.resize(self.absolute_start, Default::default());
        res[pc] = self.constants[0];
        res.extend(self.constants);
        res
            .into_iter()
            .map(|x| x.to_be_bytes())
            .flatten()
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::lexer::Token;

    #[test]
    fn after_two_collapses_only_subleq_remains() {
        let mut prog = Program::new(vec![
            Token::Operator("rts".to_string()),
            Token::Operator("jsr".to_string()),
            Token::Address(1000),
            Token::Operator("mov".to_string()),
            Token::Address(1001),
            Token::Address(1002),
            Token::Operator("subleq".to_string()),
            Token::Address(1001),
            Token::Address(1002),
            Token::Address(1001),

        ], 10);
        prog.collapse_instructions();
        prog.collapse_instructions();

        for i in prog.instructions {
            assert_eq!(i.name, "subleq");
        }
    }
}