Refactoring to make the code cleaner