Clean up argument parsing code in main