import ReactMarkdown from "react-markdown";
import "katex/dist/katex.min.css";
import RemarkMath from "remark-math";
import RemarkBreaks from "remark-breaks";
import RehypeKatex from "rehype-katex";
import RemarkGfm from "remark-gfm";
import RehypeHighlight from "rehype-highlight";
import { useRef, useState, RefObject, useEffect, useMemo } from "react";
import { copyToClipboard, useWindowSize } from "../utils";
import mermaid from "mermaid";
import Locale from "../locales";
import LoadingIcon from "../icons/three-dots.svg";
import ReloadButtonIcon from "../icons/reload.svg";
import React from "react";
import { useDebouncedCallback } from "use-debounce";
import { showImageModal, FullScreen } from "./ui-lib";
import {
  ArtifactsShareButton,
  HTMLPreview,
  HTMLPreviewHander,
} from "./artifacts";
import { useChatStore } from "../store";
import { IconButton } from "./button";

import { useAppConfig } from "../store/config";
import clsx from "clsx";

import { ChartComponentProps, extractDataFromText } from "./chart/getChartData";
import { ChartComponent } from "./chart";
import styles from "./chat.module.scss";
import { Collapse } from "antd";
import { useDebounce } from "../utils/hooks";
import { createDeepThink } from "../utils/deepThink";

export function Mermaid(props: { code: string }) {
  const ref = useRef<HTMLDivElement>(null);
  const [hasError, setHasError] = useState(false);

  useEffect(() => {
    if (props.code && ref.current) {
      mermaid
        .run({
          nodes: [ref.current],
          suppressErrors: true,
        })
        .catch((e) => {
          setHasError(true);
          console.error("[Mermaid] ", e.message);
        });
    }
    // eslint-disable-next-line react-hooks/exhaustive-deps
  }, [props.code]);

  function viewSvgInNewWindow() {
    const svg = ref.current?.querySelector("svg");
    if (!svg) return;
    const text = new XMLSerializer().serializeToString(svg);
    const blob = new Blob([text], { type: "image/svg+xml" });
    showImageModal(URL.createObjectURL(blob));
  }

  if (hasError) {
    return null;
  }

  return (
    <div
      className={clsx("no-dark", "mermaid")}
      style={{
        cursor: "pointer",
        overflow: "auto",
      }}
      ref={ref}
      onClick={() => viewSvgInNewWindow()}
    >
      {props.code}
    </div>
  );
}

export function PreCode(props: { children: any }) {
  const ref = useRef<HTMLPreElement>(null);
  const previewRef = useRef<HTMLPreviewHander>(null);
  const [mermaidCode, setMermaidCode] = useState("");
  const [htmlCode, setHtmlCode] = useState("");
  const { height } = useWindowSize();
  const chatStore = useChatStore();
  const session = chatStore.currentSession();
  const [chartData, setChartData] = useState<ChartComponentProps | null>(null);

  const renderArtifacts = useDebouncedCallback(() => {
    if (!ref.current) return;
    const mermaidDom = ref.current.querySelector("code.language-mermaid");
    if (mermaidDom) {
      setMermaidCode((mermaidDom as HTMLElement).innerText);
    }
    const htmlDom = ref.current.querySelector("code.language-html");
    const refText = ref.current.querySelector("code")?.innerText;
    if (htmlDom) {
      setHtmlCode((htmlDom as HTMLElement).innerText);
    } else if (
      refText?.startsWith("<!DOCTYPE") ||
      refText?.startsWith("<svg") ||
      refText?.startsWith("<?xml")
    ) {
      setHtmlCode(refText);
    }
  }, 600);

  const config = useAppConfig();
  const enableArtifacts =
    session.mask?.enableArtifacts !== false && config.enableArtifacts;
  //Wrap the paragraph for plain-text
  useEffect(() => {
    if (ref.current) {
      const textContent = ref.current.innerText;

      textContent &&
        extractDataFromText(textContent) &&
        setChartData(extractDataFromText(textContent));

      const codeElements = ref.current.querySelectorAll(
        "code",
      ) as NodeListOf<HTMLElement>;
      const wrapLanguages = [
        "",
        "md",
        "markdown",
        "text",
        "txt",
        "plaintext",
        "tex",
        "latex",
      ];
      codeElements.forEach((codeElement) => {
        let languageClass = codeElement.className.match(/language-(\w+)/);
        let name = languageClass ? languageClass[1] : "";
        if (wrapLanguages.includes(name)) {
          codeElement.style.whiteSpace = "pre-wrap";
        }
      });
      setTimeout(renderArtifacts, 1);
    }
  }, []);

  return (
    <>
      {mermaidCode.length > 0 && (
        <Mermaid code={mermaidCode} key={mermaidCode} />
      )}
      {htmlCode.length > 0 && enableArtifacts && (
        <FullScreen className="no-dark html" right={70}>
          <ArtifactsShareButton
            style={{ position: "absolute", right: 20, top: 10 }}
            getCode={() => htmlCode}
          />
          <IconButton
            style={{ position: "absolute", right: 120, top: 10 }}
            bordered
            icon={<ReloadButtonIcon />}
            shadow
            onClick={() => previewRef.current?.reload()}
          />
          <HTMLPreview
            ref={previewRef}
            code={htmlCode}
            autoHeight={!document.fullscreenElement}
            height={!document.fullscreenElement ? 600 : height}
          />
        </FullScreen>
      )}
      {chartData && <ChartComponent data={chartData.data} />}
      <pre ref={ref}>
        <span
          className="copy-code-button"
          onClick={() => {
            if (ref.current) {
              copyToClipboard(
                ref.current.querySelector("code")?.innerText ?? "",
              );
            }
          }}
        ></span>
        {props.children}
      </pre>
    </>
  );
}

