Rewrite to class based architecture