More refactoring, need to still look at tests