Refactored to strategy pattern