Code refactor and cleanup