Code refactor based on PR feedback