Further refactoring and fixes...