diff --git a/packages/mdx/src/plugins/rehype/shiki-constants.ts b/packages/mdx/src/plugins/rehype/constants.ts
similarity index 87%
rename from packages/mdx/src/plugins/rehype/shiki-constants.ts
rename to packages/mdx/src/plugins/rehype/constants.ts
index b90ae22..c2774b4 100644
--- a/packages/mdx/src/plugins/rehype/shiki-constants.ts
+++ b/packages/mdx/src/plugins/rehype/constants.ts
@@ -6,17 +6,16 @@ import {
 } from '@shikijs/transformers';
 import type { ShikiTransformer } from '@shikijs/types';
 import { createCssVariablesTheme } from 'shiki/core';
-import type { BundledLanguage, ThemeRegistration } from 'shiki/types';
+import type { ThemeRegistration } from 'shiki/types';
 
-export const LINE_HIGHLIGHT_CLASS_NAME = 'line-highlight';
-export const LINE_FOCUS_CLASS_NAME = 'line-focus';
-export const LINE_DIFF_ADD_CLASS_NAME = 'line-diff line-add';
-export const LINE_DIFF_REMOVE_CLASS_NAME = 'line-diff line-remove';
+import { ShikiLang, ShikiTheme } from './types.js';
 
-export type ShikiLang = BundledLanguage | 'ansi' | 'text';
-export type ShikiTheme = (typeof SHIKI_THEMES)[number];
+const LINE_HIGHLIGHT_CLASS_NAME = 'line-highlight';
+const LINE_FOCUS_CLASS_NAME = 'line-focus';
+const LINE_DIFF_ADD_CLASS_NAME = 'line-diff line-add';
+const LINE_DIFF_REMOVE_CLASS_NAME = 'line-diff line-remove';
 
