refator code and fix a bug