From 100317ae4ed29dd820d6b03b56304c0255b3302b Mon Sep 17 00:00:00 2001
From: Dominic Go <18517029+dominicstop@users.noreply.github.com>
Date: Thu, 13 Apr 2023 07:20:56 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=92=AB=20Update:=20`UIViewController+Swiz?=
 =?UTF-8?q?zling`?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../UIViewController+Swizzling.swift          | 105 ++++++++++++------
 1 file changed, 73 insertions(+), 32 deletions(-)

diff --git a/ios/src_library/Helpers+Utilities/UIViewController+Swizzling.swift b/ios/src_library/Helpers+Utilities/UIViewController+Swizzling.swift
index b5fef7bb..28a027d9 100644
--- a/ios/src_library/Helpers+Utilities/UIViewController+Swizzling.swift
+++ b/ios/src_library/Helpers+Utilities/UIViewController+Swizzling.swift
@@ -11,6 +11,62 @@ extension UIViewController {
   
   static var isSwizzled = false;
   
+  fileprivate var modalWrapper: RNIModalViewControllerWrapper? {
+    RNIModalManagerShared.getModalInstance(
+      forPresentingViewController: self
+    ) as? RNIModalViewControllerWrapper;
+  };
+  
+  @discardableResult
+  func registerIfNeeded(
+    viewControllerToPresent: UIViewController? = nil
+  ) -> RNIModalViewControllerWrapper? {
+    guard !(self is RNIModalViewController),
+          !RNIModalManagerShared.isRegistered(viewController: self)
+    else { return nil };
+    
+    let modalWrapper = self.modalWrapper ?? RNIModalViewControllerWrapper();
+    
+    modalWrapper.presentingViewController = self;
+    modalWrapper.modalViewController = viewControllerToPresent;
+    
+    return modalWrapper;
+  };
+  
+  func getPresentedModal(
+    viewControllerToPresent: UIViewController? = nil
+  ) -> (any RNIModal)? {
+    if let presentedModalVC = viewControllerToPresent as? RNIModalViewController {
+      return presentedModalVC.modalViewRef;
+      
+    } else if let presentedVC = self.presentedViewController,
+              let presentedModalVC = presentedVC as? RNIModalViewController {
+      return presentedModalVC.modalViewRef;
+      
+    } else if let presentingModalVC = self as? RNIModalViewController,
+              let presentingModal = presentingModalVC.modalViewRef,
+              let presentedModalVC = presentingModal.modalVC {
+      return presentedModalVC.modalViewRef;
+      
+    } else if let viewControllerToPresent = viewControllerToPresent,
+              let presentedModal =
+                RNIModalManagerShared.getModalInstance(
+                  forPresentedViewController: viewControllerToPresent
+                ) {
+      return presentedModal;
+      
+    } else if let presentingModalWrapper = self.modalWrapper,
+              let presentingModalVC = presentingModalWrapper.modalViewController,
+              let presentedModal =
+                RNIModalManagerShared.getModalInstance(
+                  forPresentedViewController: presentingModalVC
+                ) {
+      return presentedModal;
+    };
+    
+    return nil;
+  };
+  
   @objc fileprivate func _swizzled_present(
     _ viewControllerToPresent: UIViewController,
     animated flag: Bool,
@@ -25,27 +81,14 @@ extension UIViewController {
     );
     #endif
     
-    let presentingModal: (any RNIModal)? = {
-      if let modalVC = self as? RNIModalViewController {
-        return modalVC.modalViewRef
-      };
-      
-      let matchingModalWrapper = RNIModalManagerShared.getModalInstance(
-        forPresentingViewController: viewControllerToPresent
-      ) as? RNIModalViewControllerWrapper;
-      
-      let modalWrapper =
-        matchingModalWrapper ?? RNIModalViewControllerWrapper();
-      
-      modalWrapper.presentingViewController = self;
-      modalWrapper.modalViewController = viewControllerToPresent;
-      
-      return modalWrapper;
-    }();
+    self.registerIfNeeded(viewControllerToPresent: viewControllerToPresent);
     
-    if let presentingModal = presentingModal {
-      presentingModal.modalPresentationNotificationDelegate
-        .notifyOnModalWillShow(sender: presentingModal);
+    let presentedModal =
+      self.getPresentedModal(viewControllerToPresent: viewControllerToPresent);
+    
+    if let presentedModal = presentedModal {
+      presentedModal.modalPresentationNotificationDelegate
+        .notifyOnModalWillShow(sender: presentedModal);
     };
     
     // call original impl.
@@ -58,9 +101,9 @@ extension UIViewController {
       );
       #endif
       
-      if let presentingModal = presentingModal {
-        presentingModal.modalPresentationNotificationDelegate
-          .notifyOnModalDidShow(sender: presentingModal);
+      if let presentedModal = presentedModal {
+        presentedModal.modalPresentationNotificationDelegate
+          .notifyOnModalDidShow(sender: presentedModal);
       };
       
       completion?();
@@ -80,13 +123,11 @@ extension UIViewController {
     );
     #endif
     
-    let presentingModal = RNIModalManagerShared.getModalInstance(
-      forPresentingViewController: self
-    ) as? RNIModalViewControllerWrapper;
+    let presentedModal = self.getPresentedModal();
     
-    if let presentingModal = presentingModal {
-      presentingModal.modalPresentationNotificationDelegate
-        .notifyOnModalWillHide(sender: presentingModal);
+    if let presentedModal = presentedModal {
+      presentedModal.modalPresentationNotificationDelegate
+        .notifyOnModalWillHide(sender: presentedModal);
     };
     
     // call original impl.
@@ -99,9 +140,9 @@ extension UIViewController {
       );
       #endif
       
-      if let presentingModal = presentingModal {
-        presentingModal.modalPresentationNotificationDelegate
-          .notifyOnModalDidHide(sender: presentingModal);
+      if let presentedModal = presentedModal {
+        presentedModal.modalPresentationNotificationDelegate
+          .notifyOnModalDidHide(sender: presentedModal);
       };
       
       completion?();