refactor the ifs a bit