[FEATURE] Done some refactoring for better inheritance