"use client";
import {
  Popover,
  PopoverTrigger,
  PopoverContent,
} from "@/components/ui/popover";
import { useMemo } from "react";
import ReactMarkdown, { Components } from "react-markdown";
import rehypeRaw from "rehype-raw";
import { useTTS } from "./TTSProvider";
import rehypeHighlight from "@/lib/utils";
import { Database } from "@/utils/supabase/types";

// Utility to escape regex special characters:
function escapeRegExp(text: string) {
  return text.replace(/[-[\]{}()*+?.,\\^$|#\s]/g, "\\$&");
}

export type OCRData = {
  index: number;
  images: string[];
  markdown: string;
  citations: {
    text: string;
    number: string;
  }[];
  dimensions: {
    dpi: number;
    width: number;
    height: number;
  };
};

export default function MarkdownRenderer({
  document,
}: {
  document: Database["public"]["Tables"]["documents"]["Row"];
}) {
  // Obtain TTS info from context.
  // TTSProvider is already wrapping this component higher in the tree.
  const { currentSentence, sentences } = useTTS();

  // Determine the text to highlight.
  const textToHighlight = useMemo(() => {
    if (!sentences || sentences.length === 0) return "";
    return sentences[currentSentence] || "";
  }, [sentences, currentSentence]);

  // Setup rehype plugins including our highlight plugin.
  const rehypePlugins = useMemo(
    () => [rehypeRaw, [rehypeHighlight, { textToHighlight }] as any],
    [textToHighlight]
  );

  const ocr = document?.ocr_data as OCRData[];

  const rawContent = ocr.map((page) => page.markdown).join("\n") || "";

  const citations: {
    text: string;
    page: number;
    index: string;
    number: number;
  }[] = [];
  const totalPages = ocr.length;
  const totalSentences = sentences.length;

  let totalCitations = 0;
  ocr.forEach((page) => {
    // each page has its own citations (1-N), so we need to map them correctly
    page.citations.forEach((citation, index) => {
      totalCitations += 1;
      citations.push({
        text: citation.text,
        page: page.index,
        index: (totalCitations + index).toString(), // unique index across all pages
        number: totalCitations + index + 1, // 1-based numbering
      });
    });
  });

  const components: Components = {
    h1: ({ node, ...props }) => (
      <h1 className="text-2xl font-semibold mb-4 text-white" {...props} />
    ),
    h2: ({ node, ...props }) => (
      <h2 className="text-xl font-medium mb-3 text-white" {...props} />
    ),
    h3: ({ node, ...props }) => (
      <h3 className="text-lg font-medium mb-2 text-gray-300" {...props} />
    ),
    h4: ({ node, ...props }) => (
      <h4 className="text-lg font-medium mb-2 text-gray-300" {...props} />
    ),
    p: ({ node, ...props }) => (
      <p className="leading-7 text-gray-200" {...props} />
    ),
    img: ({ node, ...props }) => (
      <img
        className="rounded-lg shadow-sm"
        style={{ maxWidth: "100%", height: "auto" }}
        {...props}
      />
    ),
    a: ({ node, ...props }) => (
      <a className="text-blue-400 hover:underline" {...props} />
    ),
    strong: ({ node, ...props }) => (
      <strong className="text-gray-200 font-semibold" {...props} />
    ),
    blockquote: ({ node, ...props }) => (
      <blockquote
        className="italic border-l-4 pl-4 border-gray-600 text-gray-300"
        {...props}
      />
    ),
    code: ({ node, ...props }) => (
      <code
        className="bg-gray-800 rounded px-1 py-0.5 text-gray-200"
        {...props}
      />
    ),
    sup: ({ node, ...props }) => {
      // Check if the text contains a reference number
      const text = props.children!.toString();

      const referenceNumber = text;
      if (!referenceNumber) {
        return <sup {...props} data-error="no reference number" />;
      }

      const citation = citations.find(
        (c) =>
          c.index === referenceNumber || c.number.toString() === referenceNumber
      );

      if (!citation) {
        return <sup {...props} data-error="no citation found" />;
      }

      return (
        // TODO: get the references from the document and display them in a popover
        <Popover>
          <PopoverTrigger asChild>
            <sup
              className="text-gray-200 cursor-pointer underline hover:cursor-pointer"
              {...props}
            />
          </PopoverTrigger>
          <PopoverContent className="w-56 overflow-hidden rounded-lg p-0">
            <div className="p-4">
              <p>{citation.text}</p>
            </div>
          </PopoverContent>
        </Popover>
      );
    },
  };

  return (
    <ReactMarkdown
      children={rawContent}
      components={components}
      rehypePlugins={rehypePlugins}
    />
  );
}