refactor(providers,hooks): rewrite PrismThemeProvider & usePrismTheme

* reuse Theme provider logic
* move DOM mutation from provider to hook
* add a script to init theme before page load

3 files changed, 189 insertions, 0 deletions

diff --git a/src/utils/hooks/use-prism-theme/index.ts b/src/utils/hooks/use-prism-theme/index.ts
new file mode 100644
index 0000000..5972519
--- /dev/null
+++ b/src/utils/hooks/use-prism-theme/index.ts
@@ -0,0 +1 @@
+export * from './use-prism-theme';
diff --git a/src/utils/hooks/use-prism-theme/use-prism-theme.test.tsx b/src/utils/hooks/use-prism-theme/use-prism-theme.test.tsx
new file mode 100644
index 0000000..8a178c9
--- /dev/null
+++ b/src/utils/hooks/use-prism-theme/use-prism-theme.test.tsx
@@ -0,0 +1,134 @@
+import { describe, expect, it } from '@jest/globals';
+import {
+  act,
+  render,
+  renderHook,
+  screen as rtlScreen,
+} from '@testing-library/react';
+import type { FC, ReactNode } from 'react';
+import {
+  PrismThemeProvider,
+  type PrismThemeProviderProps,
+} from '../../providers/prism-theme-provider';
+import { usePrismTheme } from './use-prism-theme';
+
+const codeSample1 = `const foo = 42;`;
+const codeSample2 = `const bar = "baz";`;
+const codeSample3 = `const baz = () => false;`;
+
+const ComponentTest: FC = () => {
+  usePrismTheme();
+
+  return (
+    <div>
+      <pre className="language-js">{codeSample1}</pre>
+      <pre className="language-js">{codeSample2}</pre>
+      <pre>{codeSample3}</pre>
+    </div>
+  );
+};
+
+const createWrapper = (
+  Wrapper: FC<PrismThemeProviderProps>,
+  config: PrismThemeProviderProps
+) =>
+  function CreatedWrapper({ children }: { children: ReactNode }) {
+    return <Wrapper {...config}>{children}</Wrapper>;
+  };
+
+describe('usePrismTheme', () => {
+  it('should return the default value without provider and prevent update', () => {
+    const defaultTheme = 'system';
+    const { result } = renderHook(() => usePrismTheme());
+
+    expect(result.current.theme).toBe(defaultTheme);
+
+    act(() => result.current.setTheme('dark'));
+
+    expect(result.current.theme).toBe(defaultTheme);
+  });
+
+  it('should add an attribute on pre elements when matching a Prism block', () => {
+    const defaultTheme = 'light';
+    const attribute = 'data-debitis';
+
+    render(
+      <PrismThemeProvider
+        attribute={attribute}
+        defaultTheme={defaultTheme}
+        storageKey="soluta"
+      >
+        <ComponentTest />
+      </PrismThemeProvider>
+    );
+
+    expect(rtlScreen.getByText(codeSample1)).toHaveAttribute(
+      attribute,
+      defaultTheme
+    );
+    expect(rtlScreen.getByText(codeSample2)).toHaveAttribute(
+      attribute,
+      defaultTheme
+    );
+    expect(rtlScreen.getByText(codeSample3)).not.toHaveAttribute(
+      attribute,
+      defaultTheme
+    );
+  });
+
+  it('can update the theme value using a setter', () => {
+    const defaultTheme = 'dark';
+
+    const { result } = renderHook(() => usePrismTheme(), {
+      wrapper: createWrapper(PrismThemeProvider, {
+        attribute: 'consequuntur',
+        defaultTheme,
+        storageKey: 'deleniti',
+      }),
+    });
+
+    expect(result.current.theme).toBe(defaultTheme);
+
+    const newTheme = 'light';
+
+    act(() => result.current.setTheme(newTheme));
+
+    expect(result.current.theme).toBe(newTheme);
+  });
+
+  it('can toggle the theme from light to dark', () => {
+    const defaultTheme = 'light';
+
+    const { result } = renderHook(() => usePrismTheme(), {
+      wrapper: createWrapper(PrismThemeProvider, {
+        attribute: 'et',
+        defaultTheme,
+        storageKey: 'velit',
+      }),
+    });
+
+    expect(result.current.theme).toBe(defaultTheme);
+
+    act(() => result.current.toggleTheme());
+
+    expect(result.current.theme).toBe('dark');
+  });
+
+  it('can toggle the theme from dark to light', () => {
+    const defaultTheme = 'dark';
+
+    const { result } = renderHook(() => usePrismTheme(), {
+      wrapper: createWrapper(PrismThemeProvider, {
+        attribute: 'et',
+        defaultTheme,
+        storageKey: 'velit',
+      }),
+    });
+
+    expect(result.current.theme).toBe(defaultTheme);
+
+    act(() => result.current.toggleTheme());
+
+    expect(result.current.theme).toBe('light');
+  });
+});
diff --git a/src/utils/hooks/use-prism-theme/use-prism-theme.ts b/src/utils/hooks/use-prism-theme/use-prism-theme.ts
new file mode 100644
index 0000000..4caec6c
--- /dev/null
+++ b/src/utils/hooks/use-prism-theme/use-prism-theme.ts
@@ -0,0 +1,54 @@
+import { useCallback, useContext, useEffect } from 'react';
+import { themeValidator as isValidTheme } from '../../helpers';
+import { PrismThemeContext } from '../../providers/prism-theme-provider';
+
+export const usePrismTheme = () => {
+  const { attribute, resolvedTheme, setTheme, theme } =
+    useContext(PrismThemeContext);
+  const currentTheme = theme === 'system' ? resolvedTheme : theme;
+
+  const handleMutations: MutationCallback = useCallback(
+    (mutations) => {
+      for (const mutation of mutations) {
+        if (mutation.target.nodeName.toLowerCase() !== 'pre') return;
+
+        const newTheme = (mutation.target as HTMLPreElement).getAttribute(
+          attribute
+        );
+
+        if (isValidTheme(newTheme) && newTheme !== theme) setTheme(newTheme);
+      }
+    },
+    [attribute, setTheme, theme]
+  );
+
+  useEffect(() => {
+    if (typeof window === 'undefined') return undefined;
+
+    const preElements = document.getElementsByTagName('pre');
+    const observer = new MutationObserver(handleMutations);
+
+    for (const preEl of Array.from(preElements)) {
+      if (preEl.className.includes('language-')) {
+        preEl.setAttribute(attribute, theme);
+        observer.observe(preEl, {
+          attributes: true,
+          attributeFilter: [attribute],
+        });
+      }
+    }
+
+    return () => {
+      observer.disconnect();
+    };
+  }, [attribute, handleMutations, theme]);
+
+  const toggleTheme = useCallback(() => {
+    setTheme(() => {
+      if (currentTheme === 'dark') return 'light';
+      return 'dark';
+    });
+  }, [currentTheme, setTheme]);
+
+  return { currentTheme, resolvedTheme, setTheme, theme, toggleTheme };
+};