Streaming text

"use client";

import { useEffect, useState, type HTMLAttributes, type ReactNode } from "react";
import "./streaming-text.css";

type StreamingTextBase = HTMLAttributes<HTMLElement> & {
  /** True while more text is on its way. */
  streaming: boolean;
  /** Shows a block cursor at the end while streaming. */
  cursor?: boolean;
  /** Element to render. Use span when the text sits inside other prose. */
  as?: "p" | "div" | "span";
};

export type StreamingTextProps = StreamingTextBase &
  (
    | {
        /** All the text so far. Append to it as chunks arrive; a string that does not extend the last one starts over. */
        text: string;
        segments?: undefined;
      }
    | {
        text?: undefined;
        /**
         * Rich content so far: strings of prose with inline elements, such as citations, between them.
         * Append entries or extend the last string as the stream arrives. Each element counts as one
         * step of the stream and fades in when it arrives.
         */
        segments: ReactNode[];
      }
  );

/** Stands in for one inline element when rich content is measured as text. */
const NODE = "\uFFFC";

/** The segments as one string: prose as is, each element as a single placeholder. */
function project(segments: ReactNode[]): string {
  let out = "";
  for (const seg of segments) {
    if (seg == null || typeof seg === "boolean") continue;
    out += typeof seg === "string" || typeof seg === "number" ? String(seg) : NODE;
  }
  return out;
}

/** The segments cut at the settled point and each fading chunk. Elements keep one stable wrapper. */
function renderSegments(segments: ReactNode[], chunks: number[], settledEnd: number): ReactNode[] {
  const out: ReactNode[] = [];
  let pos = 0;
  segments.forEach((seg, i) => {
    if (seg == null || typeof seg === "boolean") return;
    if (typeof seg === "string" || typeof seg === "number") {
      const str = String(seg);
      const start = pos;
      const end = pos + str.length;
      pos = end;
      const cuts = [settledEnd, ...chunks].filter((c) => c > start && c < end).sort((a, b) => a - b);
      const points = [start, ...cuts, end];
      for (let k = 0; k < points.length - 1; k++) {
        const from = points[k];
        const to = points[k + 1];
        if (from === to) continue;
        const piece = str.slice(from - start, to - start);
        if (from < settledEnd) {
          out.push(piece);
        } else {
          const chunk = chunks.filter((c) => c <= from).pop() ?? settledEnd;
          out.push(
            <span key={`${chunk}:${i}`} className="eq-streaming-text-new">
              {piece}
            </span>,
          );
        }
      }
      return;
    }
    const at = pos;
    pos += 1;
    out.push(
      <span
        key={`node:${i}`}
        className={at >= settledEnd ? "eq-streaming-text-node eq-streaming-text-new" : "eq-streaming-text-node"}
      >
        {seg}
      </span>,
    );
  });
  return out;
}

/** Recently arrived chunks keep their own span so their fade can finish; older ones merge into plain text. */
const FADING_CHUNKS = 8;

/** End of the last finished sentence: a terminator followed by whitespace, so "3." waits for "3.5". */
function lastSentenceEnd(text: string): number {
  let end = 0;
  // A placeholder for an inline element (a citation after the full stop) may sit before the space.
  for (const match of text.matchAll(/[.!?…]+["'”’)\]]*\uFFFC*(?=\s)/g)) end = (match.index ?? 0) + match[0].length;
  return end;
}

type Track = {
  text: string;
  /** Offsets where each still-fading chunk starts. */
  chunks: number[];
  /** How much of the text has been handed to the live region. */
  spoken: number;
  /** Sentences in the live region, keyed by where they start. */
  lines: { at: number; text: string }[];
};

export function StreamingText({
  text: plain,
  segments,
  streaming,
  cursor = true,
  as: Tag = "p",
  className,
  ...rest
}: StreamingTextProps) {
  const rich = segments !== undefined;
  // Rich content is tracked through its text projection, so a new array every render is fine.
  const text = rich ? project(segments) : (plain ?? "");
  const [track, setTrack] = useState<Track>(() => ({
    text,
    chunks: [],
    spoken: streaming ? 0 : text.length,
    lines: [],
  }));

  // Adjust during render so the new tail and the announcement land in the same frame as the text.
  let next = track;
  if (text !== next.text) {
    next = text.startsWith(next.text)
      ? { ...next, text, chunks: [...next.chunks, next.text.length].slice(-FADING_CHUNKS) }
      : { text, chunks: [], spoken: 0, lines: [] };
  }
  const target = streaming ? lastSentenceEnd(text) : text.length;
  if (target > next.spoken) {
    const said = text.slice(next.spoken, target);
    const sentence = (rich ? said.replaceAll(NODE, "") : said).trim();
    next = {
      ...next,
      spoken: target,
      lines: sentence ? [...next.lines, { at: next.spoken, text: sentence }] : next.lines,
    };
  }
  if (next !== track) setTrack(next);

  // Once the stream ends the visible text is exposed, so the live copy is cleared to avoid reading it twice.
  const finished = !streaming && next.lines.length > 0;
  useEffect(() => {
    if (!finished) return;
    const timer = setTimeout(() => setTrack((t) => ({ ...t, lines: [] })), 4000);
    return () => clearTimeout(timer);
  }, [finished]);

  const settledEnd = next.chunks[0] ?? text.length;

  return (
    <Tag
      {...rest}
      className={["eq-streaming-text", className].filter(Boolean).join(" ")}
      data-streaming={streaming || undefined}
    >
      {rich ? (
        // Links inside are hidden from assistive tech while streaming, so they are inert until it ends.
        <span className="eq-streaming-text-visible" aria-hidden={streaming || undefined} inert={streaming || undefined}>
          {renderSegments(segments, next.chunks, settledEnd)}
          {streaming && cursor && <span className="eq-streaming-text-cursor" />}
        </span>
      ) : (
        <span className="eq-streaming-text-visible" aria-hidden={streaming || undefined}>
          {text.slice(0, settledEnd)}
          {next.chunks.map((start, i) => (
            <span key={start} className="eq-streaming-text-new">
              {text.slice(start, next.chunks[i + 1] ?? text.length)}
            </span>
          ))}
          {streaming && cursor && <span className="eq-streaming-text-cursor" />}
        </span>
      )}
      <span className="eq-streaming-text-live" aria-live="polite">
        {next.lines.map((line) => (
          <span key={line.at}>{line.text} </span>
        ))}
      </span>
    </Tag>
  );
}