refactoring + some clean up