More refactoring work!

Still need to refactor more.