lobe-chat

1
'use client';
2

3
import { Pre, PreSingleLine } from '@lobehub/ui/mdx';
4
import { FC, PropsWithChildren } from 'react';
5

6
const countLines = (str: string): number => {
7
  const regex = /\n/g;
8
  const matches = str.match(regex);
9
  return matches ? matches.length : 1;
10
};
11

12
const useCode = (raw: any) => {
13
  if (!raw) return;
14

15
  const { children, className } = raw.props;
16

17
  if (!children) return;
18

19
  const content = (Array.isArray(children) ? (children[0] as string) : children).trim();
20

21
  const lang = className?.replace('language-', '') || 'txt';
22

23
  const isSingleLine = countLines(content) <= 1 && content.length <= 32;
24

25
  return {
26
    content,
27
    isSingleLine,
28
    lang,
29
  };
30
};
31

32
const CodeBlock: FC<PropsWithChildren> = ({ children }) => {
33
  const code = useCode(children);
34

35
  if (!code) return;
36

37
  if (code.isSingleLine) return <PreSingleLine language={code.lang}>{code.content}</PreSingleLine>;
38

39
  return (
40
    <Pre allowChangeLanguage={false} fullFeatured language={code.lang}>
41
      {code.content}
42
    </Pre>
43
  );
44
};
45

46
export default CodeBlock;
47