Initial code refactor