Refactor and add cleanup function.