First time highlight load issue.

[gautamsarawagi]

@agentcooper This code works correctly to load the documents but the highlights are not coming in the start. When i click on the document anywhere only then I am able to see them why?

I guess somehow the highlights are not able to be present in the start but somehow they comes in one-click. That is what I am not able to understand actually.

Can anyone please help me with this.

import { useCallback, useEffect, useRef, useState } from "react"
import { XIcon } from "lucide-react"

import { Button } from "../ui/button"
import {
  RightDrawer,
  RightDrawerTitle,
  RightDrawerHeader,
  RightDrawerContent,
  RightDrawerClose,
  RightDrawerDescription,
} from "@/components/right-drawer"

import {
  PdfLoader,
  PdfHighlighter,
  Highlight,
  Popup,
  AreaHighlight,
  type IHighlight,
  Tip,
  ScaledPosition,
  NewHighlight,
  Content,
} from "react-pdf-highlighter"

import { Spinner } from "../ui/spinner"

export interface SourceWithCitations {
  title: string
  document_url: string
  citations: any[]
}

interface DocumentWithCitationsSidebarProps {
  isOpen: boolean
  onClose: () => void
  sources: SourceWithCitations | null
}

const HighlightPopup = ({ comment }: { comment?: { text?: string; emoji?: string } }) =>
  comment && comment.text ? (
    <div className="Highlight__popup">
      {comment.emoji} {comment.text}
    </div>
  ) : null

const resetHash = () => {
  document.location.hash = ""
}

const getNextId = () => String(Math.random()).slice(2)

const parseIdFromHash = () =>
  document.location.hash.slice("#highlight-".length)

export function DocumentWithCitationsSidebar({
  isOpen,
  onClose,
  sources,
}: DocumentWithCitationsSidebarProps) {

  const [highlights, setHighlights] = useState<Array<IHighlight>>(sources?.citations ?? [])

  // store scrollTo ref from PdfHighlighter
  const scrollViewerTo = useRef<(highlight: IHighlight) => void>(() => {})

  // Reset highlights whenever sources change
  useEffect(() => {
    if (sources?.citations) {
      const mapped = sources.citations.map((c, i) => ({
        id: c.id ?? String(i + 1),
        content: c.content ?? { text: "" },
        comment: c.comment ?? {},
        position: c.position,
      }))
      setHighlights(mapped)
    } else {
      setHighlights([])
    }
  }, [sources])

  // After highlights are set, force re-render/scroll
  useEffect(() => {
    if (highlights.length > 0 && scrollViewerTo.current) {
      // wait for PdfHighlighter to fully render pages
      setTimeout(() => {
        scrollViewerTo.current(highlights[0])
        // trigger a resize to ensure highlights paint correctly
        window.dispatchEvent(new Event("resize"))
      }, 50)
    }
  }, [highlights])

  const handleOpenChange = (open: boolean) => {
    if (!open) {
      onClose()
    }
  }

  const getHighlightById = useCallback(
    (id: string) => highlights.find((h) => h.id === id),
    [highlights],
  )

  const scrollToHighlightFromHash = useCallback(() => {
    const highlight = getHighlightById(parseIdFromHash())
    if (highlight) {
      scrollViewerTo.current(highlight)
    }
  }, [getHighlightById])

  useEffect(() => {
    window.addEventListener("hashchange", scrollToHighlightFromHash, false)
    return () => {
      window.removeEventListener("hashchange", scrollToHighlightFromHash, false)
    }
  }, [scrollToHighlightFromHash])

  const addHighlight = (highlight: NewHighlight) => {
    console.log("Saving highlight", highlight)
    setHighlights((prev) => [
      { ...highlight, id: getNextId() } as IHighlight,
      ...prev,
    ])
  }

  const updateHighlight = (
    highlightId: string,
    position: Partial<ScaledPosition>,
    content: Partial<Content>,
  ) => {
    console.log("Updating highlight", highlightId, position, content)
    setHighlights((prev) =>
      prev.map((h) =>
        h.id === highlightId
          ? {
              ...h,
              position: { ...h.position, ...position },
              content: { ...h.content, ...content },
            }
          : h,
      ),
    )
  }

  return (
    <RightDrawer open={isOpen} onOpenChange={handleOpenChange}>
      <RightDrawerContent className="w-[700px] max-w-[90vw]">
        <RightDrawerHeader className="flex items-center justify-between p-6 pb-4 border-b border-slate-100">
          <div>
            <RightDrawerTitle className="text-lg font-semibold text-slate-900">
              {sources?.title}
            </RightDrawerTitle>
            <RightDrawerDescription className="text-sm text-slate-500">
              {sources?.document_url}
            </RightDrawerDescription>
          </div>
          <RightDrawerClose asChild>
            <Button
              variant="ghost"
              size="sm"
              className="text-slate-400 hover:text-slate-600 hover:bg-slate-50"
            >
              <XIcon className="w-4 h-4" />
            </Button>
          </RightDrawerClose>
        </RightDrawerHeader>

        {sources?.document_url && (
          <PdfLoader url={sources.document_url} beforeLoad={<Spinner />}>
            {(pdfDocument) => (
              <PdfHighlighter
                pdfDocument={pdfDocument}
                enableAreaSelection={(event) => event.altKey}
                onScrollChange={resetHash}
                scrollRef={(scrollTo) => {
                  scrollViewerTo.current = scrollTo
                }}
                onSelectionFinished={(
                  position,
                  content,
                  hideTipAndSelection,
                  transformSelection,
                ) => (
                  <Tip
                    onOpen={transformSelection}
                    onConfirm={(comment) => {
                      addHighlight({ content, position, comment })
                      hideTipAndSelection()
                    }}
                  />
                )}
                highlightTransform={(
                  highlight,
                  index,
                  setTip,
                  hideTip,
                  viewportToScaled,
                  screenshot,
                  isScrolledTo,
                ) => {
                  const isTextHighlight = !highlight.content?.image

                  const component = isTextHighlight ? (
                    <Highlight
                      isScrolledTo={isScrolledTo}
                      position={highlight.position}
                      comment={highlight.comment}
                    />
                  ) : (
                    <AreaHighlight
                      isScrolledTo={isScrolledTo}
                      highlight={highlight}
                      onChange={(boundingRect) => {
                        updateHighlight(
                          highlight.id,
                          { boundingRect: viewportToScaled(boundingRect) },
                          { image: screenshot(boundingRect) },
                        )
                      }}
                    />
                  )

                  return (
                    <Popup
                      key={index}
                      popupContent={
                        <HighlightPopup comment={highlight.comment} />
                      }
                      onMouseOver={(popupContent) =>
                        setTip(highlight, () => popupContent)
                      }
                      onMouseOut={hideTip}
                    >
                      {component}
                    </Popup>
                  )
                }}
                highlights={highlights}
              />
            )}
          </PdfLoader>
        )}
      </RightDrawerContent>
    </RightDrawer>
  )
}