modsetter--surfsense/surfsense_web/components/assistant-ui/markdown-code-block.tsx

"use client";

import { CheckIcon, CopyIcon } from "lucide-react";
import type { CSSProperties } from "react";
import { memo, useEffect, useState } from "react";
import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
import { materialDark, materialLight } from "react-syntax-highlighter/dist/esm/styles/prism";

import { MermaidDiagram } from "@/components/assistant-ui/mermaid-diagram";
import { Button } from "@/components/ui/button";
import { cn, copyToClipboard } from "@/lib/utils";

type MarkdownCodeBlockProps = {
	className?: string;
	language: string;
	codeText: string;
	isDarkMode: boolean;
};

function stripThemeBackgrounds(
	theme: Record<string, CSSProperties>
): Record<string, CSSProperties> {
	const cleaned: Record<string, CSSProperties> = {};
	for (const key of Object.keys(theme)) {
		const { background, backgroundColor, ...rest } = theme[key] as CSSProperties & {
			background?: string;
			backgroundColor?: string;
		};
		cleaned[key] = rest;
	}
	return cleaned;
}

const cleanMaterialDark = stripThemeBackgrounds(materialDark);
const cleanMaterialLight = stripThemeBackgrounds(materialLight);

function MarkdownCodeBlockComponent({
	className,
	language,
	codeText,
	isDarkMode,
}: MarkdownCodeBlockProps) {
	const [hasCopied, setHasCopied] = useState(false);
	const normalizedLanguage = language.toLowerCase();

	useEffect(() => {
		if (!hasCopied) return;
		const timer = setTimeout(() => setHasCopied(false), 2000);
		return () => clearTimeout(timer);
	}, [hasCopied]);

	const codeBlock = (
		<div className="mt-4 overflow-hidden rounded-md bg-accent">
			<div className="flex items-center justify-between gap-4 px-4 py-2 font-semibold text-muted-foreground text-sm">
				<span className="lowercase text-xs">{language}</span>
				<Button
					variant="ghost"
					size="sm"
					className="h-8 w-8 p-0"
					type="button"
					onClick={async () => {
						const ok = await copyToClipboard(codeText);
						if (ok) setHasCopied(true);
					}}
					aria-label={hasCopied ? "Copied code" : "Copy code"}
				>
					<span className="sr-only">Copy</span>
					{hasCopied ? <CheckIcon className="!size-3" /> : <CopyIcon className="!size-3" />}
				</Button>
			</div>

			<SyntaxHighlighter
				style={isDarkMode ? cleanMaterialDark : cleanMaterialLight}
				language={language}
				PreTag="div"
				customStyle={{ margin: 0, background: "transparent" }}
				className={cn(className)}
			>
				{codeText}
			</SyntaxHighlighter>
		</div>
	);

	if (normalizedLanguage === "mermaid") {
		return <MermaidDiagram source={codeText} isDarkMode={isDarkMode} fallback={codeBlock} />;
	}

	return codeBlock;
}

export const MarkdownCodeBlock = memo(MarkdownCodeBlockComponent);