import { memo, useMemo } from 'react';
import ReactMarkdown from 'react-markdown';
import remarkGfm from 'remark-gfm';
import remarkMath from 'remark-math';
import rehypeKatex from 'rehype-katex';
import { CodeBlock } from './code-block';
import { cn } from '@/lib/cn';
import { guardCurrency } from '@/lib/math-guard';

/** Splits markdown into stable blocks so only the last block re-renders during streaming. */
function splitBlocks(md: string): string[] {
  const blocks: string[] = [];
  let buf: string[] = [];
  let inFence = false;
  for (const line of md.split('\n')) {
    if (/^\s*(```|~~~)/.test(line)) inFence = !inFence;
    buf.push(line);
    if (!inFence && line.trim() === '' && buf.length > 1) {
      blocks.push(buf.join('\n'));
      buf = [];
    }
  }
  if (buf.length) blocks.push(buf.join('\n'));
  return blocks;
}

const Block = memo(function Block({ md, conversationId, streaming }: { md: string; conversationId?: string; streaming?: boolean }) {
  return (
    <ReactMarkdown
      remarkPlugins={[remarkGfm, remarkMath]}
      rehypePlugins={[[rehypeKatex, { strict: false, throwOnError: false, output: 'html' }]]}
      components={{
        code({ className, children, ...props }) {
          const m = /language-(\w+)/.exec(className || '');
          const text = String(children).replace(/\n$/, '');
          if (m || text.includes('\n')) return <CodeBlock code={text} lang={m?.[1] || ''} conversationId={conversationId} streaming={streaming} />;
          return (
            <code className={className} {...props}>
              {children}
            </code>
          );
        },
        pre({ children }) {
          return <>{children}</>;
        },
        a({ href, children }) {
          return (
            <a href={href} target="_blank" rel="noopener noreferrer">
              {children}
            </a>
          );
        },
        table({ children }) {
          return <table>{children}</table>;
        },
      }}
    >
      {md}
    </ReactMarkdown>
  );
});

export function StreamingText({ text, streaming, conversationId, className }: { text: string; streaming?: boolean; conversationId?: string; className?: string }) {
  const blocks = useMemo(() => splitBlocks(text).map(guardCurrency), [text]);
  return (
    <div className={cn('md', streaming && text.length > 0 && 'streaming-cursor', className)} aria-live={streaming ? 'polite' : undefined}>
      {blocks.map((b, i) => (
        <Block key={i} md={b} conversationId={conversationId} streaming={streaming && i === blocks.length - 1} />
      ))}
    </div>
  );
}