From dfc63110207c8cb59ec070990f81dd1abd0dd153 Mon Sep 17 00:00:00 2001
From: EnixCoda <enixcoda@gmail.com>
Date: Thu, 30 Jan 2025 19:53:00 +0800
Subject: [PATCH] feat: auto-unmount prod version when dev version is enabled

---
 src/content.tsx               | 54 ++++++++++++++++++++++++++---------
 src/utils/waitForNextEvent.ts | 28 ++++++++++++++++++
 2 files changed, 68 insertions(+), 14 deletions(-)
 create mode 100644 src/utils/waitForNextEvent.ts

diff --git a/src/content.tsx b/src/content.tsx
index f65dfa7..f723e02 100644
--- a/src/content.tsx
+++ b/src/content.tsx
@@ -1,40 +1,66 @@
 import { Gitako } from 'components/Gitako'
+import { IN_PRODUCTION_MODE } from 'env'
 import React, { useCallback } from 'react'
 import { createRoot } from 'react-dom/client'
 import { insertMountPoint, insertSideBarMountPoint } from 'utils/DOMHelper'
 import { useAfterRedirect } from 'utils/hooks/useFastRedirect'
+import { waitForNext } from 'utils/waitForNextEvent'
 import './content.scss'
 
-if (document.readyState === 'loading') {
-  document.addEventListener('DOMContentLoaded', init)
-} else {
-  init()
-}
-
-async function init() {
-  await injectStyles(browser.runtime.getURL('content.css'))
+const renderReact = () => {
   const mountPoint = insertSideBarMountPoint()
   const MountPointWatcher = () => {
     useAfterRedirect(useCallback(() => insertMountPoint(() => mountPoint), []))
     return null
   }
-
-  createRoot(mountPoint).render(
+  const root = createRoot(mountPoint)
+  root.render(
     <>
       <MountPointWatcher />
       <Gitako />
     </>,
   )
+
+  return () => {
+    root.unmount()
+  }
 }
 
 // injects a copy of stylesheets so that other extensions(e.g. dark reader) could read
 // resolves when style is loaded to prevent render without proper styles
-async function injectStyles(url: string) {
-  return new Promise<void>(resolve => {
+const injectStyles = (url: string) =>
+  new Promise<() => void>(resolve => {
     const linkElement = document.createElement('link')
     linkElement.setAttribute('rel', 'stylesheet')
     linkElement.setAttribute('href', url)
-    linkElement.onload = () => resolve()
+    const unload = () => {
+      linkElement.remove()
+    }
+    linkElement.onload = () => resolve(unload)
     document.head.appendChild(linkElement)
   })
-}
+
+const GitakoExclusiveEventType = 'GITAKO_EXCLUSIVE_EVENT'
+const GitakoMountedEventType = 'GITAKO_MOUNTED_EVENT'
+
+Promise.resolve()
+  .then(() =>
+    document.readyState === 'loading' ? waitForNext.documentEvent('DOMContentLoaded') : null,
+  )
+  .then(() =>
+    Promise.all([injectStyles(browser.runtime.getURL('content.css')), renderReact()]).then(
+      ([unmountStyles, unmountReact]) =>
+        () =>
+          Promise.all([unmountStyles(), unmountReact()]),
+    ),
+  )
+  .then(unmount => {
+    document.dispatchEvent(new CustomEvent(GitakoMountedEventType))
+    if (IN_PRODUCTION_MODE) {
+      waitForNext.documentEvent(GitakoExclusiveEventType).then(unmount)
+    } else {
+      waitForNext
+        .documentEvent(GitakoMountedEventType)
+        .then(() => document.dispatchEvent(new CustomEvent(GitakoExclusiveEventType)))
+    }
+  })
diff --git a/src/utils/waitForNextEvent.ts b/src/utils/waitForNextEvent.ts
new file mode 100644
index 0000000..43c7823
--- /dev/null
+++ b/src/utils/waitForNextEvent.ts
@@ -0,0 +1,28 @@
+export const waitForNextDocumentEvent = <K extends keyof DocumentEventMap>(
+  type: EnumString<K>,
+  options?: boolean | AddEventListenerOptions,
+) =>
+  new Promise(resolve => {
+    const listener = (ev: DocumentEventMap[K] | Event) => {
+      document.removeEventListener(type, listener, options)
+      resolve(ev)
+    }
+    document.addEventListener(type, listener, options)
+  })
+
+export const waitForNextWindowEvent = <K extends keyof WindowEventMap>(
+  type: EnumString<K>,
+  options?: boolean | AddEventListenerOptions,
+) =>
+  new Promise(resolve => {
+    const listener = (ev: WindowEventMap[K] | Event) => {
+      window.removeEventListener(type, listener, options)
+      resolve(ev)
+    }
+    window.addEventListener(type, listener, options)
+  })
+
+export const waitForNext = {
+  documentEvent: waitForNextDocumentEvent,
+  windowEvent: waitForNextWindowEvent,
+}