Address code review and style issues