Refactor code and implement protocols