Code Block

"use client";

import * as React from "react";
import { FileCode2, WrapText } from "lucide-react";
import { cn } from "@/lib/utils";
import { CopyButton } from "@/registry/paragon/ui/copy-button";
import {
  Tooltip,
  TooltipContent,
  TooltipProvider,
  TooltipTrigger,
} from "@/registry/paragon/ui/tooltip";

/*
 * Deliberately NOT a syntax highlighter. One tasteful pass: comments recede,
 * string literals take a single desaturated tint, everything else stays
 * foreground. Line numbers are pseudo-element content (attr(data-ln)) so a
 * drag-select never captures them; diff markers are select-none for the
 * same reason.
 */

type TokenType = "code" | "comment" | "string";

interface Token {
  type: TokenType;
  text: string;
}

interface ParsedLine {
  number: number;
  /** Line content without the diff marker — what per-line copy writes. */
  text: string;
  marker: "+" | "-" | null;
  tokens: Token[];
}

const HASH_COMMENT_LANGS = new Set([
  "bash", "sh", "shell", "zsh", "python", "py", "ruby", "rb", "yaml", "yml",
  "toml", "dockerfile", "makefile", "ini", "conf", "env",
]);
const DASH_COMMENT_LANGS = new Set(["sql", "lua", "haskell"]);

