Fix buttons not getting set in Modals

* Due to race condition when showing modal, buttons were added to
previous Toolbar.

* Minor refactor, moved `updateStyles` from BaseReactActivity to
StyleHelper class since it's now also used by RnnModal

android/app/src/main/java/com/reactnativenavigation/activities/BaseReactActivity.java

 import com.facebook.react.modules.core.DefaultHardwareBackBtnHandler;
 import com.facebook.react.shell.MainReactPackage;
 import com.reactnativenavigation.BuildConfig;
+import com.reactnativenavigation.controllers.ModalController;
 import com.reactnativenavigation.core.RctManager;
 import com.reactnativenavigation.core.objects.Button;
 import com.reactnativenavigation.core.objects.Screen;
+import com.reactnativenavigation.modal.RnnModal;
 import com.reactnativenavigation.packages.RnnPackage;
 import com.reactnativenavigation.utils.ContextProvider;
 import com.reactnativenavigation.utils.StyleHelper;
         setContentView(mReactRootView);
     }
 
-    public void setNavigationStyle(Screen screen) {
-        if (mToolbar != null) {
-            mToolbar.setStyle(screen);
-        }
-
-        StyleHelper.setWindowStyle(getWindow(), this, screen);
-    }
-
     @Override
     protected void onResume() {
         super.onResume();
         }
     }
 
-    public void updateStyles() {
-        try {
-            mToolbar.update(getCurrentScreen());
-            setNavigationStyle(getCurrentScreen());
-            mToolbar.setupToolbarButtonsAsync(getCurrentScreen());
-        } catch (Exception e) {
-            Log.w("RNNavigation", "Tried to update styles with no screen!");
-        }
-    }
-
     @Override
     protected void onPause() {
         super.onPause();
 
     @CallSuper
     public void push(Screen screen) {
-        setNavigationStyle(screen);
+        StyleHelper.updateStyles(mToolbar, screen);
         if (mToolbar != null) {
             mToolbar.update(screen);
 
 
     protected abstract String getCurrentNavigatorId();
 
-    protected abstract Screen getCurrentScreen();
+    @CallSuper
+    protected Screen getCurrentScreen() {
+        ModalController modalController = ModalController.getInstance();
+        if (modalController.isModalDisplayed()) {
+            RnnModal modal = modalController.get();
+            assert modal != null;
+            return modal.getCurrentScreen();
+        }
+
+        return null;
+    }
 
     public abstract int getScreenStackSize();
 
     @Override
     public boolean onCreateOptionsMenu(Menu menu) {
         mMenu = menu;
-        if (mToolbar != null) {
-            mToolbar.setupToolbarButtonsAsync(getCurrentScreen());
+        Screen currentScreen = getCurrentScreen();
+        if (mToolbar != null && currentScreen != null) {
+            mToolbar.setupToolbarButtonsAsync(currentScreen);
         }
         return super.onCreateOptionsMenu(menu);
     }

android/app/src/main/java/com/reactnativenavigation/activities/BottomTabActivity.java

 import com.reactnativenavigation.R;
 import com.reactnativenavigation.core.RctManager;
 import com.reactnativenavigation.core.objects.Screen;
+import com.reactnativenavigation.utils.StyleHelper;
 import com.reactnativenavigation.views.RnnToolBar;
 import com.reactnativenavigation.views.ScreenStack;
 
     }
 
     private void setupPages(ArrayList<Screen> screens) {
-        new SetupTabsTask(this, screens).execute();
+        new SetupTabsTask(this, mToolbar, screens).execute();
     }
 
     private void setupToolbar(ArrayList<Screen> screens) {
         mToolbar.setScreens(screens);
         Screen initialScreen = screens.get(0);
         mToolbar.update(initialScreen);
-        setNavigationStyle(initialScreen);
+        StyleHelper.updateStyles(mToolbar, initialScreen);
     }
 
     @Override
     protected void onResume() {
         super.onResume();
         if(mScreenStacks != null) {
-            updateStyles();
+            StyleHelper.updateStyles(mToolbar, getCurrentScreen());
         }
     }
 
                 stack.push(screen);
             }
         }
-        updateStyles();
+        StyleHelper.updateStyles(mToolbar, getCurrentScreen());
     }
 
     @Override
         for (ScreenStack stack: mScreenStacks) {
             if (stack.peek().navigatorId.equals(navigatorId)) {
                 Screen popped = stack.pop();
-                updateStyles();
+                StyleHelper.updateStyles(mToolbar, getCurrentScreen());
                 return popped;
             }
         }
 
     @Override
     protected Screen getCurrentScreen() {
-        return mScreenStacks.get(mCurrentStackPosition).peek();
+        Screen currentScreen = super.getCurrentScreen();
+        if (currentScreen != null) {
+            return currentScreen;
+        }
+
+        return mScreenStacks != null ? mScreenStacks.get(mCurrentStackPosition).peek() : null;
     }
 
     @Override
         mContentFrame.removeAllViews();
         mContentFrame.addView(mScreenStacks.get(position), new FrameLayout.LayoutParams(MATCH_PARENT, MATCH_PARENT));
         mCurrentStackPosition = position;
