Finished refactor. Need to test it.