const STRING_RE = /(["'`])(?:\\.|(?!\1).)*?\1/g;

/** Splits string literals out of a code fragment. */
function tokenizeStrings(text: string): Token[] {
  const out: Token[] = [];
  let last = 0;
  STRING_RE.lastIndex = 0;
  for (const match of text.matchAll(STRING_RE)) {
    const index = match.index ?? 0;
    if (index > last) out.push({ type: "code", text: text.slice(last, index) });
    out.push({ type: "string", text: match[0] });
    last = index + match[0].length;
  }
  if (last < text.length) out.push({ type: "code", text: text.slice(last) });
  return out;
}

/** Index of a trailing line comment, or -1. Requires whitespace before the
 * delimiter so `https://…` never reads as a comment. */
function trailingCommentIndex(text: string, delimiter: string): number {
  for (let i = 1; i < text.length; i++) {
    if (
      text.startsWith(delimiter, i) &&
      /\s/.test(text[i - 1]) &&
      (i + delimiter.length >= text.length || /\s/.test(text[i + delimiter.length]))
    ) {
      return i;
    }
  }
  return -1;
}

function parseLines(
  code: string,
  language: string | undefined,
  diff: boolean,
): ParsedLine[] {
  const lang = language?.toLowerCase() ?? "";
  const hash = HASH_COMMENT_LANGS.has(lang);
  const dash = DASH_COMMENT_LANGS.has(lang);
  const slash = !hash && !dash;

  let inBlockComment = false;

  return code.replace(/\n$/, "").split("\n").map((raw, i) => {
    let marker: ParsedLine["marker"] = null;
    let text = raw;
    if (diff && (raw.startsWith("+") || raw.startsWith("-"))) {
      marker = raw[0] as "+" | "-";
      text = raw.slice(1).replace(/^ /, "");
    }

    const tokens: Token[] = [];
    let rest = text;

    if (inBlockComment) {
      const end = rest.indexOf("*/");
      if (end === -1) {
        tokens.push({ type: "comment", text: rest });
        rest = "";
      } else {
        tokens.push({ type: "comment", text: rest.slice(0, end + 2) });
        rest = rest.slice(end + 2);
        inBlockComment = false;
      }
    }

    if (rest) {
      const trimmed = rest.trimStart();
      const isFullLine =
        (slash && (trimmed.startsWith("//") || trimmed.startsWith("*"))) ||
        (hash && trimmed.startsWith("#")) ||
        (dash && trimmed.startsWith("--"));

      if (isFullLine) {
        tokens.push({ type: "comment", text: rest });
      } else {
        // Block comment opening on this line (heuristic: ignores strings).
        const blockStart = slash ? rest.indexOf("/*") : -1;
        let codePart = rest;
        let commentPart = "";
        if (blockStart !== -1) {
          const end = rest.indexOf("*/", blockStart + 2);
          if (end === -1) {
            codePart = rest.slice(0, blockStart);
            commentPart = rest.slice(blockStart);
            inBlockComment = true;
          } else {
            codePart = rest.slice(0, blockStart);
            commentPart = rest.slice(blockStart, end + 2);
            const tail = rest.slice(end + 2);
            tokens.push(...tokenizeStrings(codePart));
            tokens.push({ type: "comment", text: commentPart });
            tokens.push(...tokenizeStrings(tail));
            codePart = "";
            commentPart = "";
          }
        }
        if (codePart || commentPart) {
          const delimiter = slash ? "//" : hash ? "#" : "--";
          const at = commentPart ? -1 : trailingCommentIndex(codePart, delimiter);
          if (at !== -1) {
            tokens.push(...tokenizeStrings(codePart.slice(0, at)));
            tokens.push({ type: "comment", text: codePart.slice(at) });
          } else if (codePart) {
            tokens.push(...tokenizeStrings(codePart));
          }
          if (commentPart) tokens.push({ type: "comment", text: commentPart });
        }
      }
    }

    return { number: i + 1, text, marker, tokens };
  });
}

function normalizeHighlights(
  ranges: (number | [number, number])[] | undefined,
): Set<number> {
  const set = new Set<number>();
  for (const range of ranges ?? []) {
    if (typeof range === "number") set.add(range);
    else for (let i = range[0]; i <= range[1]; i++) set.add(i);
  }
  return set;
}

const TOKEN_CLASSES: Record<TokenType, string | undefined> = {
  code: undefined,
  comment: "text-muted-foreground/80 italic",
  string: "text-[oklch(0.52_0.1_152)] dark:text-[oklch(0.78_0.1_152)]",
};

/** Deterministic skeleton bar widths. */
const SKELETON_WIDTHS = [72, 88, 56, 78, 42, 64, 82, 36];

export interface CodeBlockProps
  extends Omit<React.ComponentProps<"div">, "children"> {
  /** Raw source. Diff markers (`+`/`-` line prefixes) are read when `diff`. */
  code: string;
  /** Used for the header tag and comment heuristics (`bash`, `sql`, …). */
  language?: string;
  /** Filename shown in the header tab. */
  filename?: string;
  /** 1-based lines or `[from, to]` ranges to accent. */
  highlightLines?: (number | [number, number])[];
  /** Tint `+`/`-` prefixed lines as an inline diff. */
  diff?: boolean;
  showLineNumbers?: boolean;
  /** Initial state of the wrap toggle. */
  defaultWrap?: boolean;
  /** Hides the filename/actions bar (per-line copy still works). */
  showHeader?: boolean;
  /** Scroll ceiling for the code area, in px. */
  maxHeight?: number;
  /** Per-line hover copy buttons. */
  lineCopy?: boolean;
  /** Skeleton lines while source loads. */
  loading?: boolean;
  /** Disables copy icon-swap motion. */
  static?: boolean;
}

/**
 * The docs workhorse: a monospace block with pseudo-element line numbers
 * (never captured by selection), highlight ranges, diff line tinting, a wrap
 * toggle, whole-block copy with the house icon swap, and a per-line copy
 * that surfaces on row hover.
 */
export function CodeBlock({
  code,
  language,
  filename,
  highlightLines,
  diff = false,
  showLineNumbers = true,
  defaultWrap = false,
  showHeader = true,
  maxHeight = 384,
  lineCopy = true,
  loading = false,
  static: isStatic = false,
  className,
  ...props
}: CodeBlockProps) {
  const [wrap, setWrap] = React.useState(defaultWrap);
  const lines = React.useMemo(
    () => parseLines(code, language, diff),
    [code, language, diff],
  );
  const highlighted = React.useMemo(
    () => normalizeHighlights(highlightLines),
    [highlightLines],
  );
  const gutterCh = String(lines.length).length;
  const empty = !loading && code.trim() === "";

  return (
    <div
      data-slot="code-block"
      className={cn(
        "w-full overflow-hidden rounded-xl bg-card text-card-foreground shadow-border",
        className,
      )}
      {...props}
    >
      {showHeader && (
        <div className="flex h-10 items-center gap-2 border-b px-3">
          <FileCode2 aria-hidden className="size-3.5 shrink-0 text-muted-foreground" />
          <span
            className="min-w-0 truncate font-mono text-xs font-medium"
            title={filename}
          >
            {filename ?? (language ? language.toUpperCase() : "Snippet")}
          </span>
          <span className="flex-1" />
          {filename && language && (
            <span className="shrink-0 text-[10px] font-medium tracking-wider text-muted-foreground uppercase">
              {language}
            </span>
          )}
          <TooltipProvider>
            <Tooltip>
              <TooltipTrigger asChild>
                <button
                  type="button"
                  aria-pressed={wrap}
                  aria-label="Wrap lines"
                  onClick={() => setWrap((w) => !w)}
                  className={cn(
                    "pressable relative flex size-7 shrink-0 items-center justify-center rounded-md transition-colors duration-150 ease-out",
                    "after:absolute after:top-1/2 after:left-1/2 after:size-10 after:-translate-1/2",
                    wrap
                      ? "bg-muted text-foreground"
                      : "text-muted-foreground hover:text-foreground",
                  )}
                >
                  <WrapText className="size-3.5" />
                </button>
              </TooltipTrigger>
              <TooltipContent>{wrap ? "Unwrap lines" : "Wrap lines"}</TooltipContent>
            </Tooltip>
            <Tooltip>
              <TooltipTrigger asChild>
                <CopyButton value={code} static={isStatic} />
              </TooltipTrigger>
              <TooltipContent>Copy code</TooltipContent>
            </Tooltip>
          </TooltipProvider>
        </div>
      )}

      {loading ? (
        <div aria-hidden className="space-y-2.5 px-4 py-4">
          {SKELETON_WIDTHS.map((width, i) => (
            <div key={i} className="flex items-center gap-4">
              {showLineNumbers && <div className="h-2.5 w-4 rounded-sm bg-muted animate-pulse motion-reduce:animate-none" />}
              <div
                className="h-2.5 rounded-sm bg-muted animate-pulse motion-reduce:animate-none"
                style={{ width: `${width}%`, animationDelay: `${i * 90}ms` }}
              />
            </div>
          ))}
        </div>
      ) : empty ? (
        <div className="flex flex-col items-center gap-1 px-4 py-10 text-center">
          <FileCode2 aria-hidden className="size-4 text-muted-foreground/60" />
          <p className="text-xs text-muted-foreground">Nothing to display</p>
        </div>
      ) : (
        <div
          tabIndex={0}
          role="region"
          aria-label={filename ? `Code: ${filename}` : "Code sample"}
          className="overflow-auto outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-inset"
          style={{ maxHeight }}
        >
          <pre
            className={cn(
              "py-3 font-mono text-[12.5px] leading-[1.7]",
              wrap ? "w-full" : "w-max min-w-full",
            )}
            style={{ "--code-gutter": `${gutterCh}ch` } as React.CSSProperties}
          >
            <code className="block">
              {lines.map((line) => (
                <CodeLine
                  key={line.number}
                  line={line}
                  wrap={wrap}
                  diff={diff}
                  lineCopy={lineCopy}
                  isStatic={isStatic}
                  showLineNumbers={showLineNumbers}
                  isHighlighted={highlighted.has(line.number)}
                />
              ))}
            </code>
          </pre>
        </div>
      )}
    </div>
  );
}

function CodeLine({
  line,
  wrap,
  diff,
  lineCopy,
  isStatic,
  showLineNumbers,
  isHighlighted,
}: {
  line: ParsedLine;
  wrap: boolean;
  diff: boolean;
  lineCopy: boolean;
  isStatic: boolean;
  showLineNumbers: boolean;
  isHighlighted: boolean;
}) {
  return (
    <span
      data-ln={line.number}
      className={cn(
        "group/ln relative grid w-full border-l-2 border-transparent pr-9 pl-4 transition-colors duration-(--duration-fast)",
        showLineNumbers
          ? "grid-cols-[var(--code-gutter)_1fr] gap-x-4 before:text-right before:text-muted-foreground/50 before:tabular-nums before:content-[attr(data-ln)]"
          : "grid-cols-[1fr]",
        isHighlighted && "border-l-primary/60 bg-primary/[0.05] dark:bg-primary/[0.08]",
        !isHighlighted && line.marker === "+" && "bg-success/[0.08] dark:bg-success/[0.12]",
        !isHighlighted && line.marker === "-" && "bg-destructive/[0.07] dark:bg-destructive/[0.12]",
        "hover:bg-muted/50",
      )}
    >
      <span
        className={cn(
          "min-w-0",
          wrap ? "break-words whitespace-pre-wrap" : "whitespace-pre",
        )}
      >
        {diff && (
          <span
            aria-hidden
            className={cn(
              "inline-block w-4 select-none",
              line.marker === "+" && "text-success",
              line.marker === "-" && "text-destructive",
              !line.marker && "text-transparent",
            )}
          >
            {line.marker ?? " "}
          </span>
        )}
        {line.tokens.length === 0 ? (
          // Keep empty lines at full height.
          "​"
        ) : (
          line.tokens.map((token, i) => (
            <span key={i} className={TOKEN_CLASSES[token.type]}>
              {token.text}
            </span>
          ))
        )}
      </span>
      {lineCopy && line.text.trim() !== "" && (
        <CopyButton
          value={line.text}
          static={isStatic}
          tabIndex={-1}
          aria-label={`Copy line ${line.number}`}
          className={cn(
            "absolute top-1/2 right-1.5 size-5 -translate-y-1/2 rounded-[5px] bg-card opacity-0 shadow-border transition-opacity duration-(--duration-fast)",
            "after:size-6 group-hover/ln:opacity-100 [&_svg]:size-3",
          )}
        />
      )}
    </span>
  );
}