refactored everything into a nice class