Refactor with new module structure