function CustomCode(props: { children: any; className?: string }) {
  const chatStore = useChatStore();
  const session = chatStore.currentSession();
  const config = useAppConfig();
  const enableCodeFold =
    session.mask?.enableCodeFold !== false && config.enableCodeFold;

  const ref = useRef<HTMLPreElement>(null);
  const [collapsed, setCollapsed] = useState(true);
  const [showToggle, setShowToggle] = useState(false);

  useEffect(() => {
    if (ref.current) {
      const codeHeight = ref.current.scrollHeight;
      setShowToggle(codeHeight > 400);
      ref.current.scrollTop = ref.current.scrollHeight;
    }
  }, [props.children]);

  const toggleCollapsed = () => {
    setCollapsed((collapsed) => !collapsed);
  };
  const renderShowMoreButton = () => {
    if (showToggle && enableCodeFold && collapsed) {
      return (
        <div
          className={clsx("show-hide-button", {
            collapsed,
            expanded: !collapsed,
          })}
        >
          <button onClick={toggleCollapsed}>{Locale.NewChat.More}</button>
        </div>
      );
    }
    return null;
  };
  return (
    <>
      <code
        className={clsx(props?.className)}
        ref={ref}
        style={{
          maxHeight: enableCodeFold && collapsed ? "400px" : "none",
          overflowY: "hidden",
        }}
      >
        {props.children}
      </code>

      {renderShowMoreButton()}
    </>
  );
}

function escapeBrackets(text: string) {
  const pattern =
    /(```[\s\S]*?```|`.*?`)|\\\[([\s\S]*?[^\\])\\\]|\\\((.*?)\\\)/g;
  return text.replace(
    pattern,
    (match, codeBlock, squareBracket, roundBracket) => {
      if (codeBlock) {
        return codeBlock;
      } else if (squareBracket) {
        return `$$${squareBracket}$$`;
      } else if (roundBracket) {
        return `$${roundBracket}$`;
      }
      return match;
    },
  );
}

