From 60c49f18389ff625177a57277ef8f292a31097bf Mon Sep 17 00:00:00 2001
From: Armand Philippot <git@armandphilippot.com>
Date: Sat, 28 Oct 2023 17:12:58 +0200
Subject: 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
---
 .../prism-theme-toggle/prism-theme-toggle.tsx      |  27 +---
 src/pages/_app.tsx                                 |   7 +-
 src/pages/_document.tsx                            |  17 ++-
 src/types/app.ts                                   |   3 +
 src/utils/constants.ts                             |   5 +
 src/utils/helpers/themes.ts                        |   7 +
 src/utils/hooks/index.ts                           |   4 +-
 src/utils/hooks/use-add-classname.tsx              |  32 -----
 src/utils/hooks/use-attributes.tsx                 |  50 -------
 src/utils/hooks/use-prism-theme/index.ts           |   1 +
 .../hooks/use-prism-theme/use-prism-theme.test.tsx | 134 +++++++++++++++++
 src/utils/hooks/use-prism-theme/use-prism-theme.ts |  54 +++++++
 src/utils/hooks/use-query-selector-all.tsx         |  22 ---
 src/utils/providers/index.ts                       |   2 +-
 src/utils/providers/prism-theme-provider/index.ts  |   1 +
 .../prism-theme-provider.test.tsx                  |  50 +++++++
 .../prism-theme-provider/prism-theme-provider.tsx  |  86 +++++++++++
 src/utils/providers/prism-theme.tsx                | 158 ---------------------
 .../providers/theme-provider/theme-provider.tsx    |  17 +--
 19 files changed, 376 insertions(+), 301 deletions(-)
 delete mode 100644 src/utils/hooks/use-add-classname.tsx
 delete mode 100644 src/utils/hooks/use-attributes.tsx
 create mode 100644 src/utils/hooks/use-prism-theme/index.ts
 create mode 100644 src/utils/hooks/use-prism-theme/use-prism-theme.test.tsx
 create mode 100644 src/utils/hooks/use-prism-theme/use-prism-theme.ts
 delete mode 100644 src/utils/hooks/use-query-selector-all.tsx
 create mode 100644 src/utils/providers/prism-theme-provider/index.ts
 create mode 100644 src/utils/providers/prism-theme-provider/prism-theme-provider.test.tsx
 create mode 100644 src/utils/providers/prism-theme-provider/prism-theme-provider.tsx
 delete mode 100644 src/utils/providers/prism-theme.tsx

(limited to 'src')

diff --git a/src/components/organisms/forms/prism-theme-toggle/prism-theme-toggle.tsx b/src/components/organisms/forms/prism-theme-toggle/prism-theme-toggle.tsx
index 2b0a179..1eba191 100644
--- a/src/components/organisms/forms/prism-theme-toggle/prism-theme-toggle.tsx
+++ b/src/components/organisms/forms/prism-theme-toggle/prism-theme-toggle.tsx
@@ -1,6 +1,6 @@
-import { useCallback, type ChangeEvent, type FC } from 'react';
+import type { FC } from 'react';
 import { useIntl } from 'react-intl';
