refactor and add more checks
refactor: pass args as well