import React from "react";
import Markdown from "react-markdown";
import remarkGfm from "remark-gfm";
import remarkMath from "remark-math";
import rehypeRaw from "rehype-raw";
import rehypeSanitize, { defaultSchema } from "rehype-sanitize";
import rehypePrismPlus from "rehype-prism-plus";
import rehypeKatex from "rehype-katex";
import remarkStreamingMarkdown, {
  type LastNodeInfo,
} from "@/utils/remarkStreamingMarkdown";
import type { PluggableList } from "unified";
import remarkCitationParser from "@/utils/remarkCitationParser";
import CopyButton from "./CopyButton";

interface StreamingMarkdownContentProps {
  content: string;
  isStreaming?: boolean;
  size?: "sm" | "md" | "lg";
  onLastNode?: (info: LastNodeInfo) => void;
  browserToolResult?: any; // TODO: proper type
}

const CodeBlock = React.memo(
  ({ children, className, ...props }: React.HTMLAttributes<HTMLPreElement>) => {
    const extractText = React.useCallback((node: React.ReactNode): string => {
      if (typeof node === "string") return node;
      if (typeof node === "number") return String(node);
      if (!node) return "";

      if (React.isValidElement(node)) {
        if (
          node.props &&
          typeof node.props === "object" &&
          "children" in node.props
        ) {
          return extractText(node.props.children as React.ReactNode);
        }
      }

      if (Array.isArray(node)) {
        return node.map(extractText).join("");
      }

      return "";
    }, []);

    const language = className?.replace(/language-/, "") || "";

    return (
      <div className="relative bg-neutral-100 dark:bg-neutral-800 rounded-2xl overflow-hidden my-6">
        <div className="flex justify-between select-none">
          <div className="text-[13px] text-neutral-500 dark:text-neutral-400 font-mono px-4 py-2">
            {language}
          </div>
          <CopyButton
            content={extractText(children)}
            showLabels={true}
            className="copy-button text-neutral-500 dark:text-neutral-400 bg-neutral-100 dark:bg-neutral-800"
          />
        </div>
        <pre className={className} {...props}>
          {children}
        </pre>
      </div>
    );
  },
);

