Refactor controller behaviour