diff --git a/package.json b/package.json
index 38b237d..c17b50c 100644
--- a/package.json
+++ b/package.json
@@ -1,7 +1,7 @@
 {
   "name": "virtual-react-json-diff",
   "type": "module",
-  "version": "1.0.13",
+  "version": "1.0.14",
   "description": "Fast, virtualized React component for visually comparing large JSON objects. Includes search, theming, and minimap.",
   "author": {
     "name": "Utku Akyüz"
diff --git a/src/components/DiffViewer/components/VirtualDiffGrid.tsx b/src/components/DiffViewer/components/VirtualDiffGrid.tsx
index a1b0146..80b6171 100644
--- a/src/components/DiffViewer/components/VirtualDiffGrid.tsx
+++ b/src/components/DiffViewer/components/VirtualDiffGrid.tsx
@@ -1,4 +1,3 @@
-// VirtualDiffGrid.tsx
 import type { InlineDiffOptions } from "json-diff-kit";
 import type { Dispatch } from "react";
 import type { ListOnScrollProps } from "react-window";
@@ -22,7 +21,6 @@ type ListDataType = {
 type VirtualDiffGridProps = {
   leftDiff: DiffRowOrCollapsed[];
   rightDiff: DiffRowOrCollapsed[];
-  outerRef: React.RefObject<Node | null>;
   listRef: React.RefObject<List<ListDataType>>;
   height: number;
   inlineDiffOptions?: InlineDiffOptions;
@@ -30,12 +28,13 @@ type VirtualDiffGridProps = {
   setScrollTop: Dispatch<React.SetStateAction<number>>;
   onExpand: (segmentIndex: number) => void;
   overScanCount?: number;
+  viewerRef?: React.RefObject<HTMLDivElement>;
+  listContainerRef?: React.RefObject<HTMLDivElement>;
 };
 
 const VirtualDiffGrid: React.FC<VirtualDiffGridProps> = ({
   leftDiff,
   rightDiff,
-  outerRef,
   listRef,
   height,
   inlineDiffOptions,
@@ -43,6 +42,8 @@ const VirtualDiffGrid: React.FC<VirtualDiffGridProps> = ({
   setScrollTop,
   onExpand,
   overScanCount = 10,
+  viewerRef,
+  listContainerRef,
 }) => {
   // Virtual List Data
   const listData = useMemo(
@@ -65,7 +66,7 @@ const VirtualDiffGrid: React.FC<VirtualDiffGridProps> = ({
 
   // ROW HEIGHT CALCULATION
   const ROW_HEIGHT = useMemo(() => getRowHeightFromCSS(), []);
-  const rowHeights = useRowHeights(leftDiff);
+  const rowHeights = useRowHeights(leftDiff, viewerRef);
   const dynamicRowHeights = useCallback(
     (index: number) => {
       const leftLine = leftDiff[index];
@@ -81,12 +82,12 @@ const VirtualDiffGrid: React.FC<VirtualDiffGridProps> = ({
   }, [rowHeights]);
 
   return (
-    <div className={classes}>
+    <div className={classes} ref={viewerRef}>
       <List
         height={height}
         width="100%"
         style={{ alignItems: "start" }}
-        outerRef={outerRef}
+        outerRef={listContainerRef}
         ref={listRef}
         className="virtual-json-diff-list-container"
         itemCount={Math.max(leftDiff.length, rightDiff.length)}
diff --git a/src/components/DiffViewer/components/VirtualizedDiffViewer.tsx b/src/components/DiffViewer/components/VirtualizedDiffViewer.tsx
index 7279cdd..d22633c 100644
--- a/src/components/DiffViewer/components/VirtualizedDiffViewer.tsx
+++ b/src/components/DiffViewer/components/VirtualizedDiffViewer.tsx
@@ -33,10 +33,11 @@ export const VirtualizedDiffViewer: React.FC<VirtualizedDiffViewerProps> = ({
   inlineDiffOptions,
   overScanCount,
 }) => {
-  const outerRef = useRef<Node>(null);
   const listRef = useRef<List>(null);
   const getDiffDataRef = useRef<typeof getDiffData>();
   const lastSent = useRef<number>();
+  const viewerRef = useRef<HTMLDivElement>(null);
+  const listContainerRef = useRef<HTMLDivElement>(null);
 
   const differ = customDiffer ?? useMemo(
     () =>
@@ -75,6 +76,8 @@ export const VirtualizedDiffViewer: React.FC<VirtualizedDiffViewerProps> = ({
       listRef.current?.scrollToItem(idx, "center");
       onSearchMatch?.(idx);
     },
+    viewerRef,
+    listContainerRef,
   );
 
   const handleExpand = useCallback(
@@ -140,7 +143,6 @@ export const VirtualizedDiffViewer: React.FC<VirtualizedDiffViewerProps> = ({
       {/* List & Minimap */}
       <div style={{ display: "flex", gap: "8px", position: "relative" }}>
         <VirtualDiffGrid
-          outerRef={outerRef}
           listRef={listRef}
           leftDiff={leftView}
           rightDiff={rightView}
@@ -150,6 +152,8 @@ export const VirtualizedDiffViewer: React.FC<VirtualizedDiffViewerProps> = ({
           onExpand={handleExpand}
           className="virtual-json-diff-list-container"
           inlineDiffOptions={inlineDiffOptions}
+          viewerRef={viewerRef}
+          listContainerRef={listContainerRef}
         />
 
         <div className="minimap-overlay">
diff --git a/src/components/DiffViewer/hooks/useRowHeights.ts b/src/components/DiffViewer/hooks/useRowHeights.ts
index 9634f99..8f0027e 100644
--- a/src/components/DiffViewer/hooks/useRowHeights.ts
+++ b/src/components/DiffViewer/hooks/useRowHeights.ts
@@ -14,11 +14,14 @@ function getWrapCount(el: Element) {
   return Math.round(el.scrollHeight / lh);
 }
 
-export function useRowHeights(leftView: DiffRowOrCollapsed[]) {
+export function useRowHeights(leftView: DiffRowOrCollapsed[], viewerRef?: React.RefObject<HTMLDivElement | null>) {
   const [rowHeights, setRowHeights] = useState<number[]>([]);
 
   const measureRows = useCallback(() => {
-    const preElements = document.querySelectorAll(".json-diff-viewer pre");
+    if (!viewerRef?.current)
+      return;
+
+    const preElements = viewerRef.current.querySelectorAll("pre");
     const newHeights: number[] = [];
     for (let i = 0; i < preElements.length; i += 2) {
       const leftWraps = getWrapCount(preElements[i]);
@@ -26,11 +29,11 @@ export function useRowHeights(leftView: DiffRowOrCollapsed[]) {
       newHeights.push(Math.max(leftWraps, rightWraps));
     }
     setRowHeights(newHeights);
-  }, []);
+  }, [viewerRef]);
 
   useLayoutEffect(() => {
     measureRows();
-  }, [leftView]);
+  }, [leftView, measureRows]);
 
   useLayoutEffect(() => {
     window.addEventListener("resize", measureRows);
diff --git a/src/components/DiffViewer/hooks/useSearch.ts b/src/components/DiffViewer/hooks/useSearch.ts
index bd617f1..fde0221 100644
--- a/src/components/DiffViewer/hooks/useSearch.ts
+++ b/src/components/DiffViewer/hooks/useSearch.ts
@@ -2,10 +2,16 @@ import { useCallback, useEffect, useRef, useState } from "react";
 
 import type { DiffRowOrCollapsed, SearchState } from "../types";
 
-import { DIFF_VIEWER_CLASS, SEARCH_DEBOUNCE_MS } from "../utils/constants";
+import { SEARCH_DEBOUNCE_MS } from "../utils/constants";
 import { highlightMatches, performSearch } from "../utils/diffSearchUtils";
 
-export function useSearch(leftView: DiffRowOrCollapsed[], initialTerm?: string, onSearchMatch?: (index: number) => void) {
+export function useSearch(
+  leftView: DiffRowOrCollapsed[],
+  initialTerm?: string,
+  onSearchMatch?: (index: number) => void,
+  viewerRef?: React.RefObject<HTMLDivElement | null>,
+  listContainerRef?: React.RefObject<HTMLDivElement | null>,
+) {
   const [searchState, setSearchState] = useState<SearchState>({
     term: initialTerm ?? "",
     results: [],
@@ -43,17 +49,18 @@ export function useSearch(leftView: DiffRowOrCollapsed[], initialTerm?: string,
   }, [searchState, onSearchMatch]);
 
   useEffect(() => {
-    highlightMatches(searchState.term, DIFF_VIEWER_CLASS);
+    if (!viewerRef?.current)
+      return;
+
+    highlightMatches(searchState.term, viewerRef.current);
 
-    const observer = new MutationObserver(() => highlightMatches(searchState.term, DIFF_VIEWER_CLASS));
+    const observer = new MutationObserver(() => viewerRef.current && highlightMatches(searchState.term, viewerRef.current));
     const config = { childList: true, subtree: true };
-    const viewer = document.querySelector(`.${DIFF_VIEWER_CLASS}`);
-    if (viewer)
-      observer.observe(viewer, config);
+    observer.observe(viewerRef.current, config);
 
-    const listContainer = document.querySelector(".virtual-json-diff-list-container");
+    const listContainer = listContainerRef?.current;
     if (listContainer) {
-      const handleScroll = () => setTimeout(() => highlightMatches(searchState.term, DIFF_VIEWER_CLASS), 100);
+      const handleScroll = () => setTimeout(() => viewerRef.current && highlightMatches(searchState.term, viewerRef.current), 100);
       listContainer.addEventListener("scroll", handleScroll);
       return () => {
         observer.disconnect();
@@ -62,7 +69,7 @@ export function useSearch(leftView: DiffRowOrCollapsed[], initialTerm?: string,
     }
 
     return () => observer.disconnect();
-  }, [searchState.term]);
+  }, [searchState.term, viewerRef, listContainerRef]);
 
   useEffect(() => {
     if (initialTerm !== undefined) {
diff --git a/src/components/DiffViewer/utils/diffSearchUtils.ts b/src/components/DiffViewer/utils/diffSearchUtils.ts
index aea81a3..322395c 100644
--- a/src/components/DiffViewer/utils/diffSearchUtils.ts
+++ b/src/components/DiffViewer/utils/diffSearchUtils.ts
@@ -16,16 +16,16 @@ export function performSearch(term: string, leftView: DiffRowOrCollapsed[]): num
   return results;
 }
 
-export function highlightMatches(term: string, className: string = "json-diff-viewer-theme-custom"): void {
+export function highlightMatches(term: string, container: HTMLDivElement): void {
   if (!term) {
-    const elements = document.querySelectorAll(`.${className} span.token.search-match`);
+    const elements = container.querySelectorAll("span.token.search-match");
     elements.forEach(element => element.classList.remove("search-match"));
     return;
   }
 
   const termToUse = term.replaceAll("(", "").replaceAll(")", "");
   const regex = new RegExp(termToUse, "gi");
-  const elements = document.querySelectorAll(`.${className} span.token`);
+  const elements = container.querySelectorAll("span.token");
 
   elements.forEach((element) => {
     const text = element.textContent || "";