more code refactoring and improvements