import { useMemo } from 'react';
import { parseBlocks } from '@/app/lib/blocks';
import { buildSankeyFromProse } from '@/app/lib/sankey/builders';
import { SqlBlock } from './SqlBlock';
import { UnverifiedChart } from './UnverifiedChart';
import { CodeBlock } from './CodeBlock';
import { TextBlock } from './TextBlock';

export function AssistantMessage({ content, isStreaming }: { content: string; isStreaming?: boolean }) {
  // Parsing is the expensive part of rendering a message; redo it only when the text changes.
  const blocks = useMemo(() => parseBlocks(content), [content]);
  // Pure-prose answer: offer a chart of any "Name $X trillion/billion" mentions (unverified: it is the model's text)
  const proseChart = useMemo(
    () => (isStreaming || blocks.some(b => b.type === 'sankey') ? null : buildSankeyFromProse(content)),
    [content, blocks, isStreaming],
  );
  return (
    <div>
      {blocks.map((block, i) => {
        const isLast = i === blocks.length - 1;
        if (block.type === 'sql')    return <SqlBlock key={i} code={block.content} />;
        if (block.type === 'sankey') return <UnverifiedChart key={i} data={block.data} />;
        if (block.type === 'code')   return <CodeBlock   key={i} lang={block.lang} content={block.content} />;
        return (
          <span key={i}>
            <TextBlock content={block.content} />
            {isStreaming && isLast && <span className="cursor-blink" />}
          </span>
        );
      })}
      {isStreaming && blocks.length === 0 && <span className="cursor-blink" />}
      {proseChart && <UnverifiedChart data={proseChart} />}
    </div>
  );
}