Large code refactor