-        updateStyles();
+        StyleHelper.updateStyles(mToolbar, getCurrentScreen());
     }
 
     private static class SetupTabsTask extends AsyncTask<Void, Void, Map<Screen, Drawable>> {
         private BottomTabActivity mActivity;
+        private RnnToolBar mToolBar;
         private ArrayList<Screen> mScreens;
 
-        public SetupTabsTask(BottomTabActivity context, ArrayList<Screen> screens) {
+        public SetupTabsTask(BottomTabActivity context, RnnToolBar toolBar, ArrayList<Screen> screens) {
             mActivity = context;
+            mToolBar = toolBar;
             mScreens = screens;
         }
 
         @Override
         protected void onPostExecute(Map<Screen, Drawable> icons) {
             mActivity.setTabsWithIcons(mScreens, icons);
-            mActivity.updateStyles();
+            StyleHelper.updateStyles(mToolBar, mActivity.getCurrentScreen());
         }
     }
 

android/app/src/main/java/com/reactnativenavigation/activities/SingleScreenActivity.java

 import com.reactnativenavigation.R;
 import com.reactnativenavigation.core.RctManager;
 import com.reactnativenavigation.core.objects.Screen;
+import com.reactnativenavigation.utils.StyleHelper;
 import com.reactnativenavigation.views.RnnToolBar;
 import com.reactnativenavigation.views.ScreenStack;
 
 
         mScreenStack = new ScreenStack(this);
         FrameLayout contentFrame = (FrameLayout) findViewById(R.id.contentFrame);
+        assert contentFrame != null;
         contentFrame.addView(mScreenStack);
         mScreenStack.push(screen);
     }
 
     protected void setupToolbar(Screen screen) {
         mToolbar.update(screen);
-        setNavigationStyle(screen);
+        StyleHelper.updateStyles(mToolbar, screen);
     }
     
     @Override
     public void push(Screen screen) {
         super.push(screen);
         mScreenStack.push(screen);
-        updateStyles();
+        StyleHelper.updateStyles(mToolbar, screen);
     }
 
     @Override
     public Screen pop(String navigatorId) {
         super.pop(navigatorId);
         Screen screen = mScreenStack.pop();
-        updateStyles();
+        StyleHelper.updateStyles(mToolbar, screen);
         return screen;
     }
 
 
     @Override
     protected Screen getCurrentScreen() {
+        Screen currentScreen = super.getCurrentScreen();
+        if (currentScreen != null) {
+            return currentScreen;
+        }
+
         return mScreenStack.peek();
     }
 

android/app/src/main/java/com/reactnativenavigation/activities/TabActivity.java

         setNavigationStyle(initialScreen);
     }
 
-    @Override
     public void setNavigationStyle(Screen screen) {
-        super.setNavigationStyle(screen);
         mTabLayout.setStyle(screen);
     }
 

android/app/src/main/java/com/reactnativenavigation/adapters/ViewPagerAdapter.java

 import com.reactnativenavigation.activities.BaseReactActivity;
 import com.reactnativenavigation.core.RctManager;
 import com.reactnativenavigation.core.objects.Screen;
-import com.reactnativenavigation.utils.ContextProvider;
+import com.reactnativenavigation.utils.StyleHelper;
 import com.reactnativenavigation.views.RnnToolBar;
 import com.reactnativenavigation.views.ScreenStack;
 
         mToolbar.setTitle(newScreen.title == null ? "" : newScreen.title);
 
         // Set navigation color
