Add state machine for main, add menu screen draw function