Refactoring classes for better maintainability