From bb4ec488d37122f7ea2856979ed292948379a6be Mon Sep 17 00:00:00 2001
From: Bongbong <tsmaka1998@gmail.com>
Date: Fri, 15 May 2026 12:17:57 +0300
Subject: [PATCH] feat(audio-compressor): add Auto Track Gain to boost quiet
 tracks

Reads YouTube's per-track loudnessDb metadata and applies compensating
gain to tracks quieter than the reference level.

Closes #3032
---
 src/i18n/resources/en.json      |   6 +
 src/plugins/audio-compressor.ts | 325 +++++++++++++++++++++++++-------
 2 files changed, 260 insertions(+), 71 deletions(-)

diff --git a/src/i18n/resources/en.json b/src/i18n/resources/en.json
index 763e854b82..e3b04e0297 100644
--- a/src/i18n/resources/en.json
+++ b/src/i18n/resources/en.json
@@ -356,6 +356,12 @@
     },
     "audio-compressor": {
       "description": "Apply compression to audio (lowers the volume of the loudest parts of the signal and raises the volume of the softest parts)",
+      "menu": {
+        "auto-track-gain": "Auto track gain (boost quiet tracks)",
+        "maximum-gain": {
+          "label": "Maximum gain"
+        }
+      },
       "name": "Audio Compressor"
     },
     "auth-proxy-adapter": {
diff --git a/src/plugins/audio-compressor.ts b/src/plugins/audio-compressor.ts
index 22cf1bc54c..2e0ece1453 100644
--- a/src/plugins/audio-compressor.ts
+++ b/src/plugins/audio-compressor.ts
@@ -2,6 +2,19 @@ import { createPlugin } from '@/utils';
 import { t } from '@/i18n';
 import { type MusicPlayer } from '@/types/music-player';
 
+import type { MenuContext } from '@/types/contexts';
+import type { MenuTemplate } from '@/menu';
+
+export type AudioCompressorPluginConfig = {
+  enabled: boolean;
+  autoTrackGain: boolean;
+  maxTrackGainDb: number;
+};
+
+const MAX_TRACK_GAIN_CHOICES = [6, 9, 12, 15, 18, 24] as const;
+
+const dbToLinear = (db: number) => Math.pow(10, db / 20);
+
 const lazySafeTry = (...fns: (() => void)[]) => {
   for (const fn of fns) {
     try {
@@ -10,93 +23,205 @@ const lazySafeTry = (...fns: (() => void)[]) => {
   }
 };
 
-const createCompressorNode = (
-  audioContext: AudioContext,
-): DynamicsCompressorNode => {
-  const compressor = audioContext.createDynamicsCompressor();
-
+const configureCompressor = (compressor: DynamicsCompressorNode) => {
   compressor.threshold.value = -50;
   compressor.ratio.value = 12;
   compressor.knee.value = 40;
   compressor.attack.value = 0;
   compressor.release.value = 0.25;
-
-  return compressor;
 };
 
-class Storage {
-  lastSource: MediaElementAudioSourceNode | null = null;
-  lastContext: AudioContext | null = null;
-  lastCompressor: DynamicsCompressorNode | null = null;
+class Chain {
+  source: MediaElementAudioSourceNode | null = null;
+  context: AudioContext | null = null;
+  compressor: DynamicsCompressorNode | null = null;
+  trackGain: GainNode | null = null;
+
+  build(source: MediaElementAudioSourceNode, context: AudioContext) {
+    if (
+      this.source === source &&
+      this.context === context &&
+      this.compressor
+    ) {
+      return; // already built
+    }
+
+    this.teardown();
+
+    this.source = source;
+    this.context = context;
+
+    const compressor = context.createDynamicsCompressor();
+    const trackGain = context.createGain();
+    configureCompressor(compressor);
+    trackGain.gain.value = 1;
 
-  connected: WeakMap<MediaElementAudioSourceNode, DynamicsCompressorNode> =
-    new WeakMap();
+    this.compressor = compressor;
+    this.trackGain = trackGain;
 
-  connectToCompressor = (
-    source: MediaElementAudioSourceNode | null = null,
-    audioContext: AudioContext | null = null,
-    compressor: DynamicsCompressorNode | null = null,
-  ): boolean => {
-    if (!(source && audioContext && compressor)) return false;
+    // Source was previously connected directly to destination by the
+    // renderer; detach that and route through our chain instead.
+    lazySafeTry(() => source.disconnect(context.destination));
 
-    const current = this.connected.get(source);
-    if (current === compressor) return false;
+    source.connect(compressor);
+    compressor.connect(trackGain);
+    trackGain.connect(context.destination);
+  }
 
-    this.lastSource = source;
-    this.lastContext = audioContext;
-    this.lastCompressor = compressor;
+  applyTrackGain(gainDb: number) {
+    if (!this.context || !this.trackGain) return;
+    this.trackGain.gain.linearRampToValueAtTime(
+      dbToLinear(gainDb),
+      this.context.currentTime + 0.1,
+    );
+  }
 
-    if (current) {
+  teardown() {
+    const { source, context, compressor } = this;
+    if (source && context && compressor) {
       lazySafeTry(
-        () => source.disconnect(current),
-        () => current.disconnect(audioContext.destination),
+        () => source.disconnect(compressor),
+        () => source.connect(context.destination),
       );
-    } else {
-      lazySafeTry(() => source.disconnect(audioContext.destination));
     }
+    lazySafeTry(
+      () => this.compressor?.disconnect(),
+      () => this.trackGain?.disconnect(),
+    );
+    this.compressor = null;
+    this.trackGain = null;
+    // Keep source/context refs so a re-enable can rebuild without waiting
+    // for the next audio-can-play event.
+  }
+}
 
-    try {
-      source.connect(compressor);
-      compressor.connect(audioContext.destination);
-      this.connected.set(source, compressor);
-      return true;
-    } catch (error) {
-      console.error('connectToCompressor failed', error);
-      return false;
+const chain = new Chain();
+
+let currentConfig: AudioCompressorPluginConfig = {
+  enabled: false,
+  autoTrackGain: false,
+  maxTrackGainDb: 12,
+};
+
+const getContentLoudnessDb = (): number | null => {
+  try {
+    const player = document.querySelector('#movie_player') as
+      | (Element & { getPlayerResponse?: () => unknown })
+      | null;
+    const response = player?.getPlayerResponse?.() as
+      | {
+          playerConfig?: {
+            audioConfig?: {
+              loudnessDb?: number;
+              perceptualLoudnessDb?: number;
+            };
+          };
+        }
+      | undefined;
+    const loudnessDb =
+      response?.playerConfig?.audioConfig?.loudnessDb ??
+      response?.playerConfig?.audioConfig?.perceptualLoudnessDb;
+    return typeof loudnessDb === 'number' ? loudnessDb : null;
+  } catch {
+    return null;
+  }
+};
+
+let pendingRetry: ReturnType<typeof setTimeout> | null = null;
+
+const cancelPendingRetry = () => {
+  if (pendingRetry !== null) {
+    clearTimeout(pendingRetry);
+    pendingRetry = null;
+  }
+};
+
+const updateTrackGain = (retriesLeft = 4) => {
+  cancelPendingRetry();
+
+  if (!currentConfig.autoTrackGain) {
+    chain.applyTrackGain(0);
+    return;
+  }
+
+  const loudnessDb = getContentLoudnessDb();
+  if (loudnessDb === null) {
+    if (retriesLeft > 0) {
+      // YT may not have populated loudness yet — retry shortly.
+      pendingRetry = setTimeout(() => updateTrackGain(retriesLeft - 1), 400);
+    } else {
+      chain.applyTrackGain(0);
     }
-  };
+    return;
+  }
 
-  disconnectCompressor = (): boolean => {
-    const source = this.lastSource;
-    const audioContext = this.lastContext;
-    if (!(source && audioContext)) return false;
-    const current = this.connected.get(source);
-    if (!current) return false;
+  // YT's loudnessDb is signed: positive = louder than reference, negative =
+  // quieter. Compensate quiet tracks; leave loud tracks alone.
+  const compensation = loudnessDb < 0 ? -loudnessDb : 0;
+  const target = Math.min(compensation, currentConfig.maxTrackGainDb);
+  chain.applyTrackGain(target);
+};
 
-    lazySafeTry(
-      () => source.connect(audioContext.destination),
-      () => source.disconnect(current),
-      () => current.disconnect(audioContext.destination),
-    );
-    this.connected.delete(source);
-    return true;
-  };
-}
+const sourceMediaElement = (
+  source: MediaElementAudioSourceNode | null,
+): HTMLVideoElement | null =>
+  ((source as
+    | (MediaElementAudioSourceNode & { mediaElement?: HTMLMediaElement })
+    | null
+  )?.mediaElement as HTMLVideoElement | undefined) ?? null;
+
+let videoSwapObserver: MutationObserver | null = null;
+
+const stopWatchingForVideoSwap = () => {
+  if (videoSwapObserver) {
+    videoSwapObserver.disconnect();
+    videoSwapObserver = null;
+  }
+};
 
-const storage = new Storage();
+const handleVideoSwap = (newEl: HTMLVideoElement) => {
+  if (!chain.context) return;
+  let newSource: MediaElementAudioSourceNode;
+  try {
+    newSource = chain.context.createMediaElementSource(newEl);
+  } catch {
+    return;
+  }
+  chain.build(newSource, chain.context);
+  updateTrackGain();
+  watchForVideoSwap();
+};
+
+const watchForVideoSwap = () => {
+  stopWatchingForVideoSwap();
+  const currentEl = sourceMediaElement(chain.source);
+  if (!currentEl) return;
+  // YT swaps the <video> element on certain seek operations (notably seek
+  // near end then back to start). Our source is permanently bound to the old
+  // element; the new one would play straight to the OS, bypassing the chain.
+  // Detect the swap and rebind.
+  const target =
+    (document.querySelector('#movie_player') as HTMLElement | null) ??
+    document.body;
+  videoSwapObserver = new MutationObserver(() => {
+    const newEl = document.querySelector<HTMLVideoElement>('video');
+    if (!newEl || newEl === currentEl) return;
+    handleVideoSwap(newEl);
+  });
+  videoSwapObserver.observe(target, { childList: true, subtree: true });
+};
 
 const audioCanPlayHandler = ({
   detail: { audioSource, audioContext },
 }: CustomEvent<Compressor>) => {
-  storage.connectToCompressor(
-    audioSource,
-    audioContext,
-    createCompressorNode(audioContext),
-  );
+  cancelPendingRetry();
+  chain.build(audioSource, audioContext);
+  updateTrackGain();
+  watchForVideoSwap();
 };
 
 const ensureAudioContextLoad = (playerApi: MusicPlayer) => {
-  if (playerApi.getPlayerState() !== 1 || storage.lastContext) return;
+  if (playerApi.getPlayerState() !== 1 || chain.context) return;
 
   playerApi.loadVideoById(
     playerApi.getPlayerResponse().videoDetails.videoId,
@@ -108,26 +233,84 @@ const ensureAudioContextLoad = (playerApi: MusicPlayer) => {
 export default createPlugin({
   name: () => t('plugins.audio-compressor.name'),
   description: () => t('plugins.audio-compressor.description'),
+  restartNeeded: false,
+  config: {
+    enabled: false,
+    autoTrackGain: false,
+    maxTrackGainDb: 12,
+  } as AudioCompressorPluginConfig,
+
+  menu: async ({
+    getConfig,
+    setConfig,
+  }: MenuContext<AudioCompressorPluginConfig>): Promise<MenuTemplate> => {
+    const config = await getConfig();
+
+    return [
+      {
+        label: t('plugins.audio-compressor.menu.auto-track-gain'),
+        type: 'checkbox',
+        checked: config.autoTrackGain,
+        click(item) {
+          setConfig({ autoTrackGain: item.checked });
+        },
+      },
+      {
+        label: t('plugins.audio-compressor.menu.maximum-gain.label'),
+        type: 'submenu',
+        submenu: MAX_TRACK_GAIN_CHOICES.map((db) => ({
+          label: `${db} dB`,
+          type: 'radio' as const,
+          checked: config.maxTrackGainDb === db,
+          click() {
+            setConfig({ maxTrackGainDb: db });
+          },
+        })),
+      },
+    ];
+  },
 
   renderer: {
+    async start({ getConfig }) {
+      // Register synchronously so we never miss an event during the await.
+      document.addEventListener('peard:audio-can-play', audioCanPlayHandler, {
+        passive: true,
+      });
+      currentConfig = await getConfig();
+      // If the chain was previously built (plugin re-enable), rebuild now
+      // rather than waiting for the next track change.
+      if (chain.source && chain.context) {
+        // YT may have swapped the <video> element while we were disabled
+        // (observer was off). Detect a stale cached source and rebind to
+        // the live element rather than rebuilding onto a dead one.
+        const currentVideo = document.querySelector<HTMLVideoElement>('video');
+        if (currentVideo && sourceMediaElement(chain.source) !== currentVideo) {
+          handleVideoSwap(currentVideo);
+        } else {
+          chain.build(chain.source, chain.context);
+          updateTrackGain();
+          watchForVideoSwap();
+        }
+      }
+    },
+
     onPlayerApiReady(playerApi) {
       ensureAudioContextLoad(playerApi);
     },
 
-    start() {
-      document.addEventListener('peard:audio-can-play', audioCanPlayHandler, {
-        passive: true,
-      });
-      storage.connectToCompressor(
-        storage.lastSource,
-        storage.lastContext,
-        storage.lastCompressor,
-      );
+    onConfigChange(newConfig: AudioCompressorPluginConfig) {
+      currentConfig = newConfig;
+      updateTrackGain();
     },
 
     stop() {
-      document.removeEventListener('peard:audio-can-play', audioCanPlayHandler);
-      storage.disconnectCompressor();
+      document.removeEventListener(
+        'peard:audio-can-play',
+        audioCanPlayHandler,
+      );
+      stopWatchingForVideoSwap();
+      cancelPendingRetry();
+      chain.teardown();
     },
   },
 });