Chat Code Block

"use client"

import * as React from "react"
import { CheckIcon, CopyIcon } from "@phosphor-icons/react"
import { cn } from "cn"

import { useCopyToClipboard } from "@/hooks/use-copy-to-clipboard"
import { Button } from "@/components/ui/button"

const LANGS = [
  "tsx",
  "typescript",
  "javascript",
  "jsx",
  "bash",
  "css",
  "json",
  "html",
  "md",
]

type Highlighter = {
  codeToHtml: (code: string, opts: Record<string, unknown>) => string
}
let highlighterPromise: Promise<Highlighter> | null = null

// Loaded on the first code block only — Shiki is far too heavy to ship eagerly.
// The JS RegExp engine avoids Oniguruma's WASM fetch, which does not load in
// the client bundle.
function getHighlighter() {
  highlighterPromise ??= (
    Promise.all([import("shiki"), import("shiki/engine/javascript")]).then(
      ([shiki, engine]) =>
        shiki.createHighlighter({
          themes: ["min-dark"],
          langs: LANGS,
          engine: engine.createJavaScriptRegexEngine(),
        })
    ) as Promise<Highlighter>
  ).catch((err) => {
    // A failed chunk load must not disable highlighting until reload.
    highlighterPromise = null
    throw err
  })
  return highlighterPromise
}

function normalizeLang(lang: string) {
  const l = lang.toLowerCase()
  if (l === "ts") return "typescript"
  if (l === "js") return "javascript"
  if (l === "sh" || l === "shell" || l === "zsh") return "bash"
  return l
}

function CopyButton({ code }: { code: string }) {
  const { copied, failed, copyToClipboard } = useCopyToClipboard({
    timeout: 1600,
  })

  return (
    <Button
      variant="ghost"
      size="icon-sm"
      aria-label="Copy code"
      onClick={() => copyToClipboard(code)}
      className="text-current/60 hover:bg-current/10 hover:text-current/90"
    >
      {copied ? <CheckIcon weight="light" /> : <CopyIcon weight="light" />}
      <span role="status" className="sr-only">
        {copied
          ? "Copied"
          : failed
            ? "Copy failed. Select and copy the code manually."
            : ""}
      </span>
    </Button>
  )
}

/**
 * A code panel inside a chat reply: the language and a copy button on top,
 * the highlighted code below, scrolling sideways when lines run long. It
 * stays dark in both light and dark mode.
 *
 * Highlighting waits for the closing fence: an open fence has no stable
 * grammar, and re-highlighting every streamed token is both slow and wrong.
 */
function ChatCodeBlock({
  code,
  lang = "tsx",
  streaming = false,
  className,
}: {
  /** The code to show. */
  code: string
  /** The language for highlighting and the label, such as "tsx" or "bash". */
  lang?: string
  /** True while the block is still arriving: plain text until it is done. */
  streaming?: boolean
  /** Extra classes for the outer `figure`, such as a different margin. */
  className?: string
}) {
  const [highlighted, setHighlighted] = React.useState<{
    code: string
    lang: string
    html: string
  } | null>(null)
  const html =
    !streaming && highlighted?.code === code && highlighted.lang === lang
      ? highlighted.html
      : null

  React.useEffect(() => {
    if (streaming) return
    let alive = true
    getHighlighter()
      .then((h) => {
        if (!alive) return
        setHighlighted({
          code,
          lang,
          html: h.codeToHtml(code, {
            lang: LANGS.includes(normalizeLang(lang))
              ? normalizeLang(lang)
              : "text",
            theme: "min-dark",
            // Lift the theme's comments to AA contrast on the scene ground.
            colorReplacements: { "#6b737c": "#8b949e" },
          }),
        })
      })
      .catch((err) => {
        console.error("[code-block] highlight failed", err)
      })
    return () => {
      alive = false
    }
  }, [code, lang, streaming])

  const label = `${normalizeLang(lang)} code`
  const scrollable =
    "scroll-fade-x overflow-x-auto px-4 pb-4 font-mono text-xs/6 focus-visible:outline-2 focus-visible:-outline-offset-2 focus-visible:outline-ring"

  return (
    <figure
      data-slot="chat-code-block"
      className={cn(
        "group/code my-4 overflow-hidden rounded-2xl bg-scene text-scene-foreground ring-1 ring-scene-foreground/10",
        className
      )}
    >
      <figcaption className="flex items-center justify-between gap-2 py-1.5 ps-4 pe-1.5">
        <span className="font-mono text-xs/5 tracking-widest text-current/60 uppercase select-none">
          {normalizeLang(lang)}
        </span>
        <CopyButton code={code} />
      </figcaption>
      {html ? (
        <div
          role="region"
          aria-label={label}
          tabIndex={0}
          className={cn(
            scrollable,
            "[&_pre]:m-0! [&_pre]:bg-transparent! [&_pre]:p-0!"
          )}
          dangerouslySetInnerHTML={{ __html: html }}
        />
      ) : (
        <pre
          role="region"
          aria-label={label}
          tabIndex={0}
          className={cn(scrollable, "text-current/80")}
        >
          <code>{code}</code>
        </pre>
      )}
    </figure>
  )
}

export { ChatCodeBlock }

pnpm dlx shadcn add @01orbital/chat-code-block

npx shadcn add @01orbital/chat-code-block

yarn shadcn add @01orbital/chat-code-block

bunx --bun shadcn add @01orbital/chat-code-block