Trying to structure the code in OOP