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

import React, { useEffect, useRef, useState } from 'react';
import Markdown from 'markdown-to-jsx';
import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter';
import { vscDarkPlus } from 'react-syntax-highlighter/dist/esm/styles/prism';
import { FcCheckmark } from 'react-icons/fc';
import { BiSolidCopy } from 'react-icons/bi';
import { LuDownload } from 'react-icons/lu';
import katex from 'katex';
import 'katex/dist/katex.min.css';
import './Streaming.css';
import './SourceRef.css';

// Math rendering components
const InlineMath = ({ children }) => {
  const mathRef = useRef(null);
  
  useEffect(() => {
    if (mathRef.current && children) {
      try {
        katex.render(children.toString(), mathRef.current, {
          throwOnError: false,
          displayMode: false
        });
      } catch (e) {
        mathRef.current.innerHTML = `<span style="color: red;">Error: ${e.message}</span>`;
      }
    }
  }, [children]);
  
  return <span ref={mathRef} className="inline-math" />;
};

const BlockMath = ({ children }) => {
  const mathRef = useRef(null);
  
  useEffect(() => {
    if (mathRef.current && children) {
      try {
        katex.render(children.toString(), mathRef.current, {
          throwOnError: false,
          displayMode: true
        });
      } catch (e) {
        mathRef.current.innerHTML = `<div style="color: red;">Error: ${e.message}</div>`;
      }
    }
  }, [children]);
  
  return <div ref={mathRef} className="block-math" />;
};

// 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 '';

  // First, temporarily replace math expressions to protect them from citation processing
  const mathPlaceholders = [];
  let mathIndex = 0;
  
  // Replace block math
  text = text.replace(/\$\$([\s\S]*?)\$\$/g, (match) => {
    const placeholder = `__BLOCK_MATH_${mathIndex}__`;
    mathPlaceholders[mathIndex] = match;
    mathIndex++;
    return placeholder;
  });
  
  // Replace inline math
  text = text.replace(/\$([^\$\n]+?)\$/g, (match) => {
    const placeholder = `__INLINE_MATH_${mathIndex}__`;
    mathPlaceholders[mathIndex] = match;
    mathIndex++;
    return placeholder;
  });

  // Process citations
  const citationRegex = /\[(\d+(?:,\s*\d+)+)\]/g;
  text = 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('');
  });

  // Restore math expressions
  mathPlaceholders.forEach((math, index) => {
    const placeholder = math.startsWith('$$') ? `__BLOCK_MATH_${index}__` : `__INLINE_MATH_${index}__`;
    text = text.replace(placeholder, math);
  });

  return text;
};

