From a8a533cb732bf9f9ba1731204ec05914b9aa990e Mon Sep 17 00:00:00 2001
From: Aria Minaei <aria.minaei@gmail.com>
Date: Sat, 5 Aug 2023 10:01:44 +0200
Subject: [PATCH] Don't cleanup StyledShadow on unmount

---
 theatre/studio/src/css.tsx | 22 ++++++++++++++++------
 1 file changed, 16 insertions(+), 6 deletions(-)

diff --git a/theatre/studio/src/css.tsx b/theatre/studio/src/css.tsx
index de8680c..302f114 100644
--- a/theatre/studio/src/css.tsx
+++ b/theatre/studio/src/css.tsx
@@ -113,14 +113,24 @@ const ProvideStyledShadow: React.FC<{
 
   useLayoutEffect(() => {
     if (!template) return
-    const shadowRoot = (template.parentNode as HTMLElement).attachShadow({
-      mode: 'open',
-    })
+    const {parentNode} = template
+    if (!parentNode) return
+
+    const hadShadowRoot =
+      // @ts-ignore
+      !!parentNode.shadowRoot
+
+    const shadowRoot = hadShadowRoot
+      ? // @ts-ignore
+        parent.shadowRoot
+      : (parentNode as HTMLElement).attachShadow({
+          mode: 'open',
+        })
+
     setShadowRoot(shadowRoot)
 
-    return () => {
-      template.parentNode?.removeChild(shadowRoot)
-    }
+    // no need to cleanup. The parent will be removed anyway if cleanup
+    // is needed.
   }, [template])
 
   const [portalLayerRef, portalLayer] = useRefAndState<HTMLDivElement>(