Refactor into classes (partial)