Complete refactor code