Refactor code base on Thong Nguyen comment