use crate::syntax::ast::{Expression, Program, Statement};

impl Program {
    pub fn simplify(mut self) -> Self {
        let mut new_statements = Vec::new();
        let mut gensym_index = 1;

        for stmt in self.statements.drain(..) {
            new_statements.append(&mut stmt.simplify(&mut gensym_index));
        }

        self.statements = new_statements;
        self
    }
}

impl Statement {
    pub fn simplify(self, gensym_index: &mut usize) -> Vec<Statement> {
        let mut new_statements = vec![];

        match self {
            Statement::Print(_, _) => new_statements.push(self),
            Statement::Binding(_, _, Expression::Reference(_, _)) => new_statements.push(self),
            Statement::Binding(_, _, Expression::Value(_, _)) => new_statements.push(self),
            Statement::Binding(loc, name, value) => {
                let (mut prereqs, new_value) = value.rebind(&name, gensym_index);
                new_statements.append(&mut prereqs);
                new_statements.push(Statement::Binding(loc, name, new_value))
            }
        }

        new_statements
    }
}

impl Expression {
    fn rebind(self, base_name: &str, gensym_index: &mut usize) -> (Vec<Statement>, Expression) {
        match self {
            Expression::Value(_, _) => (vec![], self),
            Expression::Reference(_, _) => (vec![], self),
            Expression::Primitive(loc, prim, mut expressions) => {
                let mut prereqs = Vec::new();
                let mut new_exprs = Vec::new();

                for expr in expressions.drain(..) {
                    let (mut cur_prereqs, arg) = expr.rebind(base_name, gensym_index);
                    prereqs.append(&mut cur_prereqs);
                    new_exprs.push(arg);
                }

                let new_name = format!("<{}:{}>", base_name, *gensym_index);
                *gensym_index += 1;
                prereqs.push(Statement::Binding(
                    loc.clone(),
                    new_name.clone(),
                    Expression::Primitive(loc.clone(), prim, new_exprs),
                ));

                (prereqs, Expression::Reference(loc, new_name))
            }
        }
    }
}