// 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);

  // CodeBlock component with copy functionality
  const CodeBlock = ({ language, codeString }) => {
    const [copied, setCopied] = useState(false);
    const [showCopyTooltip, setShowCopyTooltip] = useState(true);

    const handleCopy = () => {
      const textToCopy = String(codeString).replace(/\n$/, '');
      navigator.clipboard.writeText(textToCopy).then(() => {
        setCopied(true);
        setShowCopyTooltip(false);
        setTimeout(() => setCopied(false), 2000);
      }).catch(err => {
        console.error('Failed to copy:', err);
      });
    };

    return (
      <div className="code-block-container">
        <div className="code-block-header">
          <span>{language}</span>
          <button className={`code-copy-button ${copied ? 'copied' : ''}`} onClick={handleCopy}>
            {copied ? <FcCheckmark /> : <BiSolidCopy />}
            {showCopyTooltip && <span className="tooltip">Copy Code</span>}
          </button>
        </div>
        <SyntaxHighlighter
          style={vscDarkPlus}
          language={language}
          PreTag="div"
        >
          {String(codeString).replace(/\n$/, '')}
        </SyntaxHighlighter>
      </div>
    );
  };

  // TableWithCopy component with copy and download functionality
  const TableWithCopy = ({ children, ...props }) => {
    const [copied, setCopied] = useState(false);
    const [fadingCopy, setFadingCopy] = useState(false);
    const [showCopyTooltip, setShowCopyTooltip] = useState(true);
    const [downloaded, setDownloaded] = useState(false);
    const [fadingDownload, setFadingDownload] = useState(false);
    const [showDownloadTooltip, setShowDownloadTooltip] = useState(true);
    const tableRef = useRef(null);

    const handleCopy = () => {
      if (tableRef.current) {
        const rows = Array.from(tableRef.current.querySelectorAll('tr'));
        const csvData = rows.map(row => {
          const cells = Array.from(row.querySelectorAll('th, td'));
          return cells.map(cell => {
            let text = cell.innerText.trim();
            if (text.includes('"') || text.includes(',') || text.includes('\n')) {
              text = '"' + text.replace(/"/g, '""') + '"';
            }
            return text;
          }).join(',');
        }).join('\n');
        navigator.clipboard.writeText(csvData).then(() => {
          setCopied(true);
          setShowCopyTooltip(false);
          setTimeout(() => {
            setFadingCopy(true);
            setTimeout(() => {
              setCopied(false);
              setFadingCopy(false);
            }, 200);
          }, 2000);
        }).catch(err => {
          console.error('Failed to copy table:', err);
        });
      }
    };

    const handleDownload = () => {
      if (tableRef.current) {
        const rows = Array.from(tableRef.current.querySelectorAll('tr'));
        const tsvData = rows.map(row => {
          const cells = Array.from(row.querySelectorAll('th, td'));
          return cells.map(cell => {
            let text = cell.innerText.trim();
            if (text.includes('"') || text.includes('\t') || text.includes('\n')) {
              text = '"' + text.replace(/"/g, '""') + '"';
            }
            return text;
          }).join('\t');
        }).join('\n');
        const blob = new Blob([tsvData], { type: 'application/vnd.ms-excel' });
        const url = URL.createObjectURL(blob);
        const a = document.createElement('a');
        a.href = url;
        a.download = 'table.xls';
        a.click();
        URL.revokeObjectURL(url);
        setDownloaded(true);
        setShowDownloadTooltip(false);
        setTimeout(() => {
          setFadingDownload(true);
          setTimeout(() => {
            setDownloaded(false);
            setFadingDownload(false);
          }, 200);
        }, 2000);
      }
    };

    return (
      <div className="table-container">
        <button className={`table-copy-button ${copied ? 'copied' : ''} ${fadingCopy ? 'fading' : ''}`} onClick={handleCopy}>
          {copied ? <FcCheckmark /> : <BiSolidCopy />}
          {showCopyTooltip && <span className="tooltip">Copy Table</span>}
        </button>
        <table ref={tableRef} {...props}>{children}</table>
        <button className={`table-download-button ${downloaded ? 'downloaded' : ''} ${fadingDownload ? 'fading' : ''}`} onClick={handleDownload}>
          {downloaded ? <FcCheckmark /> : <LuDownload />}
          {showDownloadTooltip && <span className="tooltip">Download Table</span>}
        </button>
      </div>
    );
  };

  // Custom renderer for text nodes to handle source references and math
  const renderWithSourceRefsAndMath = (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 ({ children, ...props }) => {
      // Plain‑text version of this block (paragraph / list‑item)
      const fullText = getFullText(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 and adds math support
      const processNode = (node, keyPrefix = 'node') => {
        if (typeof node === 'string') {
          const parts = [];
          let lastIndex = 0;
          
          // First process block math
          const blockMathRegex = /\$\$([\s\S]*?)\$\$/g;
          let blockMatch;
          
          while ((blockMatch = blockMathRegex.exec(node))) {
            // Process text before the math expression for citations
            const textBefore = node.slice(lastIndex, blockMatch.index);
            if (textBefore) {
              parts.push(...processCitations(textBefore, keyPrefix, lastIndex));
            }
            
            // Add the block math component
            parts.push(
              <BlockMath key={`${keyPrefix}-block-math-${blockMatch.index}`}>
                {blockMatch[1]}
              </BlockMath>
            );
            
            lastIndex = blockMatch.index + blockMatch[0].length;
          }
          
          // Process remaining text for inline math and citations
          const remainingText = node.slice(lastIndex);
          if (remainingText) {
            parts.push(...processInlineMathAndCitations(remainingText, keyPrefix, lastIndex));
          }
          
          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
      };
      
      // Helper function to process inline math and citations
      const processInlineMathAndCitations = (text, keyPrefix, offset) => {
        const parts = [];
        let lastIndex = 0;
        
        // Combined regex for inline math and citations
        const combinedRegex = /\$([^\$\n]+?)\$|\[(\d+)\]/g;
        let match;
        
        while ((match = combinedRegex.exec(text))) {
          // Add text before the match
          if (match.index > lastIndex) {
            parts.push(text.slice(lastIndex, match.index));
          }
          
          if (match[1] !== undefined) {
            // It's inline math
            parts.push(
              <InlineMath key={`${keyPrefix}-inline-math-${match.index}`}>
                {match[1]}
              </InlineMath>
            );
          } else if (match[2] !== undefined) {
            // It's a citation
            const num = parseInt(match[2], 10);
            const absIdx = fullText.indexOf(match[0], searchCursor);
            if (absIdx !== -1) searchCursor = absIdx + match[0].length;
            const sentenceForPopup = sentenceByPos(absIdx);
            
            parts.push(
              <sup
                key={`${keyPrefix}-ref-${num}-${match.index}`}
                className="source-reference"
                onMouseEnter={(e) =>
                  showSourcePopup &&
                  showSourcePopup(num - 1, e.target, sentenceForPopup)
                }
                onMouseLeave={hideSourcePopup}
              >
                {num}
              </sup>
            );
          }
          
          lastIndex = match.index + match[0].length;
        }
        
        // Add remaining text
        if (lastIndex < text.length) {
          parts.push(text.slice(lastIndex));
        }
        
        return parts;
      };
      
      // Helper function to process only citations
      const processCitations = (text, keyPrefix, offset) => {
        const citationRegex = /\[(\d+)\]/g;
        let last = 0;
        let parts = [];
        let m;
        
        while ((m = citationRegex.exec(text))) {
          const sliceBefore = text.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 < text.length) parts.push(text.slice(last));
        return parts;
      };
  
      const processedChildren = React.Children.map(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}>
      <Markdown
        options={{
          wrapper: React.Fragment, // Use Fragment to avoid wrapper div
          forceBlock: false,
          forceInline: false,
          overrides: {
            p: {
              component: renderWithSourceRefsAndMath('p')
            },
            li: {
              component: renderWithSourceRefsAndMath('li')
            },
            h1: {
              component: renderWithSourceRefsAndMath('h1')
            },
            h2: {
              component: renderWithSourceRefsAndMath('h2')
            },
            h3: {
              component: renderWithSourceRefsAndMath('h3')
            },
            h4: {
              component: renderWithSourceRefsAndMath('h4')
            },
            h5: {
              component: renderWithSourceRefsAndMath('h5')
            },
            h6: {
              component: renderWithSourceRefsAndMath('h6')
            },
            pre: {
              component: ({ children, ...props }) => {
                let codeElement = null;
                let codeString = '';
                let language = '';

                React.Children.forEach(children, (child) => {
                  if (React.isValidElement(child)) {  // Check for element, no type restriction
                    codeElement = child;
                    const className = child.props?.className || '';
                    const match = /(?:lang|language)-(\w+)/.exec(className);
                    if (match) {
                      language = match[1];
                    }
                    // Flatten children safely
                    codeString = React.Children.toArray(child.props?.children).join('');
                  }
                });

                // Always start code blocks on a new line
                return (
                  <>
                    <div style={{ display: 'block', width: '100%', height: 0, margin: 0, padding: 0 }} />
                    {codeElement ? (
                      <CodeBlock 
                        language={language || 'plaintext'}  // Default for no-lang blocks
                        codeString={codeString} 
                      />
                    ) : (
                      <pre {...props}>{children}</pre>
                    )}
                  </>
                );
              }
            },
            code: {
              component: ({ className, children, ...props }) => {
                // This handles inline code only
                // Block code is handled by the pre component above
                return (
                  <code className={className} {...props}>
                    {children}
                  </code>
                );
              }
            },
            table: {
              component: TableWithCopy
            },
            a: {
              component: ({ children, href, ...props }) => {
                return (
                  <a 
                    href={href}
                    target="_blank"
                    rel="noopener noreferrer"
                    className="markdown-link"
                    {...props}
                  >
                    {children}
                  </a>
                );
              }
            },
            blockquote: {
              component: ({ children, ...props }) => {
                return (
                  <blockquote className="markdown-blockquote" {...props}>
                    {children}
                  </blockquote>
                );
              }
            }
          }
        }}
      >
        {normalizedContent}
      </Markdown>
    </div>
  );
};

export default Streaming;