diff --git a/apps/docs/src/components/Providers.tsx b/apps/docs/src/components/Providers.tsx
index 3c378ad..b463f3d 100644
--- a/apps/docs/src/components/Providers.tsx
+++ b/apps/docs/src/components/Providers.tsx
@@ -1,12 +1,12 @@
 "use client";
 
-import { createContext, useCallback, useContext, useEffect, useState } from "react";
+import React, { createContext, useCallback, useContext, useEffect, useState } from "react";
 
 type SidebarProps = {
   isSidebarOpen: boolean;
   isSearchOpen: boolean;
   toggleSidebar: () => void;
-  toggleSearch: () => void;
+  toggleSearch: React.MouseEventHandler<HTMLButtonElement>;
 };
 
 const SidebarContext = createContext<SidebarProps>({
diff --git a/apps/docs/src/components/Search.tsx b/apps/docs/src/components/Search.tsx
index 7886888..1ce4ea7 100644
--- a/apps/docs/src/components/Search.tsx
+++ b/apps/docs/src/components/Search.tsx
@@ -5,14 +5,14 @@
 /* eslint-disable @typescript-eslint/no-unsafe-assignment */
 /* eslint-disable @typescript-eslint/no-unsafe-call */
 // @ts-nocheck
+
 "use client";
 
-import { useEffect, useState } from "react";
+import { forwardRef, useEffect, useState, type HTMLAttributes } from "react";
 import Link from "next/link";
-import Router from "next/router";
 import { DocSearchModal, useDocSearchKeyboardEvents } from "@docsearch/react";
 import { SearchIcon } from "@iconicicons/react";
-import { Button, Kbd } from "@lemonsqueezy/wedges";
+import { Button, Kbd, type KbdKey } from "@lemonsqueezy/wedges";
 import { createPortal } from "react-dom";
 
 import { cn } from "@/lib/utils";
@@ -20,9 +20,17 @@ import { cn } from "@/lib/utils";
 import { useSidebar } from "./Providers";
 
 const docSearchConfig = {
-  appId: process.env.NEXT_PUBLIC_DOCSEARCH_APP_ID,
-  apiKey: process.env.NEXT_PUBLIC_DOCSEARCH_API_KEY,
-  indexName: process.env.NEXT_PUBLIC_DOCSEARCH_INDEX_NAME,
+  appId: process.env.NEXT_PUBLIC_DOCSEARCH_APP_ID!,
+  apiKey: process.env.NEXT_PUBLIC_DOCSEARCH_API_KEY!,
+  indexName: process.env.NEXT_PUBLIC_DOCSEARCH_INDEX_NAME!,
+};
+
+const suffixTitle = (title: string) => {
+  if (!title) {
+    return title;
+  }
+
+  return `${title} - Wedges Docs`;
 };
 
 function Hit({ hit, children }) {
@@ -42,19 +50,10 @@ function Hit({ hit, children }) {
   );
 }
 
-export function Search({ className, ...otherProps }: { className?: string; otherProps?: any }) {
+const Search = forwardRef<HTMLDivElement, HTMLAttributes<HTMLDivElement>>((props, ref) => {
+  const { className, ...otherProps } = props;
   const { isSearchOpen, toggleSearch } = useSidebar();
-  const [modifierKey, setModifierKey] = useState("command");
-
-  const suffixTitle = (title) => {
-    if (!title) {
-      return title;
-    }
-
-    const section = "Wedges Docs";
-
-    return `${title} - ${section}`;
-  };
+  const [modifierKey, setModifierKey] = useState<KbdKey>("command");
 
   useDocSearchKeyboardEvents({ isOpen: isSearchOpen, onOpen: toggleSearch, onClose: toggleSearch });
 
@@ -63,7 +62,7 @@ export function Search({ className, ...otherProps }: { className?: string; other
   }, []);
 
   return (
-    <>
+    <div ref={ref} className={cn(className)}>
       <Button
         asChild
         className={cn(
@@ -92,70 +91,33 @@ export function Search({ className, ...otherProps }: { className?: string; other
       {isSearchOpen &&
         createPortal(
           <DocSearchModal
-            {...docSearchConfig}
             hitComponent={Hit}
+            onClose={toggleSearch}
             initialScrollY={window.scrollY}
-            navigator={{
-              navigate({ itemUrl }) {
-                Router.push(itemUrl);
-              },
-            }}
+            {...docSearchConfig}
+            {...otherProps}
             transformItems={(items) => {
               return items.map((item, index) => {
-                // We transform the absolute URL into a relative URL to
-                // leverage Next's preloading.
+                // We transform the absolute URL into a relative URL to leverage Next's preloading.
                 const a = document.createElement("a");
-
                 a.href = item.url;
+                const hash = a.hash;
 
-                const hash = a.hash === "#content-wrapper" || a.hash === "#header" ? "" : a.hash;
-
-                if (item.hierarchy?.lvl0) {
-                  item.hierarchy.lvl0 = item.hierarchy.lvl0.replace(/&amp;/g, "&");
-                  item.hierarchy.lvl0 = item.hierarchy.lvl0.replace("Documentation ", "");
+                if (item?.content) {
+                  item.content = suffixTitle(item.content);
                 }
 
-                if (item._highlightResult?.hierarchy?.lvl0?.value) {
-                  item._highlightResult.hierarchy.lvl0.value =
-                    item._highlightResult.hierarchy.lvl0.value.replace(/&amp;/g, "&");
-
-                  item._highlightResult.hierarchy.lvl0.value =
-                    item._highlightResult.hierarchy.lvl0.value.replace("Documentation ", "");
-                }
-
-                if (item.content) {
-                  item.content = suffixTitle(item.content, item);
-                }
-
-                if (item._snippetResult?.content?.value) {
+                if (item?._snippetResult?.content?.value) {
                   item._snippetResult.content.value = suffixTitle(
-                    item._snippetResult.content.value,
-                    item
+                    item._snippetResult.content.value
                   );
                 }
-                if (item._highlightResult?.content?.value) {
+
+                if (item?._highlightResult?.content?.value) {
                   item._highlightResult.content.value = suffixTitle(
-                    item._highlightResult.content.value,
-                    item
-                  );
-                }
-                if (
-                  item._snippetResult?.hierarchy?.[item.type]?.value &&
-                  item._snippetResult?.hierarchy?.[item.type]?.value !== null
-                ) {
-                  item._snippetResult.hierarchy[item.type].value = suffixTitle(
-                    item._snippetResult?.hierarchy?.lvl2?.value,
-                    item
+                    item._highlightResult.content.value
                   );
                 }
-                if (item._highlightResult?.hierarchy_camel[0]?.[item.type]?.value) {
-                  item._highlightResult.hierarchy_camel[0][item.type].value = suffixTitle(
-                    item._highlightResult.hierarchy_camel[0][item.type].value,
-                    item
-                  );
-                }
-
-                item.hierarchy[item.type] = suffixTitle(item.hierarchy[item.type], item);
 
                 return {
                   ...item,
@@ -165,22 +127,21 @@ export function Search({ className, ...otherProps }: { className?: string; other
                   __is_child: () =>
                     item.type !== "lvl1" &&
                     items.length > 1 &&
-                    items[0].type === "lvl1" &&
+                    items[0]?.type === "lvl1" &&
                     index !== 0,
                   __is_first: () => index === 1,
                   __is_last: () => index === items.length - 1 && index !== 0,
                 };
               });
             }}
-            onClose={toggleSearch}
           />,
           document.body
         )}
-    </>
+    </div>
   );
-}
+});
 
-export function SearchButton() {
+function SearchButton() {
   const { toggleSearch } = useSidebar();
 
   return (
@@ -195,3 +156,5 @@ export function SearchButton() {
     </Button>
   );
 }
+
+export { Search, SearchButton };