Code block

"use client"

import * as React from "react"
import { cn } from "cn"

import { CopyButton } from "@/components/ui/copy-button"

type CodeBlockProps = Omit<React.ComponentProps<"figure">, "children"> & {
  code: string
  /** Any language id shiki knows, e.g. `tsx`, `bash`, `json`. */
  language?: string
  filename?: string
  showLineNumbers?: boolean
  /** 1-based line numbers to emphasize. */
  highlightLines?: number[]
  copyable?: boolean
}

async function highlight(code: string, language: string, highlightLines: number[]) {
  const { codeToHtml } = await import("shiki")
  const highlighted = new Set(highlightLines)
  return codeToHtml(code, {
    lang: language,
    themes: { light: "github-light", dark: "github-dark" },
    defaultColor: false,
    transformers: [
      {
        line(node, line) {
          if (highlighted.has(line)) this.addClassToHast(node, "highlighted")
        },
      },
    ],
  })
}

function CodeBlock({
  code,
  language = "text",
  filename,
  showLineNumbers = false,
  highlightLines = [],
  copyable = true,
  className,
  ...props
}: CodeBlockProps) {
  const trimmed = code.replace(/\n+$/, "")
  const [html, setHtml] = React.useState<string | null>(null)
  const highlightKey = highlightLines.join(",")

  React.useEffect(() => {
    let cancelled = false
    highlight(trimmed, language, highlightKey ? highlightKey.split(",").map(Number) : [])
      .then((result) => {
        if (!cancelled) setHtml(result)
      })
      .catch((error: unknown) => {
        console.error(`CodeBlock: highlighting "${language}" failed`, error)
      })
    return () => {
      cancelled = true
    }
  }, [trimmed, language, highlightKey])

  const header = filename ?? (language !== "text" ? language : null)

  return (
    <figure
      data-slot="code-block"
      data-line-numbers={showLineNumbers || undefined}
      className={cn(
        "group/code-block relative overflow-hidden rounded-lg border bg-card text-[0.8125rem]",
        className
      )}
      {...props}
    >
      {header && (
        <figcaption className="flex h-9 items-center justify-between gap-2 border-b bg-muted/40 pr-1.5 pl-3 font-mono text-xs text-muted-foreground">
          <span className="truncate">{header}</span>
          {copyable && <CopyButton value={trimmed} size="icon-xs" label="Copy code" />}
        </figcaption>
      )}
      {!header && copyable && (
        <CopyButton
          value={trimmed}
          size="icon-xs"
          label="Copy code"
          className="absolute top-2 right-2 z-10 bg-card/80 opacity-0 backdrop-blur transition-opacity group-hover/code-block:opacity-100 focus-visible:opacity-100"
        />
      )}
      <div
        className={cn(
          "overflow-x-auto py-3 font-mono leading-relaxed",
          "[&_pre]:bg-transparent! [&_pre]:outline-none [&_code]:grid [&_code]:min-w-max",
          "[&_.line]:px-4 [&_.line.highlighted]:bg-brand/10 [&_.line.highlighted]:shadow-[inset_2px_0_0_var(--brand)]",
          "[&_span]:text-(--shiki-light) dark:[&_span]:text-(--shiki-dark)",
          "group-data-line-numbers/code-block:[&_code]:[counter-reset:line]",
          "group-data-line-numbers/code-block:[&_.line]:before:mr-4 group-data-line-numbers/code-block:[&_.line]:before:inline-block group-data-line-numbers/code-block:[&_.line]:before:w-6 group-data-line-numbers/code-block:[&_.line]:before:text-right group-data-line-numbers/code-block:[&_.line]:before:text-muted-foreground/60 group-data-line-numbers/code-block:[&_.line]:before:content-[counter(line)] group-data-line-numbers/code-block:[&_.line]:before:[counter-increment:line]"
        )}
      >
        {html ? (
          <div dangerouslySetInnerHTML={{ __html: html }} />
        ) : (
          <pre>
            <code>
              {trimmed.split("\n").map((line, index) => (
                <span
                  key={index}
                  className={cn("line", highlightLines.includes(index + 1) && "highlighted")}
                >
                  {line || " "}
                </span>
              ))}
            </code>
          </pre>
        )}
      </div>
    </figure>
  )
}

export { CodeBlock, type CodeBlockProps }