Refactor CLI handling