deerflow2/frontend/src/components/workspace/messages/markdown-content.tsx

"use client";

import { DownloadIcon } from "lucide-react";
import { useCallback, useLayoutEffect, useMemo, useRef, useState } from "react";
import type {
  AnchorHTMLAttributes,
  ComponentPropsWithoutRef,
  ReactNode,
} from "react";

import {
  MessageResponse,
  type MessageResponseProps,
} from "@/components/ai-elements/message";
import {
  DropdownMenu,
  DropdownMenuContent,
  DropdownMenuItem,
  DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { useI18n } from "@/core/i18n/hooks";
import { streamdownPlugins } from "@/core/streamdown";
import { CopyButton } from "@/components/workspace/copy-button";
import { Tooltip } from "@/components/workspace/tooltip";
import { cn } from "@/lib/utils";

import { CitationLink } from "../citations/citation-link";

function isExternalUrl(href: string | undefined): boolean {
  return !!href && /^https?:\/\//.test(href);
}

export type MarkdownContentProps = {
  content: string;
  isLoading: boolean;
  rehypePlugins: MessageResponseProps["rehypePlugins"];
  className?: string;
  remarkPlugins?: MessageResponseProps["remarkPlugins"];
  components?: MessageResponseProps["components"];
};

type TableData = {
  headers: string[];
  rows: string[][];
};

function parseTableData(table: HTMLTableElement): TableData {
  const headers = Array.from(table.querySelectorAll("thead th")).map((cell) =>
    (cell.textContent ?? "").trim(),
  );
  const rows = Array.from(table.querySelectorAll("tbody tr")).map((row) =>
    Array.from(row.querySelectorAll("td")).map((cell) =>
      (cell.textContent ?? "").trim(),
    ),
  );
  return { headers, rows };
}

function toMarkdownTable(data: TableData): string {
  if (data.headers.length === 0) return "";
  const headerLine = `| ${data.headers.join(" | ")} |`;
  const dividerLine = `| ${data.headers.map(() => "---").join(" | ")} |`;
  const rowLines = data.rows.map((row) => `| ${row.join(" | ")} |`);
  return [headerLine, dividerLine, ...rowLines].join("\n");
}

function escapeCsvCell(cell: string): string {
  const normalized = cell.replace(/\r\n/g, "\n");
  if (!/["\n,]/.test(normalized)) return normalized;
  return `"${normalized.replace(/"/g, '""')}"`;
}

function toCsvTable(data: TableData): string {
  if (data.headers.length === 0) return "";
  const headerLine = data.headers.map(escapeCsvCell).join(",");
  const rowLines = data.rows.map((row) => row.map(escapeCsvCell).join(","));
  return [headerLine, ...rowLines].join("\r\n");
}

function downloadTextFile(content: string, filename: string, mimeType: string) {
  const blob = new Blob(["\uFEFF", content], {
    type: mimeType,
  });
  const url = URL.createObjectURL(blob);
  const anchor = document.createElement("a");
  anchor.href = url;
  anchor.download = filename;
  anchor.click();
  URL.revokeObjectURL(url);
}

export function MarkdownTable({
  className,
  children,
  copyLabel: _copyLabel,
  downloadLabel,
  ...props
}: ComponentPropsWithoutRef<"table"> & {
  copyLabel: string;
  downloadLabel: string;
}) {
  const tableRef = useRef<HTMLTableElement>(null);
  const [, forceUpdate] = useState(0);

  // 首次 mount 后 tableRef 才被赋值，用 useLayoutEffect 在 paint 前强制刷新
  useLayoutEffect(() => {
    forceUpdate((n) => n + 1);
  }, []);

  // 在 render 阶段直接从 DOM ref 计算，不依赖 effect 异步更新
  // tableRef 在上一次渲染的 commit 阶段已设置，本次渲染可用
  const clipboardData = (() => {
    const table = tableRef.current;
    if (!table) return "";
    const data = parseTableData(table);
    if (!data) return "";
    return toMarkdownTable(data);
  })();

  const handleDownload = useCallback(() => {
    const table = tableRef.current;
    if (!table) return;
    const data = parseTableData(table);
    if (!data) return;
    const csv = toCsvTable(data);
    if (!csv) return;
    downloadTextFile(csv, "table.csv", "text/csv;charset=utf-8");
  }, []);

  const handleDownloadMarkdown = useCallback(() => {
    const table = tableRef.current;
    if (!table) return;
    const data = parseTableData(table);
    if (!data) return;
    const markdown = toMarkdownTable(data);
    if (!markdown) return;
    downloadTextFile(markdown, "table.md", "text/markdown;charset=utf-8");
  }, []);

  return (
    <div
      className="my-4 flex flex-col space-y-2"
      data-streamdown="table-wrapper"
    >
      <div className="flex items-center justify-end gap-1">
        <CopyButton className="text-muted-foreground hover:bg-transparent hover:text-foreground cursor-pointer p-1 transition-all" clipboardData={clipboardData} />
        <DropdownMenu>
          <Tooltip content={downloadLabel}>
            <DropdownMenuTrigger asChild>
              <button
                className="h-[32px] w-[32px] text-muted-foreground hover:text-foreground cursor-pointer p-1 transition-all"
                type="button"
              >
                <DownloadIcon size={16} />
              </button>
            </DropdownMenuTrigger>
          </Tooltip>
          <DropdownMenuContent align="end" className="min-w-[140px] p-1">
            <DropdownMenuItem onSelect={handleDownload}>CSV</DropdownMenuItem>
            <DropdownMenuItem onSelect={handleDownloadMarkdown}>
              Markdown
            </DropdownMenuItem>
          </DropdownMenuContent>
        </DropdownMenu>
      </div>
      <div className="overflow-x-auto">
        <table
          ref={tableRef}
          className={cn(
            "border-border w-full border-collapse border",
            className,
          )}
          data-streamdown="table"
          {...props}
        >
          {children}
        </table>
      </div>
    </div>
  );
}

/** Renders markdown content. */
export function MarkdownContent({
  content,
  isLoading,
  rehypePlugins,
  className,
  remarkPlugins = streamdownPlugins.remarkPlugins,
  components: componentsFromProps,
}: MarkdownContentProps) {
  const { t } = useI18n();

  const components = useMemo(() => {
    return {
      a: (props: AnchorHTMLAttributes<HTMLAnchorElement>) => {
        if (typeof props.children === "string") {
          const match = /^citation:(.+)$/.exec(props.children);
          if (match) {
            const [, text] = match;
            return <CitationLink {...props}>{text}</CitationLink>;
          }
        }
        const { className, target, rel, ...rest } = props;
        const external = isExternalUrl(props.href);
        return (
          <a
            {...rest}
            className={cn(
              "text-primary decoration-primary/30 hover:decoration-primary/60 underline underline-offset-2 transition-colors",
              className,
            )}
            target={target ?? (external ? "_blank" : undefined)}
            rel={rel ?? (external ? "noopener noreferrer" : undefined)}
          />
        );
      },
      table: ({
        children,
        className,
        ...props
      }: ComponentPropsWithoutRef<"table"> & { children?: ReactNode }) => (
        <MarkdownTable
          className={className}
          copyLabel={t.clipboard.copyToClipboard}
          downloadLabel={t.common.download}
          {...props}
        >
          {children}
        </MarkdownTable>
      ),
      ...componentsFromProps,
    };
  }, [
    componentsFromProps,
    isLoading,
    t.clipboard.copyToClipboard,
    t.common.download,
  ]);

  if (!content) return null;

  return (
    <MessageResponse
      className={className}
      isAnimating={isLoading}
      controls={{ table: false }}
      parseIncompleteMarkdown={!isLoading}
      remarkPlugins={remarkPlugins}
      rehypePlugins={rehypePlugins}
      components={components}
    >
      {content}
    </MessageResponse>
  );
}