Revert "Only added one feature while refactoring"

This reverts commit 9726b366fc535e16f317a22d1a0a09a3223d8072.