Refactor and improve user experience