Code Block

registry/corr/components/code-block.tsx

"use client"

import * as React from "react"
import { CheckIcon, CopyIcon, TerminalIcon } from "lucide-react"

import { Button } from "@/components/ui/button"
import {
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
} from "@/components/ui/select"
import { cn } from "@/lib/utils"

export type CodeBlockLanguage =
  | "tsx"
  | "ts"
  | "jsx"
  | "js"
  | "json"
  | "css"
  | "bash"
  | "python"
  | "curl"
  | "text"

export type CodeBlockSnippet = {
  language: CodeBlockLanguage
  label?: string
  code: string
}

export type CodeBlockProps = {
  code?: string
  language?: CodeBlockLanguage
  snippets?: CodeBlockSnippet[]
  defaultLanguage?: CodeBlockLanguage
  label?: string
  showLineNumbers?: boolean
  allowLanguageSelection?: boolean
  className?: string
  contentClassName?: string
}

const languageLabels: Record<CodeBlockLanguage, string> = {
  tsx: "TSX",
  ts: "TypeScript",
  jsx: "JSX",
  js: "JavaScript",
  json: "JSON",
  css: "CSS",
  bash: "Shell",
  python: "Python",
  curl: "cURL",
  text: "Text",
}

const languageKeywords: Record<CodeBlockLanguage, string[]> = {
  tsx: [
    "import",
    "from",
    "export",
    "function",
    "return",
    "const",
    "let",
    "type",
    "interface",
    "extends",
    "satisfies",
    "as",
  ],
  ts: [
    "import",
    "from",
    "export",
    "function",
    "return",
    "const",
    "let",
    "type",
    "interface",
    "extends",
    "satisfies",
    "as",
  ],
  jsx: ["import", "from", "export", "function", "return", "const", "let"],
  js: ["import", "from", "export", "function", "return", "const", "let"],
  json: ["true", "false", "null"],
  css: [
    "display",
    "grid",
    "flex",
    "color",
    "background",
    "border",
    "padding",
    "margin",
  ],
  bash: ["npm", "npx", "pnpm", "bun", "cd", "cp", "ssh", "sudo"],
  python: [
    "import",
    "from",
    "def",
    "return",
    "if",
    "else",
    "for",
    "in",
    "with",
    "as",
    "True",
    "False",
    "None",
  ],
  curl: ["curl", "GET", "POST", "PUT", "PATCH", "DELETE"],
  text: [],
}

