nexus/ui/src/components/ChatCodeBlock.tsx

import { useState, type ReactNode, type HTMLAttributes } from "react";
import { Copy, Check } from "lucide-react";
import { Button } from "./ui/button";
import type { ExtraProps } from "react-markdown";

type ChatCodeBlockProps = HTMLAttributes<HTMLPreElement> & ExtraProps;

/**
 * Recursively flatten React children into a plain text string.
 */
function flattenText(value: ReactNode): string {
  if (value == null) return "";
  if (typeof value === "string" || typeof value === "number") return String(value);
  if (Array.isArray(value)) return value.map((item) => flattenText(item)).join("");
  if (typeof value === "object" && "props" in (value as object)) {
    const el = value as { props?: { children?: ReactNode } };
    return flattenText(el.props?.children);
  }
  return "";
}

/**
 * Extract the language name from a className like "language-typescript" -> "typescript".
 */
function extractLanguage(className?: string | null): string | null {
  if (!className) return null;
  const match = /\blanguage-(\w+)\b/.exec(className);
  return match ? match[1] : null;
}

/**
 * Get the className of the first child code element, if any.
 */
function getCodeClassName(children: ReactNode): string | null {
  if (children == null) return null;
  if (typeof children === "object" && "props" in (children as object)) {
    const el = children as { type?: unknown; props?: { className?: string } };
    if (el.type === "code" || String(el.type) === "code") {
      return el.props?.className ?? null;
    }
  }
  if (Array.isArray(children)) {
    for (const child of children) {
      const result = getCodeClassName(child);
      if (result !== null) return result;
    }
  }
  return null;
}

/**
 * Check whether children contain a code element (i.e., this is a fenced code block).
 */
function hasCodeChild(children: ReactNode): boolean {
  if (children == null) return false;
  if (typeof children === "object" && "props" in (children as object)) {
    const el = children as { type?: unknown };
    if (el.type === "code" || String(el.type) === "code") {
      return true;
    }
  }
  if (Array.isArray(children)) {
    return children.some((child) => hasCodeChild(child));
  }
  return false;
}

function doCopy(text: string) {
  navigator.clipboard.writeText(text).catch(() => {
    // Ignore clipboard errors in restricted environments
  });
}

export function ChatCodeBlock({ children, className, ...props }: ChatCodeBlockProps) {
  const [copied, setCopied] = useState(false);

  // If there's no code child, render a plain <pre> without the toolbar
  if (!hasCodeChild(children)) {
    return (
      <pre className={className} {...props}>
        {children}
      </pre>
    );
  }

  const codeClassName = getCodeClassName(children);
  const language = extractLanguage(codeClassName);
  const codeText = flattenText(children);

  function handleCopy() {
    doCopy(codeText);
    setCopied(true);
    setTimeout(() => setCopied(false), 1500);
  }

  return (
    <pre className={className} {...props}>
      <div className="flex items-center justify-between bg-card border-b border-border px-3 py-1">
        {language ? (
          <span className="text-xs text-muted-foreground font-mono">{language}</span>
        ) : (
          <span />
        )}
        <Button
          variant="ghost"
          size="icon"
          className="h-6 w-6"
          aria-label={copied ? "Copied!" : "Copy code"}
          onClick={handleCopy}
        >
          {copied ? (
            <Check className="h-3.5 w-3.5" />
          ) : (
            <Copy className="h-3.5 w-3.5" />
          )}
        </Button>
      </div>
      {children}
    </pre>
  );
}