Code block

"use client";

/*
 * Highlighted code that streams without the layout jumping: a header with
 * the file name, language, status and copy action; line numbers; focused
 * lines; rows that keep their colour while the next chunk is being
 * tokenised. Built on Intelligo's tokens and Base UI. Highlighting is
 * shiki, light and dark themes at once; the plain text shows until the
 * tokens are ready.
 *
 * Colour: shiki gives every token a light and a dark colour, so the
 * container declares `color-scheme` per theme (the one structural
 * `dark:` in this file) and each token paints with `light-dark()`.
 */

import * as React from "react";
import { CheckIcon, FileCode2Icon } from "lucide-react";
import {
  type BundledLanguage,
  codeToHtml,
  codeToTokensWithThemes,
  type ShikiTransformer,
  type SpecialLanguage,
} from "shiki";

import { CopyButton } from "@showcase/components/ui/copy-button";
import { Spinner } from "@showcase/components/ui/spinner";
import { cn } from "@showcase/lib/utils";

export type CodeLanguage = BundledLanguage | SpecialLanguage;
export type CodeBlockStatus = "streaming" | "complete";

const LIGHT_THEME = "one-light";
const DARK_THEME = "one-dark-pro";

/* ----------------------------------------------------------------------------
 * highlightCode: the HTML form, kept for callers that want a whole block
 * rendered by shiki (light and dark) rather than token rows.
 * ------------------------------------------------------------------------- */

const lineNumbers: ShikiTransformer = {
  name: "line-numbers",
  line(node, line) {
    node.children.unshift({
      type: "element",
      tagName: "span",
      properties: {
        className: [
          "inline-block",
          "min-w-10",
          "mr-4",
          "text-right",
          "select-none",
          "text-muted-foreground",
        ],
      },
      children: [{ type: "text", value: String(line) }],
    });
  },
};

async function highlightCode(
  code: string,
  language: CodeLanguage,
  showLineNumbers = false
): Promise<[light: string, dark: string]> {
  const transformers = showLineNumbers ? [lineNumbers] : [];
  return Promise.all([
    codeToHtml(code, { lang: language, theme: LIGHT_THEME, transformers }),
    codeToHtml(code, { lang: language, theme: DARK_THEME, transformers }),
  ]);
}

/* ----------------------------------------------------------------------------
 * Tokens: one shiki pass gives both theme colours per token. The hook
 * keeps the last result while a longer code (a stream) is tokenised, so
 * lines already on screen never flash back to plain text.
 * ------------------------------------------------------------------------- */

export interface CodeToken {
  content: string;
  offset: number;
  light?: string;
  dark?: string;
}

export type CodeTokenLines = CodeToken[][];

const TOKEN_CACHE_LIMIT = 64;
const tokenCache = new Map<string, CodeTokenLines>();

function tokenCacheKey(code: string, language: CodeLanguage) {
  return `${language}${code}`;
}

function rememberTokens(key: string, lines: CodeTokenLines) {
  if (tokenCache.size >= TOKEN_CACHE_LIMIT) {
    const oldest = tokenCache.keys().next().value;
    if (oldest !== undefined) tokenCache.delete(oldest);
  }
  tokenCache.set(key, lines);
}

async function tokenizeCode(
  code: string,
  language: CodeLanguage
): Promise<CodeTokenLines> {
  const lines = await codeToTokensWithThemes(code, {
    lang: language,
    themes: { light: LIGHT_THEME, dark: DARK_THEME },
  });
  return lines.map((line) =>
    line.map((token) => ({
      content: token.content,
      offset: token.offset,
      light: token.variants.light?.color,
      dark: token.variants.dark?.color,
    }))
  );
}

type TokenResult = {
  key: string;
  code: string;
  language: CodeLanguage;
  lines: CodeTokenLines;
};

/**
 * Shiki tokens for `code`, or `null` until they exist. While a stream
 * grows the previous tokens stay valid for the lines they cover.
 */
function useCodeTokens(
  code: string,
  language: CodeLanguage
): CodeTokenLines | null {
  const key = tokenCacheKey(code, language);
  const cached = tokenCache.get(key);
  const [result, setResult] = React.useState<TokenResult | null>(
    cached ? { key, code, language, lines: cached } : null
  );

  React.useEffect(() => {
    const current = tokenCache.get(key);
    if (current) {
      setResult({ key, code, language, lines: current });
      return;
    }

    let cancelled = false;
    tokenizeCode(code, language)
      .then((lines) => {
        if (cancelled) return;
        rememberTokens(key, lines);
        setResult({ key, code, language, lines });
      })
      .catch(() => {
        // An unknown language or a failed load: the plain text stays.
      });
    return () => {
      cancelled = true;
    };
  }, [code, key, language]);

  if (result?.key === key) return result.lines;
  if (result?.language === language && code.startsWith(result.code)) {
    return result.lines;
  }
  return null;
}

