First code refactor