Added Enum class
Small refactoring