refactoring, seems to be working properly now