Refactor classes and add stubs