Refactored some of the controllers code