Big refactoring for nicer API