Updated to use CLI argument parsing