Cleanup & refactoring