General cleanup and refactor