Refactor controllers a little bit