function tryWrapHtmlCode(text: string) {
  // try add wrap html code (fixed: html codeblock include 2 newline)
  // ignore embed codeblock
  if (text.includes("```")) {
    return text;
  }
  return text
    .replace(
      /([`]*?)(\w*?)([\n\r]*?)(<!DOCTYPE html>)/g,
      (match, quoteStart, lang, newLine, doctype) => {
        return !quoteStart ? "\n```html\n" + doctype : match;
      },
    )
    .replace(
      /(<\/body>)([\r\n\s]*?)(<\/html>)([\n\r]*)([`]*)([\n\r]*?)/g,
      (match, bodyEnd, space, htmlEnd, newLine, quoteEnd) => {
        return !quoteEnd ? bodyEnd + space + htmlEnd + "\n```\n" : match;
      },
    );
}

function _MarkDownContent(props: { content: string; isUser: boolean }) {
  const escapedContent = useMemo(() => {
    return tryWrapHtmlCode(escapeBrackets(props.content));
  }, [props.content]);

  type TextSegment = {
    type: "text";
    content: string;
    start: number;
    end: number;
  };

  type ElementSegment = {
    type: "element";
    node: React.ReactElement;
  };

  type Segment = TextSegment | ElementSegment;

  function parseChildrenToSegments(children: React.ReactNode): Segment[] {
    const segments: Segment[] = [];
    let textBuffer = "";
    let globalIndex = 0;

    React.Children.forEach(children, (child) => {
      if (typeof child === "string") {
        // 合并连续文本节点
        textBuffer += child;
        globalIndex += child.length;
      } else {
        // 遇到非文本节点时，先提交缓冲的文本
        if (textBuffer) {
          segments.push({
            type: "text",
            content: textBuffer,
            start: globalIndex - textBuffer.length,
            end: globalIndex,
          });
          textBuffer = "";
        }
        // 记录元素节点
        segments.push({
          type: "element",
          node: child as React.ReactElement,
        });
      }
    });

    // 提交最后缓冲的文本
    if (textBuffer) {
      segments.push({
        type: "text",
        content: textBuffer,
        start: globalIndex - textBuffer.length,
        end: globalIndex,
      });
    }

    return segments;
  }

  function CustomParagraph({ children }: { children: React.ReactNode }) {
    const segments = parseChildrenToSegments(children);
    const fullText = segments
      .filter((s): s is TextSegment => s.type === "text")
      .map((s) => s.content)
      .join("");
    const regex = /'''filedata\s+([^\n]+)[\s\S]*?'''filedata/g;
    const matches: Array<{
      fileName: string;
      start: number;
      end: number;
    }> = [];

    let match;
    while ((match = regex.exec(fullText)) !== null) {
      matches.push({
        fileName: match[1].trim(),
        start: match.index,
        end: match.index + match[0].length,
      });
    }
    const newChildren: React.ReactNode[] = [];
    let lastPos = 0;
    let currentSegmentIndex = 0;
    matches.forEach((match) => {
      while (currentSegmentIndex < segments.length) {
        const segment = segments[currentSegmentIndex];
        if (segment.type === "element") {
          newChildren.push(segment.node);
          currentSegmentIndex++;
          continue;
        }
        const segmentEnd = segment.end;
        if (segmentEnd <= match.start) {
          newChildren.push(segment.content.slice(lastPos - segment.start));
          lastPos = segmentEnd;
          currentSegmentIndex++;
        } else {
          break;
        }
      }
      newChildren.push(
        <div key={`file-${match.start}`} className={styles["message-file"]}>
          <span>{match.fileName}</span>
        </div>,
      );
      lastPos = match.end;
    });
    while (currentSegmentIndex < segments.length) {
      const segment = segments[currentSegmentIndex];
      if (segment.type === "element") {
        newChildren.push(segment.node);
      } else {
        newChildren.push(segment.content.slice(lastPos - segment.start));
      }
      currentSegmentIndex++;
    }

    return <p dir="auto">{newChildren}</p>;
  }

  const CollapsibleBlockquote = ({
    children,
  }: {
    children: React.ReactNode;
  }) => {
    const [isFinalized, setIsFinalized] = useState(false);
    const [lastUpdated, setLastUpdated] = useState(Date.now());
    const contentString = React.Children.toArray(children).join("");

    // 使用防抖检测内容是否稳定
    const debouncedCheck = useDebounce(() => {
      if (Date.now() - lastUpdated > 2000) {
        // 2秒内无更新视为最终
        setIsFinalized(true);
      }
    }, 2000);
    useEffect(() => {
      setIsFinalized(false);
      setLastUpdated(Date.now());
      debouncedCheck();
    }, [contentString]);

    return (
      <Collapse
        defaultActiveKey={["1"]}
        size="small"
        items={[
          {
            key: "1",
            label: isFinalized ? "已深度思考" : "思考中...",
            children: (
              <blockquote
                style={{
                  padding: "10px",
                  borderLeft: "2px solid #ccc",
                  transition: "all 0.3s ease",
                  fontSize: "12px",
                }}
              >
                {children}
              </blockquote>
            ),
          },
        ]}
      />
    );
  };

  return (
    <ReactMarkdown
      remarkPlugins={[RemarkMath, RemarkGfm, RemarkBreaks]}
      rehypePlugins={[
        RehypeKatex,
        [
          RehypeHighlight,
          {
            detect: false,
            ignoreMissing: true,
          },
        ],
      ]}
      components={{
        pre: PreCode,
        code: CustomCode,
        p: props.isUser
          ? CustomParagraph
          : (pProps) => <p {...pProps} dir="auto" />,
        blockquote: ({ children }) =>
          props.isUser ? (
            <blockquote>{children}</blockquote>
          ) : (
            <CollapsibleBlockquote>{children}</CollapsibleBlockquote>
          ),
        a: (aProps) => {
          const href = aProps.href || "";
          if (/\.(aac|mp3|opus|wav)$/.test(href)) {
            return (
              <figure>
                <audio controls src={href}></audio>
              </figure>
            );
          }
          if (/\.(3gp|3g2|webm|ogv|mpeg|mp4|avi)$/.test(href)) {
            return (
              <video controls width="99.9%">
                <source src={href} />
              </video>
            );
          }
          const isInternal = /^\/#/i.test(href);
          const target = isInternal ? "_self" : aProps.target ?? "_blank";
          return <a {...aProps} target={target} />;
        },
      }}
    >
      {escapedContent}
    </ReactMarkdown>
  );
}

export const MarkdownContent = React.memo(_MarkDownContent);

export function Markdown(
  props: {
    content: string;
    loading?: boolean;
    fontSize?: number;
    fontFamily?: string;
    parentRef?: RefObject<HTMLDivElement>;
    defaultShow?: boolean;
    isUser: boolean;
  } & React.DOMAttributes<HTMLDivElement>,
) {
  const mdRef = useRef<HTMLDivElement>(null);

  return (
    <div
      className="markdown-body"
      style={{
        fontSize: `${props.fontSize ?? 14}px`,
        fontFamily: props.fontFamily || "inherit",
      }}
      ref={mdRef}
      onContextMenu={props.onContextMenu}
      onDoubleClickCapture={props.onDoubleClickCapture}
      dir="auto"
    >
      {props.loading ? (
        <LoadingIcon />
      ) : (
        <MarkdownContent
          content={createDeepThink(props.content)}
          isUser={props.isUser}
        />
      )}
    </div>
  );
}