refactor to use enum
remove unneed fn and refactor