frontend/src/Components/AiComponents/ChatComponents/Streaming.js

import React, { useEffect, useRef } from 'react';
import ReactMarkdown from 'react-markdown';
import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
import { atomDark } from 'react-syntax-highlighter/dist/esm/styles/prism';
import remarkGfm from 'remark-gfm';
import rehypeRaw from 'rehype-raw';
import './Streaming.css';
import './SourceRef.css';

// Helper function to normalize various citation formats (e.g., [1,2], [1, 2]) into the standard [1][2] format
const normalizeCitations = (text) => {
  if (!text) return '';

  const citationRegex = /\[(\d+(?:,\s*\d+)+)\]/g;

  return text.replace(citationRegex, (match, capturedNumbers) => {
    const numbers = capturedNumbers
                      .split(/,\s*/)
                      .map(numStr => numStr.trim())
                      .filter(Boolean);

    if (numbers.length <= 1) {
      return match;
    }

    return numbers.map(num => `[${num}]`).join('');
  });
};

// Streaming component for rendering markdown content
const Streaming = ({ content, isStreaming, onContentRef, showSourcePopup, hideSourcePopup }) => {
  const contentRef = useRef(null);
  
  useEffect(() => {
    if (contentRef.current && onContentRef) {
      onContentRef(contentRef.current);
    }
  }, [content, onContentRef]);

  const displayContent = isStreaming ? `${content}▌` : (content || '');
  const normalizedContent = normalizeCitations(displayContent);

  // Custom renderer for text nodes to handle source references
  const renderWithSourceRefs = (elementType) => {
    const ElementComponent = elementType; // e.g., 'p', 'li'
  
    // Helper to gather plain text
    const getFullText = (something) => {
      if (typeof something === 'string') return something;
      if (Array.isArray(something)) return something.map(getFullText).join('');
      if (React.isValidElement(something) && something.props?.children)
        return getFullText(React.Children.toArray(something.props.children));
      return '';
    };
  
    return (props) => {
      // Plain‑text version of this block (paragraph / list‑item)
      const fullText = getFullText(props.children);
      // Same regex the backend used
      const sentenceRegex = /[^.!?\n]+[.!?]+[\])'"`’”]*|[^.!?\n]+$/g;
      const sentencesArr = fullText.match(sentenceRegex) || [fullText];
  
      // Helper function to find the sentence that contains position `pos`
      const sentenceByPos = (pos) => {
        let run = 0;
        for (const s of sentencesArr) {
          const end = run + s.length;
          if (pos >= run && pos < end) return s.trim();
          run = end;
        }
        return fullText.trim();
      };
  
      // Cursor that advances through fullText so each subsequent 
      // indexOf search starts AFTER the previous match
      let searchCursor = 0;                                         
  
      // Recursive renderer that preserves existing markup
      const processNode = (node, keyPrefix = 'node') => {
        if (typeof node === 'string') {
          const citationRegex = /\[(\d+)\]/g;
          let last = 0;
          let parts = [];
          let m;
          while ((m = citationRegex.exec(node))) {
            const sliceBefore = node.slice(last, m.index);
            if (sliceBefore) parts.push(sliceBefore);
  
            const localIdx  = m.index;
            const num       = parseInt(m[1], 10);
            const citStr    = m[0];
  
            // Find this specific occurrence in fullText, starting at searchCursor
            const absIdx    = fullText.indexOf(citStr, searchCursor);
            if (absIdx !== -1) searchCursor = absIdx + citStr.length;
  
            const sentenceForPopup = sentenceByPos(absIdx);          
  
            parts.push(
              <sup
                key={`${keyPrefix}-ref-${num}-${localIdx}`}
                className="source-reference"
                onMouseEnter={(e) =>
                  showSourcePopup &&
                  showSourcePopup(num - 1, e.target, sentenceForPopup)
                }                                                  
                onMouseLeave={hideSourcePopup}
              >
                {num}
              </sup>
            );
            last = localIdx + citStr.length;
          }
          if (last < node.length) parts.push(node.slice(last));
          return parts;
        }
  
        // For non‑string children, recurse (preserves <em>, <strong>, links, etc.)
        if (React.isValidElement(node) && node.props?.children) {
          const processed = React.Children.map(node.props.children, (child, i) =>
            processNode(child, `${keyPrefix}-${i}`)
          );
          return React.cloneElement(node, { children: processed });
        }
  
        return node; // element without children or unknown type
      };
  
      const processedChildren = React.Children.map(props.children, (child, i) =>
        processNode(child, `root-${i}`)
      );
  
      // Render original element (p, li, …) with processed children
      return <ElementComponent {...props}>{processedChildren}</ElementComponent>;
    };
  };  

  return (
    <div className="streaming-content" ref={contentRef}>
      <ReactMarkdown
        remarkPlugins={[remarkGfm]}
        rehypePlugins={[rehypeRaw]}
        components={{
          p: renderWithSourceRefs('p'),
          li: renderWithSourceRefs('li'),
          // Add other block elements if citations might appear directly within them
          // blockquote: renderWithSourceRefs('blockquote'),
          // div: renderWithSourceRefs('div'), // Be cautious with generic divs
          code({node, inline, className, children, ...props}) {
            const match = /language-(\w+)/.exec(className || '');
            return !inline ? (
              <div className="code-block-container">
                <div className="code-block-header">
                  <span>{match ? match[1] : 'code'}</span>
                </div>
                <SyntaxHighlighter
                  style={atomDark}
                  language={match ? match[1] : 'text'}
                  PreTag="div"
                  {...props}
                >
                  {String(children).replace(/\n$/, '')}
                </SyntaxHighlighter>
              </div>
            ) : (
              <code className={className} {...props}>
                {children}
              </code>
            );
          },
          table({node, ...props}) {
            return (
              <div className="table-container">
                <table {...props} />
              </div>
            );
          },
          a({node, children, href, ...props}) {
            return (
              <a 
                href={href}
                target="_blank"
                rel="noopener noreferrer"
                className="markdown-link"
                {...props}
              >
                {children}
              </a>
            );
          },
          blockquote({node, ...props}) {
            return (
              <blockquote className="markdown-blockquote" {...props} />
            );
          }
        }}
      >
        {normalizedContent}
      </ReactMarkdown>
    </div>
  );
};

export default Streaming;