Changed design to finite state machine