Refactor some stuff fixed bug