Refactor state transition logic

Implement a more robust version of the state transition logic (where
exiting states is managed in the same way as entering states)