refactoring on state machine