const StreamingMarkdownContent: React.FC<StreamingMarkdownContentProps> =
  React.memo(
    ({ content, isStreaming = false, size, onLastNode, browserToolResult }) => {
      // Build the remark plugins array
      const remarkPlugins = React.useMemo(() => {
        const plugins: PluggableList = [
          remarkGfm,
          [remarkMath, { singleDollarTextMath: false }],
          remarkCitationParser,
        ];

        // Add streaming plugin when in streaming mode
        if (isStreaming) {
          plugins.push([remarkStreamingMarkdown, { debug: true, onLastNode }]);
        }

        return plugins;
      }, [isStreaming, onLastNode]);

      // Create a custom sanitization schema that allows math elements
      const sanitizeSchema = React.useMemo(() => {
        return {
          ...defaultSchema,
          attributes: {
            ...defaultSchema.attributes,
            span: [
              ...(defaultSchema.attributes?.span || []),
              ["className", /^katex/],
            ],
            div: [
              ...(defaultSchema.attributes?.div || []),
              ["className", /^katex/],
            ],
            "ol-citation": ["cursor", "start", "end"],
          },
          tagNames: [
            ...(defaultSchema.tagNames || []),
            "math",
            "mrow",
            "mi",
            "mo",
            "mn",
            "msup",
            "msub",
            "mfrac",
            "mover",
            "munder",
            "msqrt",
            "mroot",
            "merror",
            "mspace",
            "mpadded",
            "ol-citation",
          ],
        };
      }, []);

      return (
        <div
          className={`
          max-w-full
          ${size === "sm" ? "prose-sm" : size === "lg" ? "prose-lg" : ""}
          prose
          prose-neutral
          prose-headings:font-semibold
          prose:text-neutral-800
          prose-strong:font-semibold
          prose-ul:marker:text-neutral-700
          prose-li:marker:text-neutral-700
          prose-headings:text-neutral-800
          prose-pre:bg-transparent
          prose-pre:rounded-xl
          prose-pre:text-neutral-800
          prose-pre:font-normal
          prose-pre:my-0
          prose-pre:max-w-full
          prose-pre:pt-1
           [&_code:not(pre_code)]:text-neutral-700
          [&_code:not(pre_code)]:bg-neutral-100
          [&_code:not(pre_code)]:font-normal
          [&_code:not(pre_code)]:px-1.5
          [&_code:not(pre_code)]:py-0.5
          [&_code:not(pre_code)]:rounded-md
          [&_code:not(pre_code)]:text-[90%]
          [&_code:not(pre_code)]:before:hidden
          [&_code:not(pre_code)]:after:hidden
          dark:prose-invert
          dark:prose:text-neutral-200
          dark:prose-pre:bg-none
          dark:prose-headings:text-neutral-200
          dark:prose-strong:text-neutral-200
          dark:prose-pre:text-neutral-200
          dark:prose:pre:text-neutral-200
          dark:[&_code:not(pre_code)]:text-neutral-200
          dark:[&_code:not(pre_code)]:bg-neutral-800
          dark:[&_code:not(pre_code)]:font-normal
          dark:prose-ul:marker:text-neutral-300
          dark:prose-li:marker:text-neutral-300
          break-words
        `}
        >
          <StreamingMarkdownErrorBoundary
            content={content}
            isStreaming={isStreaming}
          >
            <Markdown
              remarkPlugins={remarkPlugins}
              rehypePlugins={
                [
                  [rehypeRaw, { allowDangerousHtml: true }],
                  [rehypeSanitize, sanitizeSchema],
                  [rehypePrismPlus, { ignoreMissing: true }],
                  [
                    rehypeKatex,
                    {
                      errorColor: "#000000", // Black instead of red for errors
                      strict: false, // Be more lenient with parsing
                      throwOnError: false,
                    },
                  ],
                ] as PluggableList
              }
              components={{
                pre: CodeBlock,
                table: ({
                  children,
                  ...props
                }: React.HTMLAttributes<HTMLTableElement>) => (
                  <div className="overflow-x-auto max-w-full">
                    <table {...props}>{children}</table>
                  </div>
                ),
                // @ts-expect-error: custom type
                "ol-citation": ({
                  cursor,
                  // start,
                  // end,
                }: {
                  cursor: number;
                  start: number;
                  end: number;
                }) => {
                  // Check if we have a page_stack and if the cursor is valid
                  const pageStack = browserToolResult?.page_stack;
                  const hasValidPage = pageStack && cursor < pageStack.length;
                  const pageUrl = hasValidPage ? pageStack[cursor] : null;

                  // Extract a readable title from the URL if possible
                  const getPageTitle = (url: string) => {
                    if (url.startsWith("search_results_")) {
                      const searchTerm = url.substring(
                        "search_results_".length,
                      );
                      return `Search: ${searchTerm}`;
                    }
                    // For regular URLs, try to extract domain or use full URL
                    try {
                      const urlObj = new URL(url);
                      return urlObj.hostname;
                    } catch {
                      // If not a valid URL, return as is
                      return url;
                    }
                  };

                  const citationElement = (
                    <span className="text-xs text-neutral-500 dark:text-neutral-400 bg-neutral-100 dark:bg-neutral-800 rounded-full px-2 py-1 ml-1">
                      [{cursor}]
                    </span>
                  );

                  // If we have a valid page URL, wrap in a link
                  if (pageUrl && pageUrl.startsWith("http")) {
                    return (
                      <a
                        href={pageUrl}
                        target="_blank"
                        rel="noopener noreferrer"
                        className="inline-flex items-center hover:opacity-80 transition-opacity no-underline"
                        title={getPageTitle(pageUrl)}
                      >
                        {citationElement}
                      </a>
                    );
                  }

                  // Otherwise, just return the citation without a link
                  return citationElement;
                },
              }}
            >
              {content}
            </Markdown>
          </StreamingMarkdownErrorBoundary>
        </div>
      );
    },
  );

interface StreamingMarkdownErrorBoundaryProps {
  content: string;
  children: React.ReactNode;
  isStreaming: boolean;
}

// Sometimes remark will throw errors, particularly when rendering math. We add
// this fallback to show the plain text content if there's an error, and then we
// retry rendering the content if the content changes OR if we change our
// streaming state (because we render things differently when in streaming mode
// v. not, so for some cases we'll automatically recover once streaming is over)
//
// This should not be relied on for anything known to be broken (any known
// errors should be fixed!), but it's necessary to not break the full UI because
// of some bad markdown
class StreamingMarkdownErrorBoundary extends React.Component<
  StreamingMarkdownErrorBoundaryProps,
  { hasError: boolean }
> {
  constructor(props: StreamingMarkdownErrorBoundaryProps) {
    super(props);
    this.state = { hasError: false };
  }

  componentDidUpdate(prevProps: StreamingMarkdownErrorBoundaryProps) {
    if (
      prevProps.isStreaming !== this.props.isStreaming ||
      prevProps.content !== this.props.content
    ) {
      this.setState({ hasError: false });
    }
  }

  static getDerivedStateFromError(/*_error: Error*/) {
    return { hasError: true };
  }

  componentDidCatch(error: Error, info: { componentStack: string }) {
    console.error(
      "StreamingMarkdownContent: caught rendering error",
      error,
      info,
    );
  }

  render() {
    if (this.state.hasError) {
      // TODO(drifkin): render this more nicely so it's not so jarring. For
      // example, probably want to render newlines, etc. But let's not get too
      // fancy because then we'll end up needing an ErrorBoundaryErrorBoundary
      // :upside_down_face:
      return <div>{this.props.content}</div>;
    }

    return this.props.children;
  }
}

export default StreamingMarkdownContent;
