more clean up and refactor