-        ContextProvider.getActivityContext().setNavigationStyle(newScreen);
+        StyleHelper.updateStyles(mToolbar, newScreen);
 
         // Send tab selected event
         WritableMap params = Arguments.createMap();

android/app/src/main/java/com/reactnativenavigation/modal/RnnModal.java

 import android.app.Dialog;
 import android.content.Context;
 import android.content.DialogInterface;
+import android.support.annotation.Nullable;
 import android.view.LayoutInflater;
 import android.view.View;
 import android.view.Window;
 import com.reactnativenavigation.activities.BaseReactActivity;
 import com.reactnativenavigation.controllers.ModalController;
 import com.reactnativenavigation.core.objects.Screen;
-import com.reactnativenavigation.utils.ContextProvider;
 import com.reactnativenavigation.utils.SdkSupports;
 import com.reactnativenavigation.utils.StyleHelper;
 import com.reactnativenavigation.views.RctView;
         if (mScreenStack.isEmpty()) {
             dismiss();
         }
+        mToolBar.update(popped);
         return popped;
     }
 
+    @Nullable
+    public Screen getCurrentScreen() {
+        return mScreenStack.isEmpty() ? null : mScreenStack.peek();
+    }
+
     @Override
     public void onBackPressed() {
         pop();
     @Override
     public void onDismiss(DialogInterface dialog) {
         ModalController.getInstance().remove();
-        ContextProvider.getActivityContext().updateStyles();
+        StyleHelper.updateStyles(mToolBar, getCurrentScreen());
     }
 }

android/app/src/main/java/com/reactnativenavigation/utils/StyleHelper.java

 
 import android.content.Context;
 import android.support.v4.content.ContextCompat;
+import android.util.Log;
 import android.view.Window;
 
+import com.reactnativenavigation.activities.BaseReactActivity;
 import com.reactnativenavigation.core.objects.Screen;
+import com.reactnativenavigation.views.RnnToolBar;
 
 /**
  * Created by guyc on 07/05/16.
  */
 public class StyleHelper {
+
+    public static void updateStyles(RnnToolBar toolBar, Screen screen) {
+        try {
+            toolBar.updateAndSetButtons(screen);
+            setWindowStyle(screen);
+            toolBar.setupToolbarButtonsAsync(screen);
+        } catch (Exception e) {
+            Log.w("RNNavigation", "Tried to update styles with no screen!");
+        }
+    }
+
+    private static void setWindowStyle(Screen screen) {
+        BaseReactActivity context = ContextProvider.getActivityContext();
+        if (context != null) {
+            StyleHelper.setWindowStyle(context.getWindow(), context, screen);
+        }
+    }
+
     public static void setWindowStyle(Window window, Context context, Screen screen) {
         if (SdkSupports.lollipop()) {
             final int black = ContextCompat.getColor(context, android.R.color.black);

android/app/src/main/java/com/reactnativenavigation/views/RnnToolBar.java

         ((AppCompatActivity) getContext()).setSupportActionBar(this);
         setTitle(screen.title);
         setStyle(screen);
+    }
+
+    public void updateAndSetButtons(Screen screen) {
+        update(screen);
         setupToolbarButtonsAsync(screen);
     }
 
             // Remove prev screen buttons
             if(mOldButtons == null) {
                 menu.clear();
-            }
-            else {
+            } else {
                 for (Button btn : mOldButtons) {
                     menu.removeItem(btn.getItemId());
                 }
             for (i = 0; i < mNewButtons.size(); i++) {
                 Button button = mNewButtons.get(i);
                 MenuItem item = menu.add(Menu.NONE, button.getItemId(), i, button.title);
+                item.setShowAsAction(MenuItem.SHOW_AS_ACTION_IF_ROOM);
+
                 if (icons.containsKey(button.id)) {
                     Drawable icon = icons.get(button.id);
                     if (mTintColor != null) {
                         ImageUtils.tint(icon, mTintColor);
                     }
-                    item.setIcon(icon).setShowAsAction(MenuItem.SHOW_AS_ACTION_IF_ROOM);
+                    item.setIcon(icon);
                 }
             }