Some refactoring and add of a parameter that will be necessary.

Bug: 4392433
Change-Id: I17f001a7efd4f69f4c35f94ee1ca8e97391b81d5