Refactor menu handling