jiang13-bbs/frontend/components/markdown/MdCodeBlock.tsx

"use client";

import {
  useMemo,
  useRef,
  useState,
  useEffect,
  type HTMLAttributes,
  type ReactNode,
} from "react";
import { Check, Copy, ChevronDown, ChevronUp } from "lucide-react";
import { toast } from "@/lib/toast";

const COLLAPSE_LINE_THRESHOLD = 5;
/** 约 5 行代码高度（含行高） */
const COLLAPSED_MAX_HEIGHT = "7.75rem";

type MdCodeBlockProps = HTMLAttributes<HTMLPreElement> & {
  children?: ReactNode;
};

function extractText(node: ReactNode): string {
  if (node == null || typeof node === "boolean") return "";
  if (typeof node === "string" || typeof node === "number") return String(node);
  if (Array.isArray(node)) return node.map(extractText).join("");
  if (typeof node === "object" && "props" in node) {
    const el = node as { props?: { children?: ReactNode } };
    return extractText(el.props?.children);
  }
  return "";
}

function countLines(pre: HTMLElement | null, fallbackText: string): number {
  if (pre) {
    const lineEls = pre.querySelectorAll(".line, [data-line]");
    if (lineEls.length > 0) return lineEls.length;
  }
  const text = fallbackText.replace(/\n$/, "");
  if (!text) return 0;
  return text.split("\n").length;
}

function readLanguage(
  el: HTMLElement | null,
  className?: string,
  dataLanguage?: string
): string {
  const fromData = dataLanguage || el?.getAttribute("data-language") || "";
  if (fromData) return fromData;
  const cls = [className, el?.className].filter(Boolean).join(" ");
  const m = cls.match(/language-([\w+-]+)/);
  return m?.[1] || "";
}

/** Markdown 代码块：语言标签、复制、>5 行默认折叠；客户端可补 Shiki 双主题 */
export default function MdCodeBlock({ children, className, ...rest }: MdCodeBlockProps) {
  const hostRef = useRef<HTMLDivElement>(null);

  const plainText = useMemo(() => extractText(children), [children]);
  const dataLanguage =
    typeof (rest as Record<string, unknown>)["data-language"] === "string"
      ? String((rest as Record<string, unknown>)["data-language"])
      : undefined;

  const approxLines = useMemo(() => countLines(null, plainText), [plainText]);
  const alreadyShiki =
    typeof className === "string" && className.includes("shiki");

  const [copied, setCopied] = useState(false);
  const [expanded, setExpanded] = useState(false);
  const [lineCount, setLineCount] = useState(approxLines);
  const [language, setLanguage] = useState(() =>
    readLanguage(null, className, dataLanguage)
  );
  const [clientHtml, setClientHtml] = useState<string | null>(null);

  useEffect(() => {
    const host = hostRef.current;
    const pre = host?.querySelector("pre") ?? null;
    setLineCount(countLines(pre, plainText));
    setLanguage(readLanguage(pre, className, dataLanguage));
  }, [plainText, className, dataLanguage, clientHtml]);

  // 服务端已着色时跳过；客户端预览动态补双主题
  useEffect(() => {
    if (alreadyShiki || !plainText.trim()) return;
    const lang = readLanguage(null, className, dataLanguage) || "plaintext";
    let cancelled = false;
    (async () => {
      try {
        const { codeToHtml } = await import("shiki");
        const html = await codeToHtml(plainText.replace(/\n$/, ""), {
          lang: lang === "text" ? "plaintext" : lang,
          themes: {
            light: "github-light",
            dark: "github-dark-dimmed",
          },
          defaultColor: false,
        });
        if (!cancelled) setClientHtml(html);
      } catch {
        // 未知语言：保持纯文本
      }
    })();
    return () => {
      cancelled = true;
    };
  }, [plainText, className, dataLanguage, alreadyShiki]);

  const effectiveLines = lineCount > 0 ? lineCount : approxLines;
  const canCollapse = effectiveLines > COLLAPSE_LINE_THRESHOLD;
  const collapsed = canCollapse && !expanded;

  const onCopy = async () => {
    const pre = hostRef.current?.querySelector("pre");
    const text =
      pre?.innerText?.replace(/\n$/, "") || plainText.replace(/\n$/, "");
    if (!text) return;
    try {
      await navigator.clipboard.writeText(text);
      setCopied(true);
      window.setTimeout(() => setCopied(false), 1600);
    } catch {
      toast("复制失败，请手动选择文本");
    }
  };

  const preClass = ["md-code-pre", className].filter(Boolean).join(" ");

  return (
    <div className="md-code-block">
      <div className="md-code-toolbar">
        <span className="md-code-lang">{language || "text"}</span>
        <button
          type="button"
          className="md-code-copy"
          onClick={onCopy}
          aria-label={copied ? "已复制" : "复制代码"}
        >
          {copied ? <Check size={14} /> : <Copy size={14} />}
          <span>{copied ? "已复制" : "复制"}</span>
        </button>
      </div>
      <div
        ref={hostRef}
        className={`md-code-body${collapsed ? " md-code-collapsed" : ""}`}
        style={collapsed ? { maxHeight: COLLAPSED_MAX_HEIGHT } : undefined}
      >
        {clientHtml ? (
          <div
            className="md-code-shiki-host"
            dangerouslySetInnerHTML={{ __html: clientHtml }}
          />
        ) : (
          <pre className={preClass} {...rest}>
            {children}
          </pre>
        )}
        {collapsed && <div className="md-code-fade" aria-hidden="true" />}
      </div>
      {canCollapse && (
        <button
          type="button"
          className="md-code-fold"
          onClick={() => setExpanded((v) => !v)}
        >
          {expanded ? (
            <>
              <ChevronUp size={14} />
              收起
            </>
          ) : (
            <>
              <ChevronDown size={14} />
              展开全部（{effectiveLines} 行）
            </>
          )}
        </button>
      )}
    </div>
  );
}