Refactor argument handling