Update to use command-line arguments