boocode/apps/web/src/components/CodeBlock.tsx

import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
import { Check, Copy, Moon, Sun, WrapText } from 'lucide-react';
import { codeToHtml } from 'shiki';

// NOTE: spec calls for syntax-highlighted code blocks. Added Shiki in v1.1.
// Shiki output is compiler-generated and does not contain user input; setting
// it via a ref is safe here.
interface Props {
  code: string;
  lang?: string;
}

const LANG_MAP: Record<string, string> = {
  ts: 'typescript',
  tsx: 'tsx',
  typescript: 'typescript',
  js: 'javascript',
  jsx: 'jsx',
  javascript: 'javascript',
  py: 'python',
  python: 'python',
  go: 'go',
  rs: 'rust',
  rust: 'rust',
  rb: 'ruby',
  ruby: 'ruby',
  java: 'java',
  c: 'c',
  cpp: 'cpp',
  cs: 'csharp',
  csharp: 'csharp',
  php: 'php',
  sh: 'bash',
  bash: 'bash',
  shell: 'bash',
  yaml: 'yaml',
  yml: 'yaml',
  json: 'json',
  toml: 'toml',
  md: 'markdown',
  markdown: 'markdown',
  sql: 'sql',
  dockerfile: 'dockerfile',
  html: 'html',
  css: 'css',
};

// ── LRU highlight cache (module-scoped) ──────────────────────────
// Key = `${code}|${theme}|${mappedLang}`, max 50 entries.
// Avoids redundant codeToHtml calls when the same code/theme/lang
// combination is rendered multiple times (e.g. across messages).
const HIGHLIGHT_CACHE = new Map<string, string>();
const MAX_CACHE_ENTRIES = 50;

function cacheGet(key: string): string | undefined {
  if (!HIGHLIGHT_CACHE.has(key)) return undefined;
  const val = HIGHLIGHT_CACHE.get(key)!;
  // LRU touch — delete & re-set to move to end (most recently used)
  HIGHLIGHT_CACHE.delete(key);
  HIGHLIGHT_CACHE.set(key, val);
  return val;
}

function cacheSet(key: string, html: string): void {
  if (HIGHLIGHT_CACHE.size >= MAX_CACHE_ENTRIES) {
    const oldest = HIGHLIGHT_CACHE.keys().next().value;
    if (oldest !== undefined) HIGHLIGHT_CACHE.delete(oldest);
  }
  HIGHLIGHT_CACHE.set(key, html);
}