function splitLines(code: string) {
  let offset = 0;
  return code.split("\n").map((content) => {
    const line = { content, offset };
    offset += content.length + 1;
    return line;
  });
}

/* ----------------------------------------------------------------------------
 * CodeLine and HighlightedCode: the token rows, without chrome.
 * ------------------------------------------------------------------------- */

export interface CodeLineProps extends Omit<React.ComponentProps<"span">, "children"> {
  code: string;
  tokens?: CodeToken[];
}

function CodeLine({ code, tokens, className, ...props }: CodeLineProps) {
  return (
    <span data-slot="code-line" className={className} {...props}>
      {tokens
        ? tokens.map((token) => (
            <span
              key={`${token.offset}-${token.content}`}
              style={{
                color: `light-dark(${token.light ?? "currentColor"}, ${
                  token.dark ?? token.light ?? "currentColor"
                })`,
              }}
            >
              {token.content}
            </span>
          ))
        : code}
    </span>
  );
}

/** The `color-scheme` switch every token colour reads through. */
const SCHEME = "[color-scheme:light] dark:[color-scheme:dark]";

export interface HighlightedCodeProps
  extends Omit<React.ComponentProps<"pre">, "children"> {
  code: string;
  language?: CodeLanguage;
}

/** Highlighted code as a bare `<pre>`, for output panes and inline results. */
function HighlightedCode({
  code,
  language = "bash",
  className,
  ...props
}: HighlightedCodeProps) {
  const tokens = useCodeTokens(code, language);
  const lines = splitLines(code);

  return (
    <pre
      data-slot="highlighted-code"
      className={cn(
        "m-0 overflow-x-auto font-mono text-xs leading-5 whitespace-pre text-foreground/85",
        SCHEME,
        className
      )}
      {...props}
    >
      <code>
        {lines.map((line, index) => (
          <React.Fragment key={line.offset}>
            <CodeLine code={line.content} tokens={tokens?.[index]} />
            {index < lines.length - 1 ? "\n" : null}
          </React.Fragment>
        ))}
      </code>
    </pre>
  );
}

/* ----------------------------------------------------------------------------
 * CodeBlock: the framed block with header, line numbers and actions.
 * ------------------------------------------------------------------------- */

type CodeBlockContextValue = { code: string };

const CodeBlockContext = React.createContext<CodeBlockContextValue>({
  code: "",
});

export interface CodeBlockProps
  extends Omit<React.ComponentProps<"div">, "children"> {
  code: string;
  language?: CodeLanguage;
  /** Shown in the header; the header appears when there is one. */
  filename?: React.ReactNode;
  status?: CodeBlockStatus;
  /** Force the header on (or off) regardless of filename and status. */
  showHeader?: boolean;
  showLineNumbers?: boolean;
  /** 1-based line numbers to emphasise. */
  highlightLines?: number[];
  /** Pixel height the viewport scrolls within; unset grows with the code. */
  maxHeight?: number;
  wrap?: boolean;
  /** Render the built-in copy action (in the header, or top-right). */
  copyable?: boolean;
  onCopy?: () => void;
  copyLabel?: string;
  copiedLabel?: string;
  statusLabels?: Partial<Record<CodeBlockStatus, string>>;
  /** Extra actions — `<CodeBlockCopyButton />` or the caller's own. */
  children?: React.ReactNode;
}

const STATUS_LABELS: Record<CodeBlockStatus, string> = {
  streaming: "Writing",
  complete: "Ready",
};

