Refactor client code to have proper structure