Refactor and prettify