Refactor: Use states

This pull request splits the code into several states.