Make category pill work on invidious and mobile youtube

2 files changed, 26 insertions, 5 deletions

diff --git a/src/content.ts b/src/content.ts
index 93e3fb8c..23d3b9d3 100644
--- a/src/content.ts
+++ b/src/content.ts
@@ -647,7 +647,7 @@ function setupCategoryPill() {
         categoryPill = new CategoryPill();
     }
 
-    categoryPill.attachToPage();
+    categoryPill.attachToPage(onMobileYouTube);
 }
 
 async function sponsorsLookup(id: string, keepOldSubmissions = true) {
diff --git a/src/render/CategoryPill.tsx b/src/render/CategoryPill.tsx
index 6c6695e2..4fe846d0 100644
--- a/src/render/CategoryPill.tsx
+++ b/src/render/CategoryPill.tsx
@@ -9,15 +9,19 @@ export class CategoryPill {
     ref: React.RefObject<CategoryPillComponent>;
 
     unsavedState: CategoryPillState;
+
+    mutationObserver?: MutationObserver;
     
     constructor() {
         this.ref = React.createRef();
     }
 
-    async attachToPage(): Promise<void> {
-        // TODO: Mobile and invidious
-        const referenceNode = await GenericUtils.wait(() => document.querySelector(".ytd-video-primary-info-renderer.title") as HTMLElement);
-        
+    async attachToPage(onMobileYouTube: boolean): Promise<void> {
+        const referenceNode = 
+            await GenericUtils.wait(() => 
+                // YouTube, Mobile YouTube, Invidious
+                document.querySelector(".ytd-video-primary-info-renderer.title, .slim-video-information-title, #player-container + .h-box > h1") as HTMLElement);
+
         if (referenceNode && !referenceNode.contains(this.container)) {
             this.container = document.createElement('span');
             this.container.id = "categoryPill";
@@ -26,6 +30,10 @@ export class CategoryPill {
             referenceNode.prepend(this.container);
             referenceNode.style.display = "flex";
 
+            if (this.ref.current) {
+                this.unsavedState = this.ref.current.state;
+            }
+
             ReactDOM.render(
                 <CategoryPillComponent ref={this.ref} />,
                 this.container
@@ -35,6 +43,19 @@ export class CategoryPill {
                 this.ref.current?.setState(this.unsavedState);
                 this.unsavedState = null;
             }
+
+            if (onMobileYouTube) {
+                if (this.mutationObserver) {
+                    this.mutationObserver.disconnect();
+                }
+                
+                this.mutationObserver = new MutationObserver(() => this.attachToPage(onMobileYouTube));
+
+                this.mutationObserver.observe(referenceNode, { 
+                    childList: true,
+                    subtree: true
+                });
+            }
         }
     }