Refactored controller code (again :D)