Code refactoring and updated features

cr: https://code.amazon.com/reviews/CR-107207477
