Fix crash when rotating custom theme screen

Initialize custom steps before fragments are restored,
and save current step for reinitializing in the right one.

Fixes: 130574920
Change-Id: I7b3ece938271cf6d155da70891ee94b1497ec927
diff --git a/src/com/android/customization/picker/theme/CustomThemeActivity.java b/src/com/android/customization/picker/theme/CustomThemeActivity.java
index c2d4e51..6c21d4a 100644
--- a/src/com/android/customization/picker/theme/CustomThemeActivity.java
+++ b/src/com/android/customization/picker/theme/CustomThemeActivity.java
@@ -67,6 +67,7 @@
     public static final int RESULT_THEME_APPLIED = 20;
 
     private static final String TAG = "CustomThemeActivity";
+    private static final String KEY_STATE_CURRENT_STEP = "CustomThemeActivity.currentStep";
 
     private ThemesUserEventLogger mUserEventLogger;
     private List<ComponentStep<?>> mSteps;
@@ -77,7 +78,6 @@
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
-        super.onCreate(savedInstanceState);
         CustomizationInjector injector = (CustomizationInjector) InjectorProvider.getInjector();
         mUserEventLogger = (ThemesUserEventLogger) injector.getUserEventLogger(this);
         Intent intent = getIntent();
@@ -103,10 +103,17 @@
                 new OverlayManagerCompat(this),
                 mUserEventLogger);
         mThemeManager.fetchOptions(null, false);
+
+        int currentStep = 0;
+        if (savedInstanceState != null) {
+            currentStep = savedInstanceState.getInt(KEY_STATE_CURRENT_STEP);
+        }
+        initSteps(currentStep);
+
+        super.onCreate(savedInstanceState);
         setContentView(R.layout.activity_custom_theme);
         mApplyButton = findViewById(R.id.next_button);
         mApplyButton.setOnClickListener(view -> onNextOrApply());
-        initSteps();
 
         FragmentManager fm = getSupportFragmentManager();
         Fragment fragment = fm.findFragmentById(R.id.fragment_container);
@@ -116,6 +123,12 @@
         }
     }
 
+    @Override
+    protected void onSaveInstanceState(Bundle outState) {
+        super.onSaveInstanceState(outState);
+        outState.putInt(KEY_STATE_CURRENT_STEP, mCurrentStep);
+    }
+
     private void navigateToStep(int i) {
         FragmentManager fragmentManager = getSupportFragmentManager();
         ComponentStep step = mSteps.get(i);
@@ -132,7 +145,7 @@
         updateApplyButtonLabel();
     }
 
-    private void initSteps() {
+    private void initSteps(int currentStep) {
         mSteps = new ArrayList<>();
         OverlayManagerCompat manager = new OverlayManagerCompat(this);
         mSteps.add(new FontStep(new FontOptionsProvider(this, manager), 0, 4));
@@ -140,7 +153,7 @@
         mSteps.add(new ColorStep(new ColorOptionsProvider(this, manager, mCustomThemeManager),
                 2, 4));
         mSteps.add(new ShapeStep(new ShapeOptionsProvider(this, manager), 3, 4));        
-        mCurrentStep = 0;
+        mCurrentStep = currentStep;
     }
 
     private void onNextOrApply() {