spb/polyllm

1"use client";2import * as React from "react";3import { Check, ChevronsDownUp, ChevronsUpDown, Copy, Download, WrapText } from "lucide-react";4import { extensionForLanguage } from "@/lib/chat/markdown-blocks";5import { useDebounced } from "@/lib/client/hooks";6import { cn } from "@/lib/utils";78type Highlighter = { codeToHtml: (code: string, opts: { lang: string; themes: { light: string; dark: string }; defaultColor: false | "light" }) => string; getLoadedLanguages: () => string[]; loadLanguage: (lang: never) => Promise<void> };910let highlighterPromise: Promise<Highlighter> | null = null;11const KNOWN = new Set<string>();1213async function getHighlighter(): Promise<Highlighter> {14  if (!highlighterPromise) {15    highlighterPromise = (async () => {16      const { createHighlighter } = await import("shiki/bundle/web");17      const { createJavaScriptRegexEngine } = await import("shiki/engine/javascript");18      const hl = await createHighlighter({ themes: ["github-light-default", "github-dark-default"], langs: ["javascript", "typescript", "tsx", "jsx", "json", "bash", "shell", "python", "markdown", "html", "css", "sql", "yaml"], engine: createJavaScriptRegexEngine({ forgiving: true }) });19      for (const l of hl.getLoadedLanguages()) KNOWN.add(l);20      return hl as unknown as Highlighter;21    })();22  }23  return highlighterPromise;24}2526const ALIASES: Record<string, string> = { js: "javascript", ts: "typescript", sh: "bash", zsh: "bash", shell: "bash", py: "python", yml: "yaml", md: "markdown", rb: "ruby", rs: "rust", golang: "go", "c++": "cpp", cs: "csharp", kt: "kotlin", plaintext: "text", txt: "text" };2728/** Lines above which a block starts collapsed (Expand reveals everything). */29const COLLAPSE_LINES = 28;3031export function CodeBlock({ code, lang, wrap }: { code: string; lang?: string; wrap?: boolean }) {32  const [html, setHtml] = React.useState<string | null>(null);33  const [copied, setCopied] = React.useState(false);34  const [localWrap, setLocalWrap] = React.useState<boolean | null>(null);35  const [expanded, setExpanded] = React.useState(false);36  const language = React.useMemo(() => {37    const l = (lang ?? "").toLowerCase().trim();38    return ALIASES[l] ?? l;39  }, [lang]);40  const effectiveWrap = localWrap ?? wrap ?? false;41  // While streaming the block changes on every token; highlight the settled version only.42  const settled = useDebounced(code, 120);43  const lineCount = React.useMemo(() => (code ? code.split("\n").length : 0), [code]);44  const collapsible = lineCount > COLLAPSE_LINES;4546  React.useEffect(() => {47    let cancelled = false;48    if (!language || language === "text") {49      // eslint-disable-next-line react-hooks/set-state-in-effect50      setHtml(null);51      return;52    }53    (async () => {54      try {55        const hl = await getHighlighter();56        if (!KNOWN.has(language)) {57          try {58            await hl.loadLanguage(language as never);59            KNOWN.add(language);60          } catch {61            if (!cancelled) setHtml(null);62            return;63          }64        }65        const out = hl.codeToHtml(settled, { lang: language, themes: { light: "github-light-default", dark: "github-dark-default" }, defaultColor: "light" });66        if (!cancelled) setHtml(out);67      } catch {68        if (!cancelled) setHtml(null);69      }70    })();71    return () => {72      cancelled = true;73    };74  }, [settled, language]);7576  const copy = async () => {77    try {78      await navigator.clipboard.writeText(code);79      setCopied(true);80      setTimeout(() => setCopied(false), 1400);81    } catch {82      /* ignore */83    }84  };8586  const download = () => {87    const ext = extensionForLanguage(language);88    const blob = new Blob([code], { type: "text/plain;charset=utf-8" });89    const url = URL.createObjectURL(blob);90    const a = document.createElement("a");91    a.href = url;92    a.download = ext === "Dockerfile" ? "Dockerfile" : `snippet.${ext}`;93    document.body.appendChild(a);94    a.click();95    a.remove();96    setTimeout(() => URL.revokeObjectURL(url), 1000);97  };9899  // Show the highlighted HTML only when it matches the current code (avoids stale colouring while streaming).100  const showHtml = html && settled === code;101102  return (103    <div className="group/code my-2 overflow-hidden rounded-lg border border-border bg-bg-subtle text-[13px]">104      <div className="flex h-9 items-center justify-between gap-2 border-b border-border px-2.5 sm:h-8 sm:px-3">105        <span className="truncate font-mono text-[11px] text-fg-subtle">106          {language || "text"}107          {lineCount > 1 ? <span className="hidden sm:inline"> · {lineCount} lines</span> : null}108        </span>109        <div className="flex items-center gap-0.5">110          {collapsible ? (111            <ToolBtn label={expanded ? "Collapse" : "Expand"} onClick={() => setExpanded((e) => !e)}>112              {expanded ? <ChevronsDownUp className="size-3.5" /> : <ChevronsUpDown className="size-3.5" />}113            </ToolBtn>114          ) : null}115          <ToolBtn label={effectiveWrap ? "Disable line wrap" : "Wrap long lines"} onClick={() => setLocalWrap((w) => !(w ?? wrap ?? false))} active={effectiveWrap}>116            <WrapText className="size-3.5" />117          </ToolBtn>118          <ToolBtn label="Download" onClick={download}>119            <Download className="size-3.5" />120          </ToolBtn>121          <button onClick={copy} className="tap inline-flex h-7 items-center gap-1 rounded px-1.5 text-[11px] text-fg-subtle hover:bg-bg-muted hover:text-fg" aria-label="Copy code">122            {copied ? <Check className="size-3.5 text-success" /> : <Copy className="size-3.5" />}123            <span className="hidden sm:inline">{copied ? "Copied" : "Copy"}</span>124          </button>125        </div>126      </div>127      <div className={cn("relative", collapsible && !expanded && "max-h-[420px] overflow-hidden")}>128        <div className={cn("overflow-x-auto scrollbar-thin", effectiveWrap && "[&_pre]:whitespace-pre-wrap [&_pre]:break-words [&_code]:whitespace-pre-wrap")}>129          {showHtml ? (130            <div className="[&_pre]:m-0 [&_pre]:bg-transparent! [&_pre]:p-3 [&_pre]:leading-[1.55] [&_code]:font-mono [&_code]:text-[12.5px]" dangerouslySetInnerHTML={{ __html: html }} />131          ) : (132            <pre className="m-0 p-3 leading-[1.55]">133              <code className="font-mono text-[12.5px]">{code}</code>134            </pre>135          )}136        </div>137        {collapsible && !expanded ? (138          <button type="button" onClick={() => setExpanded(true)} className="absolute inset-x-0 bottom-0 flex h-16 items-end justify-center bg-gradient-to-t from-bg-subtle via-bg-subtle/80 to-transparent pb-2 text-[12px] font-medium text-fg-muted hover:text-fg">139            Show all {lineCount} lines140          </button>141        ) : null}142      </div>143    </div>144  );145}146147function ToolBtn({ label, onClick, active, children }: { label: string; onClick: () => void; active?: boolean; children: React.ReactNode }) {148  return (149    <button type="button" onClick={onClick} className={cn("tap rounded p-1.5 text-fg-subtle hover:bg-bg-muted hover:text-fg", active && "text-accent")} aria-label={label} title={label}>150      {children}151    </button>152  );153}154