code syntax highlights

components/Chat/ChatMessage.tsx

@@ -1,6 +1,7 @@
 import { Message } from "@/types";
 import { FC } from "react";
 import ReactMarkdown from "react-markdown";
+import { CodeBlock } from "../Markdown/CodeBlock";
 
 interface Props {
   message: Message;
@@ -16,7 +17,30 @@ export const ChatMessage: FC<Props> = ({ message }) => {
         <div className="mr-4 font-bold min-w-[40px]">{message.role === "assistant" ? "AI:" : "You:"}</div>
 
         <div className="prose dark:prose-invert">
-          <ReactMarkdown>{message.content}</ReactMarkdown>
+          <ReactMarkdown
+            components={{
+              code({ node, inline, className, children, ...props }) {
+                const match = /language-(\w+)/.exec(className || "");
+                return !inline && match ? (
+                  <CodeBlock
+                    key={Math.random()}
+                    language={match[1]}
+                    value={String(children).replace(/\n$/, "")}
+                    {...props}
+                  />
+                ) : (
+                  <code
+                    className={className}
+                    {...props}
+                  >
+                    {children}
+                  </code>
+                );
+              }
+            }}
+          >
+            {message.content}
+          </ReactMarkdown>
         </div>
       </div>
     </div>

components/Markdown/CodeBlock.tsx

@@ -0,0 +1,40 @@
+import { FC, useState } from "react";
+import { Prism as SyntaxHighlighter } from "react-syntax-highlighter";
+import { tomorrow } from "react-syntax-highlighter/dist/cjs/styles/prism";
+
+interface Props {
+  language: string;
+  value: string;
+}
+
+export const CodeBlock: FC<Props> = ({ language, value }) => {
+  const [buttonText, setButtonText] = useState("Copy");
+
+  const copyToClipboard = () => {
+    navigator.clipboard.writeText(value).then(() => {
+      setButtonText("Copied!");
+
+      setTimeout(() => {
+        setButtonText("Copy");
+      }, 2000);
+    });
+  };
+
+  return (
+    <div className="relative">
+      <SyntaxHighlighter
+        language={language}
+        style={tomorrow}
+      >
+        {value}
+      </SyntaxHighlighter>
+
+      <button
+        className="absolute top-2 right-2 text-white bg-blue-600 py-1 px-2 rounded focus:outline-none hover:bg-blue-700"
+        onClick={copyToClipboard}
+      >
+        {buttonText}
+      </button>
+    </div>
+  );
+};