Streaming answer

'use client';

import * as React from 'react';
import type { StreamingContent, StreamingTextStatus } from '../streaming-text/streaming-text';
import { StreamingAnswer, type StreamingAnswerSource } from './streaming-answer';

type Step =
  | { at: number; add: StreamingContent[number] }
  | { at: number; status: StreamingTextStatus }
  | { at: number; source: StreamingAnswerSource };

const SOURCES: StreamingAnswerSource[] = [
  { name: 'Build log', href: 'https://example.com/builds/218' },
  { name: 'Bundle report', href: 'https://example.com/bundle' },
  { name: 'Release notes', href: 'https://example.com/release' },
];

const BUILD: Step[] = [
  { at: 300, add: 'The new build is ready to test. ' },
  { at: 560, add: { type: 'citation', label: 'build-log', href: 'https://example.com/builds/218' } },
  { at: 600, source: SOURCES[0] as StreamingAnswerSource },
  { at: 700, add: ' The first screen now loads in 1.2 seconds, about 30 percent faster' },
  { at: 1700, add: ' than last week, and the main bundle is 18 KB smaller.' },
  { at: 1750, source: SOURCES[1] as StreamingAnswerSource },
  { at: 2100, add: '\n\nTo try it, open the preview link on your phone ' },
  { at: 2350, add: { type: 'citation', label: 'release-notes', href: 'https://example.com/release' } },
  { at: 2400, source: SOURCES[2] as StreamingAnswerSource },
  { at: 2500, add: ' and sign in with your test account.' },
  { at: 2900, status: 'done' },
];

const CUT: Step[] = [
  { at: 400, add: 'The new build is ready to test. ' },
  { at: 700, add: { type: 'citation', label: 'build-log' } },
  { at: 720, source: SOURCES[0] as StreamingAnswerSource },
  { at: 900, add: ' The first screen now loads in 1.2 seconds, about 30 percent faster' },
  { at: 1700, status: 'error' },
];

const FOLLOW_UPS = ['How do I test it on my phone?', 'What made the bundle smaller?'];

function finalState(steps: Step[]) {
  const state = {
    content: [] as StreamingContent,
    status: 'streaming' as StreamingTextStatus,
    sources: [] as StreamingAnswerSource[],
  };
  for (const step of steps) apply(state, step);
  return state;
}

function apply(state: ReturnType<typeof finalState>, step: Step) {
  if ('add' in step) state.content = [...state.content, step.add];
  else if ('status' in step) state.status = step.status;
  else state.sources = [...state.sources, step.source];
}

function useAnswer(steps: Step[]) {
  const [state, setState] = React.useState(() => finalState(steps));
  const timers = React.useRef<number[]>([]);
  const play = React.useCallback(() => {
    for (const id of timers.current) window.clearTimeout(id);
    setState({ content: [], status: 'streaming', sources: [] });
    timers.current = steps.map((step) =>
      window.setTimeout(
        () =>
          setState((s) => {
            const next = { ...s };
            apply(next, step);
            return next;
          }),
        step.at,
      ),
    );
  }, [steps]);
  React.useEffect(
    () => () => {
      for (const id of timers.current) window.clearTimeout(id);
    },
    [],
  );
  return { state, play };
}

const plainText = (content: StreamingContent) =>
  content.map((item) => (typeof item === 'string' ? item : '')).join('');

export default function StreamingAnswerDemo({ replay = 0 }: { replay?: number }) {
  const build = useAnswer(BUILD);
  const cut = useAnswer(CUT);
  const plays = React.useRef({ build: build.play, cut: cut.play });
  plays.current = { build: build.play, cut: cut.play };

  React.useEffect(() => {
    if (replay === 0) return;
    plays.current.build();
    plays.current.cut();
  }, [replay]);

  return (
    <div className="flex w-full max-w-[480px] flex-col gap-12">
      <StreamingAnswer
        content={build.state.content}
        status={build.state.status}
        sources={build.state.sources}
        copyText={plainText(build.state.content)}
        followUps={FOLLOW_UPS}
        actionsDisabled={build.state.status === 'streaming'}
        onRegenerate={build.play}
      />
      <StreamingAnswer
        content={cut.state.content}
        status={cut.state.status}
        sources={cut.state.sources}
        copyText={plainText(cut.state.content)}
        actionsDisabled={cut.state.status === 'streaming'}
        onRetry={cut.play}
        onRegenerate={cut.play}
      />
    </div>
  );
}