Add class, methods, refactor and working!