strix/strix/viewer/frontend/src/components/vulnerability/MdCodeBlock.tsx

"use client";

import { useState } from "react";
import hljs from "@/lib/hljs";
import "highlight.js/styles/github-dark.css";
import { Copy, Check } from "lucide-react";
import { copyToClipboard } from "@/lib/vulnerability-utils";

export function MdCodeBlock({
  className,
  children,
  node,
}: {
  className?: string;
  children?: React.ReactNode;
  node?: { data?: { meta?: string }; properties?: { metastring?: string } };
}) {
  const [copied, setCopied] = useState(false);
  const raw = String(children).replace(/\n$/, "");
  const match = /language-(\w+)/.exec(className || "");
  const isBlock = raw.includes("\n") || match;

  if (!isBlock) {
    return <code className={`${className || ""} bg-white/8 px-1.5 py-0.5 rounded text-[13px]`}>{children}</code>;
  }

  const meta = node?.data?.meta || node?.properties?.metastring || "";
  const titleMatch = /title=["']?([^"'\s}]+)["']?/.exec(meta);
  const startMatch = /startLineNumber=(\d+)/.exec(meta);
  const fileName = titleMatch?.[1] || null;
  const startLine = startMatch ? parseInt(startMatch[1], 10) : 1;
  const headerLabel = fileName
    ? startMatch
      ? `${fileName}:${startLine}`
      : fileName
    : null;

  let highlighted: string;
  if (match) {
    try {
      highlighted = hljs.highlight(raw, { language: match[1], ignoreIllegals: true }).value;
    } catch {
      highlighted = hljs.highlightAuto(raw).value;
    }
  } else {
    highlighted = hljs.highlightAuto(raw).value;
  }

  const lines = highlighted.split("\n");

  const copy = () => {
    copyToClipboard(raw);
    setCopied(true);
    setTimeout(() => setCopied(false), 2000);
  };

  return (
    <div className="group/code relative rounded-md border border-[#2a2a2a] my-4 text-[#ddd] overflow-hidden">
      {headerLabel ? (
        <div className="flex items-stretch">
          <span className="relative flex items-center text-[13px] text-[#999] font-mono px-4 py-2 border-r border-[#2a2a2a]">{headerLabel}<span className="absolute top-0 inset-x-0 h-0.5 bg-white/60 rounded-full" /></span>
          <div className="flex-1 border-b border-[#2a2a2a]" />
          <button
            onClick={copy}
            className="px-3 py-2 text-[#555] hover:text-white transition-colors border-b border-[#2a2a2a]"
            aria-label="Copy code"
          >
            {copied ? (
              <Check className="w-3.5 h-3.5 text-emerald-400" />
            ) : (
              <Copy className="w-3.5 h-3.5" />
            )}
          </button>
        </div>
      ) : (
        <button
          onClick={copy}
          className="absolute top-2 right-2 z-10 p-1 rounded text-[#444] hover:text-white opacity-0 group-hover/code:opacity-100 transition-opacity"
          aria-label="Copy code"
        >
          {copied ? (
            <Check className="w-3.5 h-3.5 text-emerald-400" />
          ) : (
            <Copy className="w-3.5 h-3.5" />
          )}
        </button>
      )}
      <div className="overflow-auto max-h-[400px]">
        <table className="w-full border-collapse font-mono text-[12px] leading-[22px] [font-variant-ligatures:none]">
          <tbody>
            {lines.map((line, i) => (
              <tr key={i}>
                <td className="select-none w-[1px] whitespace-nowrap px-4 text-right text-[#555] align-top text-[12px] leading-[22px] border-r border-[#2a2a2a]">
                  {startLine + i}
                </td>
                <td
                  className="pl-4 pr-4 whitespace-pre"
                  dangerouslySetInnerHTML={{ __html: line || "\n" }}
                />
              </tr>
            ))}
          </tbody>
        </table>
      </div>
    </div>
  );
}

// rehype plugin: pass code fence meta string through to code element properties
type HastNode = {
  type: string;
  tagName?: string;
  children?: HastNode[];
  properties?: Record<string, unknown>;
  data?: { meta?: string };
};

export function rehypeCodeMeta() {
  return (tree: HastNode) => {
    const visit = (node: HastNode) => {
      if (node.type === "element" && node.tagName === "pre" && node.children) {
        const codeEl = node.children.find(
          (c) => c.type === "element" && c.tagName === "code"
        );
        if (codeEl?.data?.meta) {
          codeEl.properties = codeEl.properties || {};
          codeEl.properties.metastring = codeEl.data.meta;
        }
      }
      if (node.children) {
        node.children.forEach((child) => visit(child));
      }
    };
    visit(tree);
  };
}

export const mdComponents = {
  code: MdCodeBlock as React.ComponentType<React.HTMLAttributes<HTMLElement>>,
  pre: ({ children }: { children?: React.ReactNode }) => <>{children}</>,
};