More refactoring and small fixes