Refactored input handling in MainMenuScreen and LobbyScreen