Alfred/src/web/features/chat/parts/CodeBlockWithCopy.tsx

import { CopyOutlined } from "@ant-design/icons";
import { CodeHighlighter } from "@ant-design/x";
import { App, Button, Flex, Typography } from "antd";
import React from "react";
import { oneDark, oneLight } from "react-syntax-highlighter/dist/esm/styles/prism";

import { useIsDark } from "../../../shared/hooks/use-is-dark";

type SyntaxTheme = Record<string, Record<string, null | number | string>>;

const customOneDark: SyntaxTheme = {
  ...(oneDark as SyntaxTheme),
  'pre[class*="language-"]': {
    ...(oneDark as SyntaxTheme)['pre[class*="language-"]'],
    margin: 0,
  },
};

const customOneLight: SyntaxTheme = {
  ...(oneLight as SyntaxTheme),
  'pre[class*="language-"]': {
    ...(oneLight as SyntaxTheme)['pre[class*="language-"]'],
    margin: 0,
  },
};

interface CodeBlockWithCopyProps {
  block?: boolean;
  children?: React.ReactNode;
  className?: string;
  lang?: string;
  streamStatus?: "done" | "loading";
}

export function CodeBlockWithCopy({ block, children, className, lang }: CodeBlockWithCopyProps) {
  const { message } = App.useApp();
  const isDark = useIsDark();

  if (!block) {
    return <code className={className}>{children}</code>;
  }

  const codeText = extractText(children);
  const displayLang = lang ?? "plaintext";

  const handleCopy = () => {
    void navigator.clipboard.writeText(codeText).then(() => {
      void message.success("已复制");
    });
  };

  const header = (
    <Flex align="center" justify="space-between" style={{ padding: "0 4px" }}>
      <Typography.Text style={{ color: "var(--ant-color-text-quaternary)", fontSize: 12 }}>
        {displayLang}
      </Typography.Text>
      <Button
        icon={<CopyOutlined />}
        onClick={handleCopy}
        size="small"
        style={{ color: "var(--ant-color-text-quaternary)" }}
        type="text"
      />
    </Flex>
  );

  return (
    <CodeHighlighter
      header={header}
      highlightProps={{ style: (isDark ? customOneDark : customOneLight) as React.CSSProperties }}
      lang={displayLang}
    >
      {codeText}
    </CodeHighlighter>
  );
}

function extractText(children: React.ReactNode): string {
  return React.Children.toArray(children)
    .map((child) => (typeof child === "string" ? child : ""))
    .join("");
}