function CodeBlock({
  code,
  language = "typescript",
  filename,
  status = "complete",
  showHeader,
  showLineNumbers = false,
  highlightLines,
  maxHeight,
  wrap = false,
  copyable = false,
  onCopy,
  copyLabel = "Copy code",
  copiedLabel = "Copied",
  statusLabels,
  className,
  children,
  ...props
}: CodeBlockProps) {
  const viewportRef = React.useRef<HTMLDivElement>(null);
  const streaming = status === "streaming";
  const header = showHeader ?? (Boolean(filename) || streaming);
  const tokens = useCodeTokens(code, language);
  const highlighted = React.useMemo(
    () => new Set(highlightLines ?? []),
    [highlightLines]
  );
  const lines = splitLines(code);
  const labels = { ...STATUS_LABELS, ...statusLabels };

  // A stream keeps the newest line in view; the caller's scroll position
  // is left alone once the code is complete.
  React.useLayoutEffect(() => {
    const viewport = viewportRef.current;
    if (!viewport || !streaming) return;

    const frame = requestAnimationFrame(() => {
      if (viewport.scrollHeight <= viewport.clientHeight) return;
      const reduce = window.matchMedia?.(
        "(prefers-reduced-motion: reduce)"
      ).matches;
      if (typeof viewport.scrollTo === "function") {
        viewport.scrollTo({
          top: viewport.scrollHeight,
          behavior: reduce ? "auto" : "smooth",
        });
      } else {
        viewport.scrollTop = viewport.scrollHeight;
      }
    });
    return () => cancelAnimationFrame(frame);
  });

  const actions =
    copyable || children ? (
      <div
        data-slot="code-block-actions"
        className={cn(
          "flex shrink-0 items-center gap-1",
          !header && "absolute top-2 right-2"
        )}
      >
        {copyable ? (
          <CopyButton
            data-slot="code-block-copy"
            value={code}
            label={copyLabel}
            copiedLabel={copiedLabel}
            onCopy={onCopy}
          />
        ) : null}
        {children}
      </div>
    ) : null;

  return (
    <CodeBlockContext.Provider value={{ code }}>
      <div
        data-slot="code-block"
        data-state={status}
        aria-busy={streaming || undefined}
        className={cn(
          "group/code-block relative w-full overflow-hidden rounded-lg border bg-background text-sm text-foreground",
          className
        )}
        {...props}
      >
        {header ? (
          <div
            data-slot="code-block-header"
            className="flex h-10 items-center gap-2.5 border-b px-3"
          >
            <FileCode2Icon
              aria-hidden="true"
              className="size-3.5 shrink-0 text-muted-foreground/70"
            />
            {filename ? (
              <span className="min-w-0 truncate font-mono text-xs text-foreground/80">
                {filename}
              </span>
            ) : null}
            <span className="text-xs font-medium tracking-wide text-muted-foreground/55 uppercase">
              {language}
            </span>
            <span
              data-slot="code-block-status"
              className={cn(
                "ml-auto inline-flex shrink-0 items-center gap-1 text-xs font-medium",
                streaming ? "text-info" : "text-success"
              )}
            >
              {streaming ? (
                <Spinner className="size-3" aria-label={labels.streaming} />
              ) : (
                <CheckIcon className="size-3" aria-hidden="true" />
              )}
              {streaming ? labels.streaming : labels.complete}
            </span>
            {actions}
          </div>
        ) : (
          actions
        )}

        <div
          ref={viewportRef}
          data-slot="code-block-viewport"
          role={streaming ? "log" : undefined}
          aria-live={streaming ? "polite" : undefined}
          className={cn(
            "overflow-auto py-2 [scrollbar-width:none] [&::-webkit-scrollbar]:hidden",
            SCHEME
          )}
          style={{ maxHeight }}
        >
          <pre
            className={cn(
              "m-0 font-mono text-xs leading-5 text-foreground/85",
              !wrap && "min-w-max"
            )}
          >
            <code>
              {lines.map((line, index) => {
                const lineNumber = index + 1;
                return (
                  <span
                    key={line.offset}
                    data-slot="code-block-line"
                    data-highlighted={
                      highlighted.has(lineNumber) || undefined
                    }
                    className={cn(
                      "grid min-h-5",
                      highlighted.has(lineNumber) && "bg-info/10"
                    )}
                    style={{
                      gridTemplateColumns: showLineNumbers
                        ? "2.75rem minmax(0, 1fr)"
                        : "minmax(0, 1fr)",
                    }}
                  >
                    {showLineNumbers ? (
                      <span className="pr-3 text-right text-muted-foreground/40 tabular-nums select-none">
                        {lineNumber}
                      </span>
                    ) : null}
                    <CodeLine
                      code={line.content}
                      tokens={tokens?.[index]}
                      className={cn(
                        "pr-4",
                        showLineNumbers ? "pl-1" : "pl-4",
                        wrap ? "break-words whitespace-pre-wrap" : "whitespace-pre"
                      )}
                    />
                  </span>
                );
              })}
            </code>
          </pre>
        </div>
      </div>
    </CodeBlockContext.Provider>
  );
}

function CodeBlockCopyButton(
  props: Omit<React.ComponentProps<typeof CopyButton>, "value">
) {
  const { code } = React.useContext(CodeBlockContext);
  return <CopyButton data-slot="code-block-copy" value={code} {...props} />;
}

export {
  CodeBlock,
  CodeBlockCopyButton,
  CodeLine,
  HighlightedCode,
  highlightCode,
  useCodeTokens,
};