const tokenPattern =
  /(\/\/.*|#.*|"(?:\\.|[^"\\])*"|'(?:\\.|[^'\\])*'|`(?:\\.|[^`\\])*`|\b[A-Za-z_$][A-Za-z0-9_$-]*\b|\b\d+(?:\.\d+)?\b|[{}()[\].,;:<>/=+\-*|&!])/g

function getTokenClass(token: string, language: CodeBlockLanguage) {
  const keywords = languageKeywords[language] ?? []

  if (
    token.startsWith("//") ||
    (["bash", "python", "curl"].includes(language) && token.startsWith("#"))
  ) {
    return "text-muted-foreground"
  }

  if (token.startsWith('"') || token.startsWith("'") || token.startsWith("`")) {
    return language === "json"
      ? "text-emerald-700 dark:text-emerald-300"
      : "text-green-700 dark:text-green-300"
  }

  if (keywords.includes(token)) {
    return "text-sky-700 dark:text-sky-300"
  }

  if (language === "json" && /^(true|false|null)$/.test(token)) {
    return "text-purple-700 dark:text-purple-300"
  }

  if (/^\d/.test(token)) {
    return "text-amber-700 dark:text-amber-300"
  }

  if (/^[A-Z]/.test(token) && ["tsx", "jsx", "ts", "js"].includes(language)) {
    return "text-violet-700 dark:text-violet-300"
  }

  if (/^[{}()[\].,;:<>/=+\-*|&!]$/.test(token)) {
    return "text-muted-foreground"
  }

  if (language === "css" && token.startsWith("--")) {
    return "text-purple-700 dark:text-purple-300"
  }

  return "text-foreground"
}

function highlightLine(line: string, language: CodeBlockLanguage) {
  const nodes: React.ReactNode[] = []
  let cursor = 0

  for (const match of line.matchAll(tokenPattern)) {
    const token = match[0]
    const index = match.index ?? 0

    if (index > cursor) {
      nodes.push(line.slice(cursor, index))
    }

    nodes.push(
      <span key={`${index}-${token}`} className={getTokenClass(token, language)}>
        {token}
      </span>
    )

    cursor = index + token.length
  }

  if (cursor < line.length) {
    nodes.push(line.slice(cursor))
  }

  return nodes
}

function CodeLine({
  line,
  lineNumber,
  language,
  showLineNumbers,
}: {
  line: string
  lineNumber: number
  language: CodeBlockLanguage
  showLineNumbers: boolean
}) {
  return (
    <div className="table-row">
      {showLineNumbers ? (
        <span className="table-cell w-10 pr-5 text-right text-muted-foreground select-none">
          {lineNumber}
        </span>
      ) : null}
      <span className="table-cell min-w-0 whitespace-pre">
        {highlightLine(line || " ", language)}
      </span>
    </div>
  )
}

export function CodeBlock({
  code,
  language = "tsx",
  snippets,
  defaultLanguage,
  label = "Code",
  showLineNumbers = true,
  allowLanguageSelection = true,
  className,
  contentClassName,
}: CodeBlockProps) {
  const availableSnippets = React.useMemo<CodeBlockSnippet[]>(
    () =>
      snippets?.length
        ? snippets
        : [
            {
              language,
              code: code ?? "",
            },
          ],
    [code, language, snippets]
  )
  const initialLanguage =
    defaultLanguage ?? availableSnippets[0]?.language ?? language
  const [activeLanguage, setActiveLanguage] =
    React.useState<CodeBlockLanguage>(initialLanguage)
  const [copied, setCopied] = React.useState(false)
  const activeSnippet =
    availableSnippets.find((snippet) => snippet.language === activeLanguage) ??
    availableSnippets[0]
  const activeCode = activeSnippet?.code ?? ""
  const activeLabel =
    activeSnippet?.label ?? languageLabels[activeSnippet?.language ?? language]
  const selectable = allowLanguageSelection && availableSnippets.length > 1

  async function copyCode() {
    await navigator.clipboard.writeText(activeCode)
    setCopied(true)
    window.setTimeout(() => setCopied(false), 1400)
  }

  return (
    <div
      className={cn(
        "overflow-hidden rounded-lg border bg-card text-card-foreground",
        className
      )}
    >
      <div className="flex min-w-0 flex-wrap items-center justify-between gap-2 border-b bg-muted/30 px-3 py-2">
        <div className="flex min-w-0 items-center gap-2 text-xs font-medium text-muted-foreground">
          <TerminalIcon className="size-3.5 shrink-0" />
          <span className="truncate">{label}</span>
        </div>
        <div className="flex items-center gap-2">
          {selectable ? (
            <Select
              value={activeLanguage}
              onValueChange={(value) =>
                setActiveLanguage(value as CodeBlockLanguage)
              }
            >
              <SelectTrigger size="sm" className="h-7">
                <SelectValue />
              </SelectTrigger>
              <SelectContent align="end">
                {availableSnippets.map((snippet) => (
                  <SelectItem
                    key={snippet.language}
                    value={snippet.language}
                  >
                    {snippet.label ?? languageLabels[snippet.language]}
                  </SelectItem>
                ))}
              </SelectContent>
            </Select>
          ) : (
            <span className="rounded-md border bg-background px-2 py-1 text-[11px] font-medium text-muted-foreground">
              {activeLabel}
            </span>
          )}
          <Button
            type="button"
            variant="ghost"
            size="icon-sm"
            onClick={copyCode}
            className="rounded-md"
          >
            <span
              className="t-icon-swap inline-grid size-3.5 place-items-center"
              data-state={copied ? "b" : "a"}
              aria-hidden="true"
            >
              <span className="t-icon inline-flex size-3.5 items-center justify-center" data-icon="a">
                <CopyIcon className="size-3.5" />
              </span>
              <span className="t-icon inline-flex size-3.5 items-center justify-center" data-icon="b">
                <CheckIcon className="size-3.5" />
              </span>
            </span>
            <span className="sr-only">Copy code</span>
          </Button>
        </div>
      </div>
      <pre
        className={cn(
          "max-h-[28rem] min-w-0 overflow-auto p-4 text-xs leading-6",
          contentClassName
        )}
      >
        <code className="table min-w-max border-spacing-0 font-mono">
          {activeCode.split("\n").map((line, index) => (
            <CodeLine
              key={`${index}-${line}`}
              line={line}
              lineNumber={index + 1}
              language={activeSnippet?.language ?? language}
              showLineNumbers={showLineNumbers}
            />
          ))}
        </code>
      </pre>
    </div>
  )
}