export function CodeBlock({ code, lang }: Props) {
  const [copied, setCopied] = useState(false);
  const [html, setHtml] = useState<string | null>(null);
  const [theme, setTheme] = useState<'github-dark' | 'github-light'>(() => {
    try {
      if (localStorage.getItem('codeblock-theme') === 'github-light') return 'github-light';
    } catch {
      /* localStorage unavailable */
    }
    return 'github-dark';
  });
  const [wordWrap, setWordWrap] = useState(false);
  const [expanded, setExpanded] = useState(false);
  const highlightRef = useRef<HTMLDivElement>(null);

  // ── Derived state ──────────────────────────────────────────────

  // Diff mode: detect `diff-` prefix (e.g. diff-ts, diff-py).
  // The actual lang for highlighting is the part after `diff-`.
  const isDiff = !!lang && lang.startsWith('diff-');
  const actualLang = isDiff && lang ? lang.slice('diff-'.length) : lang ?? '';
  const mappedLang = actualLang ? (LANG_MAP[actualLang.toLowerCase()] ?? null) : null;

  // Strip leading `+`/`-` from code lines when in diff mode.
  // The markers are rendered in the gutter instead.
  const cleanCode = useMemo(
    () => (isDiff ? code.replace(/^[+-]/gm, '') : code),
    [code, isDiff],
  );

  const codeLines = useMemo(() => code.split('\n'), [code]);
  const totalLines = codeLines.length;

  // Gutter is hidden entirely when code has >= 1000 lines.
  const showGutter = totalLines < 1000;

  // Collapsible: auto-collapse to 15 lines when >= 30 lines total.
  const isLong = totalLines >= 30;
  const collapsed = isLong && !expanded;
  const visibleLines = collapsed ? codeLines.slice(0, 15) : codeLines;

  // Diff marker array: '+' / '-' / '' per line for the gutter.
  const diffMarkers = useMemo(
    () => (isDiff ? codeLines.map((l) => (l[0] === '+' ? '+' : l[0] === '-' ? '-' : '')) : null),
    [isDiff, codeLines],
  );

  // ── Shiki highlighting ─────────────────────────────────────────

  useEffect(() => {
    let cancelled = false;

    if (!mappedLang) {
      setHtml(null);
      return;
    }

    const cacheKey = `${cleanCode}|${theme}|${mappedLang}`;
    const cached = cacheGet(cacheKey);
    if (cached !== undefined) {
      setHtml(cached);
      return;
    }

    (async () => {
      try {
        const result = await codeToHtml(cleanCode, { lang: mappedLang, theme });
        cacheSet(cacheKey, result);
        if (!cancelled) setHtml(result);
      } catch (err) {
        console.warn('shiki highlight failed:', err);
        if (!cancelled) setHtml(null);
      }
    })();

    return () => {
      cancelled = true;
    };
  }, [cleanCode, mappedLang, theme]);

  // Inject Shiki HTML via ref (output is compiler-generated, not user input)
  useEffect(() => {
    if (highlightRef.current) {
      // Shiki generates sanitized HTML spans — not user-supplied content.
      // eslint-disable-next-line no-unsanitized/property
      highlightRef.current.innerHTML = html ?? '';
    }
  }, [html]);

  // Sync word-wrap state to the injected <pre> element inside shiki's output
  useEffect(() => {
    const pre = highlightRef.current?.querySelector('pre');
    if (pre) {
      pre.style.whiteSpace = wordWrap ? 'pre-wrap' : 'nowrap';
    }
  }, [html, wordWrap]);

  // ── Handlers ───────────────────────────────────────────────────

  const handleToggleTheme = useCallback(() => {
    setTheme((prev) => {
      const next = prev === 'github-dark' ? 'github-light' : 'github-dark';
      try {
        localStorage.setItem('codeblock-theme', next);
      } catch {
        /* noop */
      }
      return next;
    });
  }, []);

  const handleCopy = useCallback(async () => {
    try {
      await navigator.clipboard.writeText(cleanCode);
      setCopied(true);
      setTimeout(() => setCopied(false), 1200);
    } catch {
      /* ignore */
    }
  }, [cleanCode]);

  // ── Shared class segments ──────────────────────────────────────

  const preBaseClass = 'overflow-x-auto px-3 py-2 font-mono text-xs leading-relaxed';
  const shikiWrapperClass = `${preBaseClass} [&>pre]:!bg-transparent [&>pre]:!m-0 [&>pre]:!p-0`;
  const collapsedClass = collapsed ? 'max-h-[calc(15*1.625em)] overflow-hidden' : '';

  // ── Render ─────────────────────────────────────────────────────

  return (
    <div className="rounded-md border bg-muted/40 overflow-hidden text-sm my-1">
      {/* ── Toolbar ──────────────────────────────────────────── */}
      <div className="flex items-center justify-between px-2 py-1 border-b text-xs text-muted-foreground">
        <span className="font-mono">{actualLang || lang || 'code'}</span>
        <div className="flex items-center gap-0.5">
          {/* Theme toggle — persists to localStorage key 'codeblock-theme' */}
          <button
            type="button"
            onClick={handleToggleTheme}
            className="flex items-center gap-1 px-1.5 py-0.5 rounded hover:bg-muted text-foreground"
            aria-label={`Switch to ${theme === 'github-dark' ? 'light' : 'dark'} theme`}
          >
            {theme === 'github-dark' ? <Sun className="size-3" /> : <Moon className="size-3" />}
          </button>

          {/* Word-wrap toggle */}
          <button
            type="button"
            onClick={() => setWordWrap((prev) => !prev)}
            className={`flex items-center gap-1 px-1.5 py-0.5 rounded hover:bg-muted text-foreground ${wordWrap ? 'bg-muted' : ''}`}
            aria-label={wordWrap ? 'Disable word wrap' : 'Enable word wrap'}
          >
            <WrapText className="size-3" />
          </button>

          {/* Copy button — existing behavior (Check icon, 1200ms revert) */}
          <button
            type="button"
            onClick={() => void handleCopy()}
            className="flex items-center gap-1 px-1.5 py-0.5 rounded hover:bg-muted text-foreground"
            aria-label="Copy code"
          >
            {copied ? <Check className="size-3" /> : <Copy className="size-3" />}
            <span>{copied ? 'Copied' : 'Copy'}</span>
          </button>
        </div>
      </div>

      {/* ── Code body (flex row: gutter + code) ──────────────── */}
      <div className="flex">
        {/* Gutter — line numbers or diff markers */}
        {showGutter && (
          <div
            className="flex-none select-none text-right text-muted-foreground/50 font-mono text-xs leading-relaxed py-2 border-r border-border/30"
            aria-hidden="true"
          >
            {visibleLines.map((line, i) => {
              const isPlus = diffMarkers?.[i] === '+';
              const isMinus = diffMarkers?.[i] === '-';
              let gutterCellClass = 'px-2 leading-relaxed';
              if (isPlus) gutterCellClass += ' bg-green-500/10 border-l-2 border-green-500';
              if (isMinus) gutterCellClass += ' bg-red-500/10 border-l-2 border-red-500';
              const content = diffMarkers ? diffMarkers[i] : String(i + 1);
              return (
                <div key={i} className={gutterCellClass}>
                  {content}
                </div>
              );
            })}
          </div>
        )}

        {/* Code area */}
        <div className="relative flex-1 min-w-0">
          {html !== null ? (
            <div ref={highlightRef} className={`${shikiWrapperClass} ${collapsedClass}`} />
          ) : (
            <pre
              className={`${preBaseClass} ${collapsedClass}`}
              style={{ whiteSpace: wordWrap ? 'pre-wrap' : 'nowrap' }}
            >
              {collapsed ? codeLines.slice(0, 15).join('\n') : code}
            </pre>
          )}

          {/* Gradient fade overlay for collapsed state */}
          {collapsed && (
            <div className="absolute bottom-0 left-0 right-0 h-12 bg-gradient-to-b from-transparent to-background pointer-events-none" />
          )}
        </div>
      </div>

      {/* "Show N more" button for collapsed state */}
      {collapsed && (
        <button
          type="button"
          onClick={() => setExpanded(true)}
          className="w-full text-xs text-muted-foreground hover:text-foreground py-1 border-t border-border/30 bg-muted/20"
        >
          Show {totalLines - 15} more {totalLines - 15 === 1 ? 'line' : 'lines'}
        </button>
      )}
    </div>
  );
}