-import { type PrismTheme, usePrismTheme } from '../../../../utils/providers';
+import { usePrismTheme } from '../../../../utils/hooks';
 import { Icon, Legend } from '../../../atoms';
 import {
   Switch,
@@ -20,24 +20,7 @@ export type PrismThemeToggleProps = Omit<
  */
 export const PrismThemeToggle: FC<PrismThemeToggleProps> = (props) => {
   const intl = useIntl();
-  const { theme, setTheme, resolvedTheme } = usePrismTheme();
-
-  /**
-   * Check if the resolved or chosen theme is dark theme.
-   *
-   * @returns {boolean} True if it is dark theme.
-   */
-  const isDarkTheme = (prismTheme?: PrismTheme): boolean => {
-    if (prismTheme === 'system') return resolvedTheme === 'dark';
-    return prismTheme === 'dark';
-  };
-
-  const updateTheme = useCallback(
-    (e: ChangeEvent<HTMLInputElement>) => {
-      setTheme(e.target.value === 'light' ? 'light' : 'dark');
-    },
-    [setTheme]
-  );
+  const { currentTheme, toggleTheme } = usePrismTheme();
 
   const themeLabel = intl.formatMessage({
     defaultMessage: 'Code blocks:',
@@ -85,8 +68,8 @@ export const PrismThemeToggle: FC<PrismThemeToggleProps> = (props) => {
       items={options}
       legend={<Legend>{themeLabel}</Legend>}
       name="code-blocks"
-      onSwitch={updateTheme}
-      value={isDarkTheme(theme) ? 'dark' : 'light'}
+      onSwitch={toggleTheme}
+      value={currentTheme}
     />
   );
 };
diff --git a/src/pages/_app.tsx b/src/pages/_app.tsx
index 0c92c93..caf4a96 100644
--- a/src/pages/_app.tsx
+++ b/src/pages/_app.tsx
@@ -3,7 +3,7 @@ import { IntlProvider } from 'react-intl';
 import '../styles/globals.scss';
 import type { AppPropsWithLayout } from '../types';
 import { settings } from '../utils/config';
-import { STORAGE_KEY } from '../utils/constants';
+import { PRISM_THEME_ATTRIBUTE, STORAGE_KEY } from '../utils/constants';
 import {
   AckeeProvider,
   MotionProvider,
@@ -37,7 +37,10 @@ const App = ({ Component, pageProps }: AppPropsWithLayout) => {
             attribute={STORAGE_KEY.THEME}
             storageKey={STORAGE_KEY.THEME}
           >
-            <PrismThemeProvider>
+            <PrismThemeProvider
+              attribute={PRISM_THEME_ATTRIBUTE}
+              storageKey={STORAGE_KEY.PRISM}
+            >
               {getLayout(<Component {...componentProps} />, {})}
             </PrismThemeProvider>
           </ThemeProvider>
diff --git a/src/pages/_document.tsx b/src/pages/_document.tsx
index 7e3b5e6..6d065cd 100644
--- a/src/pages/_document.tsx
+++ b/src/pages/_document.tsx
@@ -1,6 +1,12 @@
 import { Html, Head, Main, NextScript } from 'next/document';
 import Script from 'next/script';
-import { STORAGE_KEY } from '../utils/constants';
+import {
+  PRISM_THEME_ATTRIBUTE,
+  STORAGE_KEY,
+  VALID_THEMES,
+} from '../utils/constants';
+
+const validPrismThemesStr = VALID_THEMES.map((t) => `"${t}"`);
 
 // eslint-disable-next-line @typescript-eslint/no-shadow -- Required by NextJs
 export default function Document() {
@@ -25,6 +31,15 @@ export default function Document() {
           // eslint-disable-next-line react/jsx-no-literals
           strategy="beforeInteractive"
         />
+        <Script
+          dangerouslySetInnerHTML={{
+            __html: `!function(){const t=localStorage.getItem("${STORAGE_KEY.PRISM}"),e="string"==typeof t?JSON.parse(t):void 0,s=e&&[${validPrismThemesStr}].includes(e)?e:"system";document.documentElement.setAttribute("${PRISM_THEME_ATTRIBUTE}",s)}();`,
+          }}
+          // eslint-disable-next-line react/jsx-no-literals
+          id="prism-theme-hydration"
+          // eslint-disable-next-line react/jsx-no-literals
+          strategy="beforeInteractive"
+        />
       </Head>
       <body>
         <Main />
diff --git a/src/types/app.ts b/src/types/app.ts
index 565fe97..2588a7b 100644
--- a/src/types/app.ts
+++ b/src/types/app.ts
@@ -2,6 +2,7 @@ import type { NextPage } from 'next';
 import type { AppProps as NextAppProps } from 'next/app';
 import type { ReactElement, ReactNode } from 'react';
 import type { MessageFormatElement } from 'react-intl';
+import type { VALID_THEMES } from '../utils/constants';
 
 export type NextPageWithLayoutOptions = {
   withExtraPadding?: boolean;
@@ -139,3 +140,5 @@ export type Spacing = '2xs' | 'xs' | 'sm' | 'md' | 'lg' | 'xl' | '2xl' | '3xl';
 export type Validator<T> = (value: unknown) => value is T;
 
 export type AckeeTrackerValue = 'full' | 'partial';
+
+export type Theme = (typeof VALID_THEMES)[number];
diff --git a/src/utils/constants.ts b/src/utils/constants.ts
index 0c64f3a..7129624 100644
--- a/src/utils/constants.ts
+++ b/src/utils/constants.ts
@@ -34,5 +34,10 @@ export const ROUTES = {
 export const STORAGE_KEY = {
   ACKEE: 'ackee-tracking',
   MOTION: 'reduced-motion',
+  PRISM: 'prismjs-color-scheme',
   THEME: 'theme',
 } as const;
+
+export const PRISM_THEME_ATTRIBUTE = 'data-prismjs-color-scheme-current';
+
+export const VALID_THEMES = ['dark', 'light', 'system'] as const;
diff --git a/src/utils/helpers/themes.ts b/src/utils/helpers/themes.ts
index 8ef3a19..66c0652 100644
--- a/src/utils/helpers/themes.ts
+++ b/src/utils/helpers/themes.ts
@@ -1,3 +1,6 @@
+import type { Theme } from '../../types';
+import { VALID_THEMES } from '../constants';
+
 /**
  * Check if the user prefers dark color scheme.
  *
@@ -16,3 +19,7 @@ export const getThemeFromSystem = () => {
   if (prefersDarkScheme()) return 'dark';
   return 'light';
 };
+
+export const themeValidator = (value: unknown): value is Theme =>
+  typeof value === 'string' &&
+  (VALID_THEMES as readonly string[]).includes(value);
diff --git a/src/utils/hooks/index.ts b/src/utils/hooks/index.ts
index 4372ca0..f1bb31e 100644
--- a/src/utils/hooks/index.ts
+++ b/src/utils/hooks/index.ts
@@ -1,7 +1,5 @@
 export * from './use-ackee';
-export * from './use-add-classname';
 export * from './use-article';
-export * from './use-attributes';
 export * from './use-breadcrumb';
 export * from './use-comments';
 export * from './use-data-from-api';
@@ -15,7 +13,7 @@ export * from './use-mutation-observer';
 export * from './use-on-click-outside';
 export * from './use-pagination';
 export * from './use-prism';
-export * from './use-query-selector-all';
+export * from './use-prism-theme';
 export * from './use-reading-time';
 export * from './use-redirection';
 export * from './use-reduced-motion';
diff --git a/src/utils/hooks/use-add-classname.tsx b/src/utils/hooks/use-add-classname.tsx
deleted file mode 100644
index 8b0f6d6..0000000
--- a/src/utils/hooks/use-add-classname.tsx
+++ /dev/null
@@ -1,32 +0,0 @@
-import { useCallback, useEffect } from 'react';
-
-export type UseAddClassNameProps = {
-  className: string;
-  element?: HTMLElement;
-  elements?: NodeListOf<HTMLElement> | HTMLElement[];
-};
-
-/**
- * Add className to the given element(s).
- *
- * @param {UseAddClassNameProps} props - An object with classnames and one or more elements.
- */
-export const useAddClassName = ({
-  className,
-  element,
-  elements,
-}: UseAddClassNameProps) => {
-  const classNames = className.split(' ').filter((string) => string !== '');
-
-  const setClassName = useCallback(
-    (el: HTMLElement) => {
-      el.classList.add(...classNames);
-    },
-    [classNames]
-  );
-
-  useEffect(() => {
-    if (element) setClassName(element);
-    if (elements && elements.length > 0) elements.forEach(setClassName);
-  }, [element, elements, setClassName]);
-};
diff --git a/src/utils/hooks/use-attributes.tsx b/src/utils/hooks/use-attributes.tsx
deleted file mode 100644
index 20e9947..0000000
--- a/src/utils/hooks/use-attributes.tsx
+++ /dev/null
@@ -1,50 +0,0 @@
-import { useCallback, useEffect } from 'react';
-import { fromKebabCaseToCamelCase } from '../helpers';
-
-export type useAttributesProps = {
-  /**
-   * An HTML element.
-   */
-  element?: HTMLElement;
-  /**
-   * A node list of HTML Element.
-   */
-  elements?: NodeListOf<HTMLElement> | HTMLElement[];
-  /**
-   * The attribute name.
-   */
-  attribute: string;
-  /**
-   * The attribute value.
-   */
-  value: string;
-};
-
-/**
- * Set HTML attributes to the given element or to the HTML document.
- *
- * @param props - An object with element, attribute name and value.
- */
-export const useAttributes = ({
-  element,
-  elements,
-  attribute,
-  value,
-}: useAttributesProps) => {
-  const setAttribute = useCallback(
-    (el: HTMLElement) => {
-      if (attribute.startsWith('data')) {
-        el.setAttribute(attribute, value);
-      } else {
-        const camelCaseAttribute = fromKebabCaseToCamelCase(attribute);
-        el.dataset[camelCaseAttribute] = value;
-      }
-    },
-    [attribute, value]
-  );
-
-  useEffect(() => {
-    if (element) setAttribute(element);
-    if (elements && elements.length > 0) elements.forEach(setAttribute);
-  }, [element, elements, setAttribute]);
-};
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 };
+};
diff --git a/src/utils/hooks/use-query-selector-all.tsx b/src/utils/hooks/use-query-selector-all.tsx
deleted file mode 100644
index a3650ea..0000000
--- a/src/utils/hooks/use-query-selector-all.tsx
+++ /dev/null
@@ -1,22 +0,0 @@
-import { useRouter } from 'next/router';
-import { useEffect, useState } from 'react';
-
-/**
- * Use `document.querySelectorAll`.
- *
- * @param {string} query - A query.
- * @returns {NodeListOf<HTMLElementTagNameMap[T]|undefined>} - The node list.
- */
-export const useQuerySelectorAll = <T extends keyof HTMLElementTagNameMap>(
-  query: string
-): NodeListOf<HTMLElementTagNameMap[T]> | undefined => {
-  const [elements, setElements] =
-    useState<NodeListOf<HTMLElementTagNameMap[T]>>();
-  const { asPath } = useRouter();
-
-  useEffect(() => {
-    setElements(document.querySelectorAll(query));
-  }, [asPath, query]);
-
-  return elements;
-};
diff --git a/src/utils/providers/index.ts b/src/utils/providers/index.ts
index 90dca15..bc521b3 100644
--- a/src/utils/providers/index.ts
+++ b/src/utils/providers/index.ts
@@ -1,4 +1,4 @@
 export * from './ackee-provider';
 export * from './motion-provider';
-export * from './prism-theme';
+export * from './prism-theme-provider';
 export * from './theme-provider';
diff --git a/src/utils/providers/prism-theme-provider/index.ts b/src/utils/providers/prism-theme-provider/index.ts
new file mode 100644
index 0000000..fec2444
--- /dev/null
+++ b/src/utils/providers/prism-theme-provider/index.ts
@@ -0,0 +1 @@
+export * from './prism-theme-provider';
diff --git a/src/utils/providers/prism-theme-provider/prism-theme-provider.test.tsx b/src/utils/providers/prism-theme-provider/prism-theme-provider.test.tsx
new file mode 100644
index 0000000..d0c2630
--- /dev/null
+++ b/src/utils/providers/prism-theme-provider/prism-theme-provider.test.tsx
@@ -0,0 +1,50 @@
+import { describe, expect, it } from '@jest/globals';
+import { render, screen as rtlScreen } from '@testing-library/react';
+import { type FC, useContext } from 'react';
+import { PrismThemeContext, PrismThemeProvider } from './prism-theme-provider';
+
+const bodyPrefix = 'Current Prism theme is:';
+
+const ComponentTest: FC = () => {
+  const { theme } = useContext(PrismThemeContext);
+
+  return (
+    <div>
+      {bodyPrefix} {theme}
+    </div>
+  );
+};
+
+describe('PrismThemeProvider', () => {
+  it('uses the default value when the provider is not used', () => {
+    const defaultValue = 'system';
+
+    render(<ComponentTest />);
+
+    expect(rtlScreen.getByText(new RegExp(bodyPrefix))).toHaveTextContent(
+      `${bodyPrefix} ${defaultValue}`
+    );
+  });
+
+  it('provides the given value to its children and set a matching attribute', () => {
+    const attribute = 'ratione';
+    const theme = 'dark';
+
+    const { baseElement } = render(
+      <PrismThemeProvider
+        attribute={attribute}
+        defaultTheme={theme}
+        storageKey="qui"
+      >
+        <ComponentTest />
+      </PrismThemeProvider>
+    );
+
+    expect(rtlScreen.getByText(new RegExp(bodyPrefix))).toHaveTextContent(
+      `${bodyPrefix} ${theme}`
+    );
+    expect(baseElement.parentElement?.getAttribute(`data-${attribute}`)).toBe(
+      theme
+    );
+  });
+});
diff --git a/src/utils/providers/prism-theme-provider/prism-theme-provider.tsx b/src/utils/providers/prism-theme-provider/prism-theme-provider.tsx
new file mode 100644
index 0000000..0f0a4a1
--- /dev/null
+++ b/src/utils/providers/prism-theme-provider/prism-theme-provider.tsx
@@ -0,0 +1,86 @@
+import {
+  type Dispatch,
+  type FC,
+  type ReactNode,
+  type SetStateAction,
+  createContext,
+  useEffect,
+  useMemo,
+} from 'react';
+import type { Theme } from '../../../types';
+import {
+  getDataAttributeFrom,
+  getThemeFromSystem,
+  themeValidator,
+} from '../../helpers';
+import { useLocalStorage, useSystemColorScheme } from '../../hooks';
+
+export type PrismThemeContextProps = {
+  attribute: string;
+  resolvedTheme: Exclude<Theme, 'system'>;
+  setTheme: Dispatch<SetStateAction<Theme>>;
+  theme: Theme;
+};
+
+export const PrismThemeContext = createContext<PrismThemeContextProps>({
+  attribute: 'data-prism-theme',
+  resolvedTheme: getThemeFromSystem(),
+  setTheme: (value) => value,
+  theme: 'system',
+});
+
+export type PrismThemeProviderProps = {
+  /**
+   * The attribute name to append to document root.
+   */
+  attribute: string;
+  /**
+   * The provider children.
+   */
+  children?: ReactNode;
+  /**
+   * The default theme.
+   *
+   * @default 'system'
+   */
+  defaultTheme?: Theme;
+  /**
+   * The key to use in local storage.
+   */
+  storageKey: string;
+};
+
+export const PrismThemeProvider: FC<PrismThemeProviderProps> = ({
+  attribute,
+  children,
+  defaultTheme = 'system',
+  storageKey,
+}) => {
+  const [theme, setTheme] = useLocalStorage(
+    storageKey,
+    defaultTheme,
+    themeValidator
+  );
+  const resolvedTheme = useSystemColorScheme();
+  const dataAttribute = getDataAttributeFrom(attribute);
+
+  useEffect(() => {
+    if (typeof window !== 'undefined') {
+      document.documentElement.setAttribute(dataAttribute, `${theme}`);
+    }
+
+    return () => {
+      document.documentElement.removeAttribute(dataAttribute);
+    };
+  }, [dataAttribute, theme]);
+
+  const value = useMemo(() => {
+    return { attribute: dataAttribute, resolvedTheme, setTheme, theme };
+  }, [dataAttribute, resolvedTheme, setTheme, theme]);
+
+  return (
+    <PrismThemeContext.Provider value={value}>
+      {children}
+    </PrismThemeContext.Provider>
+  );
+};
diff --git a/src/utils/providers/prism-theme.tsx b/src/utils/providers/prism-theme.tsx
deleted file mode 100644
index 603e8a7..0000000
--- a/src/utils/providers/prism-theme.tsx
+++ /dev/null
@@ -1,158 +0,0 @@
-/* eslint-disable max-statements */
-import {
-  createContext,
-  type FC,
-  type ReactNode,
-  useCallback,
-  useContext,
-  useEffect,
-  useState,
-  useMemo,
-} from 'react';
-import { useAttributes, useLocalStorage, useQuerySelectorAll } from '../hooks';
-
-export type PrismTheme = 'dark' | 'light' | 'system';
-export type ResolvedPrismTheme = Exclude<PrismTheme, 'system'>;
-
-export type UsePrismThemeProps = {
-  themes: PrismTheme[];
-  theme?: PrismTheme;
-  setTheme: (theme: PrismTheme) => void;
-  resolvedTheme?: ResolvedPrismTheme;
-  codeBlocks?: NodeListOf<HTMLPreElement>;
-};
-
-export type PrismThemeProviderProps = {
-  attribute?: string;
-  children: ReactNode;
-  storageKey?: string;
-  themes?: PrismTheme[];
-};
-
-export const PrismThemeContext = createContext<UsePrismThemeProps>({
-  themes: ['dark', 'light', 'system'],
-  setTheme: (_) => {
-    // This is intentional.
-  },
-});
-
-export const usePrismTheme = () => useContext(PrismThemeContext);
-
-/**
- * Check if user prefers dark color scheme.
- *
- * @returns {boolean|undefined} True if `prefers-color-scheme` is set to `dark`.
- */
-const prefersDarkScheme = (): boolean | undefined => {
-  if (typeof window === 'undefined') return undefined;
-
-  return window.matchMedia('(prefers-color-scheme: dark)').matches;
-};
-
-/**
- * Check if a given string is a Prism theme name.
- *
- * @param {string} theme - A string.
- * @returns {boolean} True if the given string match a Prism theme name.
- */
-const isValidTheme = (theme: string): boolean =>
-  theme === 'dark' || theme === 'light' || theme === 'system';
-
-const defaultThemes = ['dark', 'light', 'system'] satisfies PrismTheme[];
-
-const validator = (value: unknown): value is PrismTheme =>
-  typeof value === 'string' && (defaultThemes as string[]).includes(value);
-
-export const PrismThemeProvider: FC<PrismThemeProviderProps> = ({
-  attribute = 'data-prismjs-color-scheme-current',
-  storageKey = 'prismjs-color-scheme',
-  themes = defaultThemes,
-  children,
-}) => {
-  /**
-   * Retrieve the theme to use depending on `prefers-color-scheme`.
-   */
-  const getThemeFromSystem = useCallback(() => {
-    if (prefersDarkScheme()) return 'dark';
-    return 'light';
-  }, []);
-
-  const [prismTheme, setPrismTheme] = useLocalStorage<PrismTheme>(
-    storageKey,
-    'system',
-    validator
-  );
-
-  useEffect(() => {
-    if (!isValidTheme(prismTheme)) setPrismTheme('system');
-  }, [prismTheme, setPrismTheme]);
-
-  const [resolvedTheme, setResolvedTheme] = useState<ResolvedPrismTheme>();
-
-  useEffect(() => {
-    if (prismTheme === 'dark' || prismTheme === 'light') {
-      setResolvedTheme(prismTheme);
-    } else {
-      setResolvedTheme(getThemeFromSystem());
-    }
-  }, [prismTheme, getThemeFromSystem]);
-
-  const updateResolvedTheme = useCallback(() => {
-    setResolvedTheme(getThemeFromSystem());
-  }, [getThemeFromSystem]);
-
-  useEffect(() => {
-    window
-      .matchMedia('(prefers-color-scheme: dark)')
-      .addEventListener('change', updateResolvedTheme);
-
-    return () =>
-      window
-        .matchMedia('(prefers-color-scheme: dark)')
-        .removeEventListener('change', updateResolvedTheme);
-  }, [updateResolvedTheme]);
-
-  const preTags = useQuerySelectorAll<'pre'>('pre');
-  useAttributes({ elements: preTags, attribute, value: prismTheme });
-
-  /**
-   * Listen for changes on pre attributes and update theme.
-   */
-  const listenAttributeChange = useCallback(
-    (pre: HTMLPreElement) => {
-      const observer = new MutationObserver((mutations) => {
-        mutations.forEach((record) => {
-          const mutatedPre = record.target as HTMLPreElement;
-          const newTheme = mutatedPre.getAttribute(attribute) as PrismTheme;
-          setPrismTheme(newTheme);
-        });
-      });
-      observer.observe(pre, {
-        attributes: true,
-        attributeFilter: [attribute],
-      });
-    },
-    [attribute, setPrismTheme]
-  );
-
-  useEffect(() => {
-    if (!preTags) return;
-    preTags.forEach(listenAttributeChange);
-  }, [preTags, listenAttributeChange]);
-
-  const value = useMemo(() => {
-    return {
-      themes,
-      theme: prismTheme,
-      setTheme: setPrismTheme,
-      codeBlocks: preTags,
-      resolvedTheme,
-    };
-  }, [preTags, prismTheme, resolvedTheme, setPrismTheme, themes]);
-
-  return (
-    <PrismThemeContext.Provider value={value}>
-      {children}
-    </PrismThemeContext.Provider>
-  );
-};
diff --git a/src/utils/providers/theme-provider/theme-provider.tsx b/src/utils/providers/theme-provider/theme-provider.tsx
index b743a7f..f09f74d 100644
--- a/src/utils/providers/theme-provider/theme-provider.tsx
+++ b/src/utils/providers/theme-provider/theme-provider.tsx
@@ -7,13 +7,14 @@ import {
   useMemo,
   useEffect,
 } from 'react';
-import { getDataAttributeFrom, getThemeFromSystem } from '../../helpers';
+import type { Theme } from '../../../types';
+import {
+  getDataAttributeFrom,
+  getThemeFromSystem,
+  themeValidator,
+} from '../../helpers';
 import { useLocalStorage, useSystemColorScheme } from '../../hooks';
 
-const validThemes = ['dark', 'light', 'system'] as const;
-
-type Theme = (typeof validThemes)[number];
-
 type ThemeContextProps = {
   resolvedTheme: Exclude<Theme, 'system'>;
   setTheme: Dispatch<SetStateAction<Theme>>;
@@ -26,10 +27,6 @@ export const ThemeContext = createContext<ThemeContextProps>({
   theme: 'system',
 });
 
-const validator = (value: unknown): value is Theme =>
-  typeof value === 'string' &&
-  (validThemes as readonly string[]).includes(value);
-
 export type ThemeProviderProps = {
   /**
    * The attribute name to append to document root.
@@ -60,7 +57,7 @@ export const ThemeProvider: FC<ThemeProviderProps> = ({
   const [theme, setTheme] = useLocalStorage(
     storageKey,
     defaultTheme,
-    validator
+    themeValidator
   );
   const userColorScheme = useSystemColorScheme();
   const resolvedTheme: Exclude<Theme, 'system'> =
-- 
cgit v1.2.3