-export const SHIKI_CSS_THEME = createCssVariablesTheme({
+const SHIKI_CSS_THEME = createCssVariablesTheme({
   name: 'css-variables',
   variablePrefix: '--mint-',
   variableDefaults: {
@@ -68,23 +67,55 @@ export const SHIKI_CSS_THEME = createCssVariablesTheme({
   fontStyle: true,
 });
 
-export const DEFAULT_LANG = 'text' as const;
-export const DEFAULT_DARK_THEME: ShikiTheme = 'dark-plus' as const;
-export const DEFAULT_LIGHT_THEME: ShikiTheme = 'github-light-default' as const;
-export const DEFAULT_THEMES: [ShikiTheme, ShikiTheme, ThemeRegistration] = [
+const DEFAULT_LANG = 'text' as const;
+const DEFAULT_DARK_THEME: ShikiTheme = 'dark-plus' as const;
+const DEFAULT_LIGHT_THEME: ShikiTheme = 'github-light-default' as const;
+const DEFAULT_THEMES: [ShikiTheme, ShikiTheme, ThemeRegistration] = [
   DEFAULT_LIGHT_THEME,
   DEFAULT_DARK_THEME,
   SHIKI_CSS_THEME,
 ] as const;
 
-export const shikiColorReplacements: Partial<Record<ShikiTheme, string | Record<string, string>>> =
-  {
-    'dark-plus': {
-      '#1e1e1e': '#0B0C0E',
-    },
-  };
+const SHIKI_COLOR_REPLACEMENTS: Partial<Record<ShikiTheme, string | Record<string, string>>> = {
+  'dark-plus': {
+    '#1e1e1e': '#0B0C0E',
+  },
+};
+
+const DEFAULT_LANGS = [
+  'bash',
+  'blade',
+  'c',
+  'css',
+  'c#',
+  'c++',
+  'dart',
+  'diff',
+  'go',
+  'html',
+  'java',
+  'javascript',
+  'jsx',
+  'json',
+  'kotlin',
+  'log',
+  'lua',
+  'markdown',
+  'mdx',
+  'php',
+  'powershell',
+  'python',
+  'ruby',
+  'rust',
+  'solidity',
+  'swift',
+  'toml',
+  'typescript',
+  'tsx',
+  'yaml',
+];
 
-export const DEFAULT_LANG_ALIASES: Record<string, ShikiLang> = {
+const DEFAULT_LANG_ALIASES: Record<string, ShikiLang> = {
   ansi: 'ansi',
   abap: 'abap',
   'actionscript-3': 'actionscript-3',
@@ -400,9 +431,9 @@ export const DEFAULT_LANG_ALIASES: Record<string, ShikiLang> = {
   zig: 'zig',
 };
 
-export const UNIQUE_LANGS = Array.from(new Set(Object.values(DEFAULT_LANG_ALIASES)));
+const UNIQUE_LANGS = Array.from(new Set(Object.values(DEFAULT_LANG_ALIASES)));
 
-export const SHIKI_THEMES = [
+const SHIKI_THEMES = [
   'andromeeda',
   'aurora-x',
   'ayu-dark',
@@ -467,58 +498,39 @@ export const SHIKI_THEMES = [
   'css-variables', // for users who want to use custom CSS to style their code blocks
 ] as const;
 
-export const DEFAULT_LANGS = [
-  'bash',
-  'blade',
-  'c',
-  'css',
-  'c#',
-  'c++',
-  'dart',
-  'diff',
-  'go',
-  'html',
-  'java',
-  'javascript',
-  'jsx',
-  'json',
-  'kotlin',
-  'log',
-  'lua',
-  'markdown',
-  'mdx',
-  'php',
-  'powershell',
-  'python',
-  'ruby',
-  'rust',
-  'solidity',
-  'swift',
-  'toml',
-  'typescript',
-  'tsx',
-  'yaml',
-];
-
-export const matchAlgorithm = {
-  matchAlgorithm: 'v3',
-} as const;
-
-export const SHIKI_TRANSFORMERS: ShikiTransformer[] = [
+const SHIKI_TRANSFORMERS: ShikiTransformer[] = [
   transformerMetaHighlight({
     className: LINE_HIGHLIGHT_CLASS_NAME,
   }),
   transformerNotationHighlight({
-    ...matchAlgorithm,
+    matchAlgorithm: 'v3',
     classActiveLine: LINE_HIGHLIGHT_CLASS_NAME,
   }),
   transformerNotationFocus({
-    ...matchAlgorithm,
+    matchAlgorithm: 'v3',
     classActiveLine: LINE_FOCUS_CLASS_NAME,
   }),
   transformerNotationDiff({
-    ...matchAlgorithm,
+    matchAlgorithm: 'v3',
     classLineAdd: LINE_DIFF_ADD_CLASS_NAME,
     classLineRemove: LINE_DIFF_REMOVE_CLASS_NAME,
   }),
 ];
+
+export {
+  LINE_HIGHLIGHT_CLASS_NAME,
+  LINE_FOCUS_CLASS_NAME,
+  LINE_DIFF_ADD_CLASS_NAME,
+  LINE_DIFF_REMOVE_CLASS_NAME,
+  SHIKI_CSS_THEME,
+  DEFAULT_LANG,
+  DEFAULT_DARK_THEME,
+  DEFAULT_LIGHT_THEME,
+  DEFAULT_THEMES,
+  SHIKI_COLOR_REPLACEMENTS,
+  DEFAULT_LANG_ALIASES,
+  UNIQUE_LANGS,
+  SHIKI_THEMES,
+  DEFAULT_LANGS,
+  SHIKI_TRANSFORMERS,
+};
diff --git a/packages/mdx/src/plugins/rehype/core.ts b/packages/mdx/src/plugins/rehype/core.ts
new file mode 100644
index 0000000..2b342cf
--- /dev/null
+++ b/packages/mdx/src/plugins/rehype/core.ts
@@ -0,0 +1,53 @@
+import type { Root } from 'hast';
+import { getSingletonHighlighter } from 'shiki';
+import type { Plugin } from 'unified';
+import { visit } from 'unist-util-visit';
+
+import { DEFAULT_THEMES, DEFAULT_LANGS } from './constants.js';
+import { handlePreCode } from './handlers.js';
+import type { RehypeSyntaxHighlightingOptions } from './types.js';
+import { getLanguagesToLoad, getThemesToLoad } from './utils.js';
+
+const highlighterPromise = getSingletonHighlighter({
+  themes: DEFAULT_THEMES,
+  langs: DEFAULT_LANGS,
+});
+
+/** Rehype plugin with syntax highlighting, twoslash and etc. */
+const rehypeSyntaxHighlighting: Plugin<[RehypeSyntaxHighlightingOptions?], Root, Root> = (
+  options = {}
+) => {
+  return async (tree) => {
+    // Declare highlighter
+    const highlighter = await highlighterPromise;
+
+    // Parse options
+    const themesToLoad = getThemesToLoad(options);
+    const languagesToLoad = getLanguagesToLoad(options);
+
+    // Load themes and languages passed in options
+    await Promise.all([
+      ...themesToLoad.map((theme) => highlighter.loadTheme(theme)),
+      ...languagesToLoad.map((lang) => highlighter.loadLanguage(lang)),
+    ]);
+
+    const queue: Promise<void>[] = [];
+
+    // Visit HAST and apply highlight & twoslash
+    visit(tree, 'element', (node, index, parent) => {
+      const child = node.children[0];
+
+      // required for HAST node replacement
+      if (!parent || index === undefined) return;
+
+      // handle `<pre><code>...</code></pre>`
+      if (node.tagName === 'pre' && child?.type === 'element' && child.tagName === 'code') {
+        handlePreCode({ node, child, index, parent, highlighter, options, queue });
+      }
+    });
+
+    await Promise.all(queue);
+  };
+};
+
+export { rehypeSyntaxHighlighting };
diff --git a/packages/mdx/src/plugins/rehype/handlers.ts b/packages/mdx/src/plugins/rehype/handlers.ts
new file mode 100644
index 0000000..2f0f810
--- /dev/null
+++ b/packages/mdx/src/plugins/rehype/handlers.ts
@@ -0,0 +1,169 @@
+import { transformerTwoslash } from '@shikijs/twoslash';
+import type { Root, Element } from 'hast';
+import { toString } from 'hast-util-to-string';
+import { MdxJsxFlowElementHast, MdxJsxTextElementHast } from 'mdast-util-mdx-jsx';
+import { Highlighter } from 'shiki';
+
+import {
+  DEFAULT_DARK_THEME,
+  DEFAULT_LANG,
+  DEFAULT_LANG_ALIASES,
+  DEFAULT_LANGS,
+  DEFAULT_LIGHT_THEME,
+  SHIKI_TRANSFORMERS,
+  SHIKI_COLOR_REPLACEMENTS,
+  UNIQUE_LANGS,
+} from './constants.js';
+import { getTwoslashOptions, parseLineComment } from './twoslash/config.js';
+import { RehypeSyntaxHighlightingOptions, ShikiLang } from './types.js';
+import { getCustomLanguagesNames, getLanguage } from './utils.js';
+
+/** Handles `<pre><code>...</code></pre>` nodes */
+function handlePreCode({
+  node,
+  child,
+  index,
+  parent,
+  highlighter,
+  options,
+  queue,
+}: {
+  node: Element;
+  child: Element;
+  index: number;
+  parent: Element | Root | MdxJsxTextElementHast | MdxJsxFlowElementHast;
+  highlighter: Highlighter;
+  options: RehypeSyntaxHighlightingOptions;
+  queue: Promise<void>[];
+}) {
+  // Determine language
+  const preNodeLang = getLanguage(node, DEFAULT_LANG_ALIASES);
+  const codeNodeLang = getLanguage(child, DEFAULT_LANG_ALIASES);
+  const lang = preNodeLang ?? codeNodeLang ?? DEFAULT_LANG;
+
+  // Set metadata on code node copied from pre node
+  if (!Object.keys(node.properties).length) {
+    node.properties = child.properties;
+  }
+
+  if (!node.data) {
+    node.data = child.data;
+  }
+
+  // Traverse node
+  const customLanguageNames = getCustomLanguagesNames(options);
+  const isLanguageLoaded = DEFAULT_LANGS.includes(lang) || customLanguageNames.includes(lang);
+
+  if (isLanguageLoaded) {
+    traverseNode({ node, index, parent, highlighter, lang, options });
+  } else if (UNIQUE_LANGS.includes(lang)) {
+    queue.push(
+      highlighter.loadLanguage(lang).then(() => {
+        traverseNode({ node, index, parent, highlighter, lang, options });
+      })
+    );
+  } else {
+    traverseNode({ node, index, parent, highlighter, lang: DEFAULT_LANG, options });
+  }
+}
+
+/** Traverse node and apply highlighting and twoslash */
+function traverseNode({
+  node,
+  index,
+  parent,
+  highlighter,
+  lang,
+  options,
+}: {
+  node: Element;
+  index: number;
+  parent: Element | Root | MdxJsxTextElementHast | MdxJsxFlowElementHast;
+  highlighter: Highlighter;
+  lang: ShikiLang;
+  options: RehypeSyntaxHighlightingOptions;
+}) {
+  try {
+    let code = toString(node);
+
+    const meta = node.data?.meta?.split(' ') ?? [];
+    const twoslashIndex = meta.findIndex((str) => str.toLowerCase() === 'twoslash');
+    const shouldUseTwoslash = twoslashIndex > -1;
+
+    // Remove twoslash from meta list
+    if (node.data && node.data.meta && shouldUseTwoslash) {
+      meta.splice(twoslashIndex, 1);
+      node.data.meta = meta.join(' ').trim() || undefined;
+    }
+
+    // Set Link Map
+    const linkMap = options.linkMap ?? new Map();
+
+    if (shouldUseTwoslash) {
+      const codeLines = code.split('\n');
+
+      for (const [i, line] of codeLines.entries()) {
+        const parsedLineComment = parseLineComment(line);
+
+        if (!parsedLineComment) continue;
+
+        const { word, href } = parsedLineComment;
+
+        linkMap.set(word, href);
+
+        // Remove twoslash line from code lines
+        codeLines.splice(i, 1);
+      }
+
+      code = codeLines.join('\n');
+    }
+
+    // Convert code to HAST and apply transformers
+    const twoslashOptions = getTwoslashOptions({ linkMap });
+
+    const hast = highlighter.codeToHast(code, {
+      lang: lang ?? DEFAULT_LANG,
+      meta: shouldUseTwoslash ? { __raw: 'twoslash' } : undefined,
+      themes: {
+        light:
+          options.themes?.light ??
+          options.theme ??
+          (options.codeStyling === 'dark' ? DEFAULT_DARK_THEME : DEFAULT_LIGHT_THEME),
+        dark: options.themes?.dark ?? options.theme ?? DEFAULT_DARK_THEME,
+      },
+      colorReplacements: SHIKI_COLOR_REPLACEMENTS,
+      tabindex: false,
+      tokenizeMaxLineLength: 1000,
+      transformers: [...SHIKI_TRANSFORMERS, transformerTwoslash(twoslashOptions)],
+    });
+
+    // Set code node metadata
+    const preCodeNode = hast.children[0] as Element;
+
+    if (!preCodeNode) return;
+
+    node.data = node.data ?? {};
+
+    preCodeNode.data = node.data;
+    preCodeNode.properties.language = lang;
+
+    // Set pre node metadata
+    const codeNode = preCodeNode.children[0] as Element;
+
+    if (codeNode) {
+      codeNode.data = node.data;
+      codeNode.properties.language = lang;
+    }
+
+    // Replace previous precode node with new precode node
+    parent.children.splice(index, 1, preCodeNode);
+  } catch (err) {
+    if (err instanceof Error && /Unknown language/.test(err.message)) {
+      return;
+    }
+
+    throw err;
+  }
+}
+
+export { handlePreCode };
diff --git a/packages/mdx/src/plugins/rehype/index.ts b/packages/mdx/src/plugins/rehype/index.ts
index 865bc66..c639c01 100644
--- a/packages/mdx/src/plugins/rehype/index.ts
+++ b/packages/mdx/src/plugins/rehype/index.ts
@@ -1 +1,2 @@
-export * from './rehypeSyntaxHighlighting.js';
+export * from './core.js';
+export * from './types.js';
diff --git a/packages/mdx/src/plugins/rehype/rehypeSyntaxHighlighting.ts b/packages/mdx/src/plugins/rehype/rehypeSyntaxHighlighting.ts
deleted file mode 100644
index 87b3d57..0000000
--- a/packages/mdx/src/plugins/rehype/rehypeSyntaxHighlighting.ts
+++ /dev/null
@@ -1,214 +0,0 @@
-import { transformerTwoslash } from '@shikijs/twoslash';
-import { type } from 'arktype';
-import type { Element, Root } from 'hast';
-import { toString } from 'hast-util-to-string';
-import type { MdxJsxFlowElementHast, MdxJsxTextElementHast } from 'mdast-util-mdx-jsx';
-import { createHighlighter, type Highlighter } from 'shiki';
-import type { Plugin } from 'unified';
-import { visit } from 'unist-util-visit';
-
-import {
-  type ShikiLang,
-  type ShikiTheme,
-  shikiColorReplacements,
-  DEFAULT_LANG_ALIASES,
-  DEFAULT_LANG,
-  DEFAULT_DARK_THEME,
-  DEFAULT_LIGHT_THEME,
-  DEFAULT_THEMES,
-  DEFAULT_LANGS,
-  SHIKI_TRANSFORMERS,
-  UNIQUE_LANGS,
-} from './shiki-constants.js';
-import { TextMateGrammar, TextMateGrammarType } from './shiki/custom-language.js';
-import { getTwoslashOptions, parseLineComment } from './twoslash/config.js';
-import { getLanguage } from './utils.js';
-
-export type RehypeSyntaxHighlightingOptions = {
-  theme?: ShikiTheme;
-  themes?: Record<'light' | 'dark', ShikiTheme>;
-  codeStyling?: 'dark' | 'system' | 'light' | Record<string, unknown> | null;
-  linkMap?: Map<string, string>;
-  customLanguages?: string[];
-};
-
-let highlighterPromise: Promise<Highlighter> | null = null;
-
-async function getHighlighter(): Promise<Highlighter> {
-  if (!highlighterPromise) {
-    highlighterPromise = createHighlighter({
-      themes: DEFAULT_THEMES,
-      langs: DEFAULT_LANGS,
-    });
-  }
-  return highlighterPromise;
-}
-
-export const rehypeSyntaxHighlighting: Plugin<[RehypeSyntaxHighlightingOptions?], Root, Root> = (
-  options = {}
-) => {
-  return async (tree) => {
-    const nodesToProcess: Promise<void>[] = [];
-    const customLanguageNames: string[] = [];
-
-    const themesToLoad: ShikiTheme[] = [];
-    if (options.themes) {
-      themesToLoad.push(options.themes.dark);
-      themesToLoad.push(options.themes.light);
-    } else if (options.theme) {
-      themesToLoad.push(options.theme);
-    }
-
-    const highlighter = await getHighlighter();
-
-    await Promise.all([
-      ...themesToLoad
-        .filter(
-          (theme): theme is Exclude<ShikiTheme, 'css-variables'> =>
-            !DEFAULT_THEMES.includes(theme) && theme !== 'css-variables'
-        )
-        .map((theme) => highlighter.loadTheme(theme)),
-      ...(options.customLanguages?.map(async (unparsedLang) => {
-        const parsedLang = JSON.parse(unparsedLang);
-        const lang = TextMateGrammar(parsedLang);
-        if (lang instanceof type.errors) {
-          console.error(lang.summary);
-          return;
-        }
-        await highlighter.loadLanguage(lang);
-        const possibleNames = [lang.name, lang.displayName, ...(lang.aliases ?? [])];
-        customLanguageNames.push(...possibleNames.filter((l) => l != undefined));
-      }) ?? []),
-    ]);
-
-    visit(tree, 'element', (node, index, parent) => {
-      const child = node.children[0];
-      if (
-        !parent ||
-        index === undefined ||
-        node.type !== 'element' ||
-        node.tagName !== 'pre' ||
-        !child ||
-        child.type !== 'element' ||
-        child.tagName !== 'code'
-      ) {
-        return;
-      }
-
-      // set the metadata of `node` (which is a pre element) to that of
-      // `child` (which is the code element that likely contains all the metadata)
-      if (!Object.keys(node.properties).length) {
-        node.properties = child.properties;
-      }
-      if (!node.data) {
-        node.data = child.data;
-      }
-
-      let lang =
-        getLanguage(node, DEFAULT_LANG_ALIASES) ??
-        getLanguage(child, DEFAULT_LANG_ALIASES) ??
-        DEFAULT_LANG;
-
-      if (
-        !DEFAULT_LANGS.includes(lang) &&
-        !customLanguageNames.includes(lang) &&
-        UNIQUE_LANGS.includes(lang)
-      ) {
-        nodesToProcess.push(
-          highlighter.loadLanguage(lang).then(() => {
-            traverseNode({ node, index, parent, highlighter, lang, options });
-          })
-        );
-      } else {
-        if (!UNIQUE_LANGS.includes(lang) && !customLanguageNames.includes(lang)) {
-          lang = DEFAULT_LANG;
-        }
-        traverseNode({ node, index, parent, highlighter, lang, options });
-      }
-    });
-    await Promise.all(nodesToProcess);
-  };
-};
-
-function traverseNode({
-  node,
-  index,
-  parent,
-  highlighter,
-  lang,
-  options,
-}: {
-  node: Element;
-  index: number;
-  parent: Element | Root | MdxJsxTextElementHast | MdxJsxFlowElementHast;
-  highlighter: Highlighter;
-  lang: ShikiLang;
-  options: RehypeSyntaxHighlightingOptions;
-}) {
-  try {
-    let code = toString(node);
-
-    const meta = node.data?.meta?.split(' ') ?? [];
-    const twoslashIndex = meta.findIndex((str) => str.toLowerCase() === 'twoslash');
-    const shouldUseTwoslash = twoslashIndex > -1;
-
-    if (node.data && node.data.meta && shouldUseTwoslash) {
-      meta.splice(twoslashIndex, 1);
-      node.data.meta = meta.join(' ').trim() || undefined;
-    }
-
-    const linkMap = options.linkMap ?? new Map();
-    if (shouldUseTwoslash) {
-      const splitCode = code.split('\n');
-
-      for (const [i, line] of splitCode.entries()) {
-        const parsedLineComment = parseLineComment(line);
-        if (!parsedLineComment) continue;
-        const { word, href } = parsedLineComment;
-        linkMap.set(word, href);
-        splitCode.splice(i, 1);
-      }
-
-      code = splitCode.join('\n');
-    }
-
-    const twoslashOptions = getTwoslashOptions({ linkMap });
-
-    const hast = highlighter.codeToHast(code, {
-      lang: lang ?? DEFAULT_LANG,
-      meta: shouldUseTwoslash ? { __raw: 'twoslash' } : undefined,
-      themes: {
-        light:
-          options.themes?.light ??
-          options.theme ??
-          (options.codeStyling === 'dark' ? DEFAULT_DARK_THEME : DEFAULT_LIGHT_THEME),
-        dark: options.themes?.dark ?? options.theme ?? DEFAULT_DARK_THEME,
-      },
-      colorReplacements: shikiColorReplacements,
-      tabindex: false,
-      tokenizeMaxLineLength: 1000,
-      transformers: [...SHIKI_TRANSFORMERS, transformerTwoslash(twoslashOptions)],
-    });
-
-    const codeElement = hast.children[0] as Element;
-    if (!codeElement) return;
-
-    const preChild = codeElement.children[0] as Element;
-
-    node.data = node.data ?? {};
-    codeElement.data = node.data;
-    codeElement.properties.language = lang;
-    if (preChild) {
-      preChild.data = node.data;
-      preChild.properties.language = lang;
-    }
-    parent.children.splice(index, 1, codeElement);
-  } catch (err) {
-    if (err instanceof Error && /Unknown language/.test(err.message)) {
-      return;
-    }
-    throw err;
-  }
-}
-
-export { TextMateGrammar, type TextMateGrammarType };
diff --git a/packages/mdx/src/plugins/rehype/types.ts b/packages/mdx/src/plugins/rehype/types.ts
new file mode 100644
index 0000000..4733376
--- /dev/null
+++ b/packages/mdx/src/plugins/rehype/types.ts
@@ -0,0 +1,17 @@
+import { BundledLanguage } from 'shiki/types';
+
+import { SHIKI_THEMES } from './constants.js';
+
+type ShikiLang = BundledLanguage | 'ansi' | 'text';
+
+type ShikiTheme = (typeof SHIKI_THEMES)[number];
+
+interface RehypeSyntaxHighlightingOptions {
+  theme?: ShikiTheme;
+  themes?: Record<'light' | 'dark', ShikiTheme>;
+  codeStyling?: 'dark' | 'system' | 'light' | Record<string, unknown> | null;
+  linkMap?: Map<string, string>;
+  customLanguages?: string[];
+}
+
+export type { ShikiLang, ShikiTheme, RehypeSyntaxHighlightingOptions };
diff --git a/packages/mdx/src/plugins/rehype/utils.ts b/packages/mdx/src/plugins/rehype/utils.ts
index 545ede9..c76efbc 100644
--- a/packages/mdx/src/plugins/rehype/utils.ts
+++ b/packages/mdx/src/plugins/rehype/utils.ts
@@ -1,25 +1,98 @@
+import { type } from 'arktype';
 import type { Element } from 'hast';
 
-import { type ShikiLang } from './shiki-constants.js';
+import { DEFAULT_THEMES } from './constants.js';
+import { TextMateGrammar, TextMateGrammarType } from './shiki/custom-language.js';
+import type { RehypeSyntaxHighlightingOptions, ShikiLang, ShikiTheme } from './types.js';
 
-export function classNameOrEmptyArray(element: Element): string[] {
-  const className = element.properties.className;
-  if (Array.isArray(className) && className.every((el) => typeof el === 'string')) return className;
-  return [];
+const LANGUAGE_PREFIX = 'language-';
+
+function getElementClassNameList(element: Element): string[] {
+  const classNameList = element.properties.className;
+
+  const isArray = Array.isArray(classNameList);
+  const isStringArray = isArray && classNameList.every((item) => typeof item === 'string');
+
+  if (!isStringArray) return [];
+
+  return classNameList;
+}
+
+/** Get language from element's classname without `language-` prefix */
+function getLanguage(node: Element, aliases: Record<string, ShikiLang>): ShikiLang | undefined {
+  const classNameList = getElementClassNameList(node);
+  const languageClassName = classNameList.find((className) =>
+    className.startsWith(LANGUAGE_PREFIX)
+  );
+
+  if (!languageClassName) return undefined;
+
+  const language = languageClassName.slice(LANGUAGE_PREFIX.length);
+
+  return aliases[language];
+}
+
+function getThemesToLoad(options: RehypeSyntaxHighlightingOptions) {
+  const themesToLoad: ShikiTheme[] = [];
+  const themesToExclude = [...DEFAULT_THEMES, 'css-variables'];
+
+  if (options.themes) {
+    themesToLoad.push(options.themes.dark);
+    themesToLoad.push(options.themes.light);
+  } else if (options.theme) {
+    themesToLoad.push(options.theme);
+  }
+
+  const filteredThemesToLoad = themesToLoad.filter(
+    (theme) => !themesToExclude.includes(theme)
+  ) as Exclude<ShikiTheme, 'css-variables'>[];
+
+  return filteredThemesToLoad;
 }
 
-export function getLanguage(
-  node: Element,
-  aliases: Record<string, ShikiLang>
-): ShikiLang | undefined {
-  const className = classNameOrEmptyArray(node);
+function getParsedLanguages(languages: string[]) {
+  const parsedLanguages: TextMateGrammarType[] = [];
+
+  languages.forEach((language) => {
+    const parsedLang = JSON.parse(language);
+    const lang = TextMateGrammar(parsedLang);
 
-  for (const classListItem of className) {
-    if (classListItem.startsWith('language-')) {
-      const lang = classListItem.slice(9).toLowerCase();
-      if (lang) return aliases[lang] ?? (lang as ShikiLang);
+    if (lang instanceof type.errors) {
+      console.error(lang.summary);
+      return;
     }
+
+    parsedLanguages.push(lang);
+  });
+
+  return parsedLanguages;
+}
+
+function getLanguagesToLoad(options: RehypeSyntaxHighlightingOptions) {
+  if (options.customLanguages) {
+    const parsedLanguages = getParsedLanguages(options.customLanguages);
+    return parsedLanguages;
   }
 
-  return undefined;
+  return [];
 }
+
+function getCustomLanguagesNames(options: RehypeSyntaxHighlightingOptions) {
+  const customLanguageNames: string[] = [];
+
+  if (options.customLanguages) {
+    const parsedLanguages = getParsedLanguages(options.customLanguages);
+
+    parsedLanguages.forEach((lang) => {
+      const possibleNames = [lang.name, lang.displayName, ...(lang.aliases ?? [])].filter(
+        (l) => l != undefined
+      );
+
+      customLanguageNames.push(...possibleNames);
+    });
+  }
+
+  return customLanguageNames;
+}
+
+export { getLanguage, getThemesToLoad, getLanguagesToLoad, getCustomLanguagesNames };