src/app/components/StreamingCodeBlock.tsx

import { CodeToTokenTransformStream, type RecallToken } from '@shikijs/stream';
import { ShikiStreamRenderer } from '@shikijs/stream/react';
import type { ThemedToken } from 'shiki/core';
import { useEffect, useRef, useState } from 'react';
import {
  codeThemeName,
  getCodeHighlighter,
  normalizeCodeLanguage,
  type CodeTheme,
} from '../../lib/code-highlighter';

interface StreamingCodeBlockProps {
  code: string;
  language: string;
  theme: CodeTheme;
  complete: boolean;
}

export function StreamingCodeBlock({
  code,
  language,
  theme,
  complete,
}: StreamingCodeBlockProps) {
  const normalizedLanguage = normalizeCodeLanguage(language);
  const [revision, setRevision] = useState(0);
  const [tokenStream, setTokenStream] = useState<ReadableStream<ThemedToken | RecallToken> | null>(null);
  const [highlightFailed, setHighlightFailed] = useState(false);
  const writerRef = useRef<WritableStreamDefaultWriter<string> | null>(null);
  const writtenRef = useRef('');
  const closedRef = useRef(false);
  const codeRef = useRef(code);
  const completeRef = useRef(complete);
  const queueRef = useRef<Promise<void>>(Promise.resolve());

  codeRef.current = code;
  completeRef.current = complete;

  const queueWrite = (
    writer: WritableStreamDefaultWriter<string>,
    chunk: string,
    closeAfter: boolean,
  ) => {
    queueRef.current = queueRef.current
      .then(async () => {
        if (chunk) await writer.write(chunk);
        if (closeAfter && !closedRef.current) {
          closedRef.current = true;
          await writer.close();
        }
      })
      .catch(() => {
        if (writerRef.current === writer) {
          writerRef.current = null;
          setHighlightFailed(true);
        }
      });
  };

  useEffect(() => {
    let disposed = false;
    let writer: WritableStreamDefaultWriter<string> | null = null;
    setTokenStream(null);
    setHighlightFailed(false);
    writtenRef.current = '';
    closedRef.current = false;
    queueRef.current = Promise.resolve();

    void getCodeHighlighter(normalizedLanguage)
      .then((highlighter) => {
        if (disposed) return;
        const transform = new CodeToTokenTransformStream({
          highlighter,
          lang: normalizedLanguage,
          theme: codeThemeName(theme),
          allowRecalls: true,
        });
        writer = transform.writable.getWriter();
        writerRef.current = writer;
        const currentCode = codeRef.current;
        writtenRef.current = currentCode;
        setTokenStream(transform.readable);
        queueWrite(writer, currentCode, completeRef.current);
      })
      .catch(() => {
        if (!disposed) setHighlightFailed(true);
      });

    return () => {
      disposed = true;
      if (writer && writerRef.current === writer) writerRef.current = null;
      if (writer && !closedRef.current) {
        closedRef.current = true;
        void writer.close().catch(() => undefined);
      }
    };
  }, [normalizedLanguage, revision, theme]);

  useEffect(() => {
    const writer = writerRef.current;
    if (!writer || closedRef.current) return;
    const written = writtenRef.current;
    if (!code.startsWith(written)) {
      writerRef.current = null;
      closedRef.current = true;
      void writer.close().catch(() => undefined);
      setRevision((current) => current + 1);
      return;
    }
    const delta = code.slice(written.length);
    writtenRef.current = code;
    if (delta || complete) queueWrite(writer, delta, complete);
  }, [code, complete]);

  const label = language.trim() || 'plain text';
  const useFallback = !tokenStream || highlightFailed;

  return (
    <figure
      className="code-block"
      data-language={normalizedLanguage}
      data-code-complete={complete ? 'true' : 'false'}
    >
      <figcaption>
        <span>{label}</span>
        {!complete && <span className="code-stream-state"><i /> streaming</span>}
      </figcaption>
      {useFallback
        ? <pre className="code-fallback"><code>{code}</code></pre>
        : <ShikiStreamRenderer stream={tokenStream} />}
    </figure>
  );
}