TJWaterFrontend_Refine/src/components/chat/AgentMarkdownBlock.tsx

"use client";

import React from "react";
import type { Element, Root, RootContent, Text } from "hast";
import ReactMarkdown, { type Components } from "react-markdown";
import remarkGfm from "remark-gfm";

import markdownStyles from "./GlobalChatboxMarkdown.module.css";

export const normalizeClipboardText = (value: string) => value.replace(/\s+$/u, "");

const isTextNode = (node: RootContent): node is Text => node.type === "text";

const isElementNode = (node: RootContent): node is Element => node.type === "element";

const createFadeSpan = (value: string, fadeKey: string): Element => ({
  type: "element",
  tagName: "span",
  properties: {
    className: [markdownStyles.streamFade],
    dataStreamFadeKey: fadeKey,
    dataStreamRevealLength: value.length,
  },
  children: [{ type: "text", value }],
});

const splitTextTail = (value: string, tailLength: number) => {
  const codePoints = Array.from(value);
  const stableText = codePoints.slice(0, -tailLength).join("");
  const animatedText = codePoints.slice(-tailLength).join("");
  return { stableText, animatedText };
};

const createStreamFadePlugin = (fadeLength: number, fadeKey: string) => {
  return () => (tree: Root) => {
    let remainingFadeLength = fadeLength;

    const visitChildren = (parent: Element | Root) => {
      const nextChildren: RootContent[] = [];

      for (let index = parent.children.length - 1; index >= 0; index -= 1) {
        const child = (parent.children as RootContent[])[index];
        if (isTextNode(child)) {
          if (!child.value.trim()) {
            nextChildren.unshift(child);
            continue;
          }

          if (remainingFadeLength <= 0) {
            nextChildren.unshift(child);
            continue;
          }

          const textLength = Array.from(child.value).length;
          const tailLength = Math.min(textLength, remainingFadeLength);
          const { stableText, animatedText } = splitTextTail(child.value, tailLength);
          remainingFadeLength -= tailLength;

          if (animatedText) {
            nextChildren.unshift(createFadeSpan(animatedText, fadeKey));
          }
          if (stableText) {
            nextChildren.unshift({ ...child, value: stableText });
          }
          continue;
        }

        if (isElementNode(child)) {
          visitChildren(child);
        }

        nextChildren.unshift(child);
      }

      parent.children = nextChildren as typeof parent.children;
    };

    visitChildren(tree);
  };
};

const StreamFadeSpan: Components["span"] = ({ node, children, ...props }) => {
  const ref = React.useRef<HTMLSpanElement>(null);
  const fadeKeyValue = node?.properties?.dataStreamFadeKey;
  const fadeKey = typeof fadeKeyValue === "string" ? fadeKeyValue : undefined;
  const revealLengthValue = node?.properties?.dataStreamRevealLength;
  const revealLength =
    typeof revealLengthValue === "number"
      ? revealLengthValue
      : typeof revealLengthValue === "string"
        ? Number.parseInt(revealLengthValue, 10)
        : 0;

  React.useLayoutEffect(() => {
    if (!fadeKey) return;

    const element = ref.current;
    if (!element) return;
    if (window.matchMedia?.("(prefers-reduced-motion: reduce)").matches) return;

    const duration = Math.min(260, Math.max(120, revealLength * 14));
    const animation = element.animate(
      [
        { clipPath: "inset(0 100% 0 0)", opacity: 0.46 },
        { clipPath: "inset(0 0% 0 0)", opacity: 1 },
      ],
      {
        duration,
        easing: "cubic-bezier(0.16, 1, 0.3, 1)",
        fill: "both",
      },
    );

    return () => {
      animation.cancel();
    };
  }, [fadeKey, revealLength]);

  return (
    <span {...props} ref={ref}>
      {children}
    </span>
  );
};

const markdownComponents: Components = {
  span: StreamFadeSpan,
};

export const MarkdownBlock = ({
  children,
  streamFadeKey,
  streamFadeLength,
}: {
  children: string;
  streamFadeKey?: string;
  streamFadeLength?: number | null;
}) => {
  const handleCopy = React.useCallback((event: React.ClipboardEvent<HTMLDivElement>) => {
    const selectedText = window.getSelection()?.toString();
    if (!selectedText) return;

    event.preventDefault();
    event.clipboardData.setData("text/plain", normalizeClipboardText(selectedText));
  }, []);
  const rehypePlugins = React.useMemo(
    () =>
      typeof streamFadeLength === "number" && streamFadeLength > 0
        ? [createStreamFadePlugin(
            streamFadeLength,
            streamFadeKey ?? `stream-tail-${children.length}`,
          )]
        : [],
    [children.length, streamFadeKey, streamFadeLength],
  );

  return (
    <div className={markdownStyles.markdown} onCopy={handleCopy}>
      <ReactMarkdown
        components={markdownComponents}
        remarkPlugins={[remarkGfm]}
        rehypePlugins={rehypePlugins}
      >
        {children}
      </ReactMarkdown>
    </div>
  );
};