frontend/src/Components/AiComponents/ChatWindow.js

import React, { useRef, useState, useCallback, useEffect } from 'react';
import Box from '@mui/material/Box';
import Snackbar from '@mui/material/Snackbar';
import Slide from '@mui/material/Slide';
import IconButton from '@mui/material/IconButton';
import { FaTimes, FaSpinner, FaCheckCircle } from 'react-icons/fa';
import GraphDialog from './ChatComponents/Graph';
import Streaming from './ChatComponents/Streaming';
import SourcePopup from './ChatComponents/SourcePopup';
import './ChatWindow.css';

import bot from '../../Icons/bot.png';
import copy from '../../Icons/copy.png';
import evaluate from '../../Icons/evaluate.png';
import sourcesIcon from '../../Icons/sources.png';
import graphIcon from '../../Icons/graph.png';
import user from '../../Icons/user.png';
import excerpts from '../../Icons/excerpts.png';

// SlideTransition function for both entry and exit transitions.
function SlideTransition(props) {
  return <Slide {...props} direction="up" />;
}

function ChatWindow({
  blockId,
  userMessage,
  tokenChunks,
  aiAnswer,
  thinkingTime,
  thoughtLabel,
  sourcesRead,
  finalSources,
  excerptsData,
  isLoadingExcerpts,
  onFetchExcerpts,
  actions,
  tasks,
  openRightSidebar,
  // openLeftSidebar,
  isError,
  errorMessage
}) {
  console.log(`[ChatWindow ${blockId}] Received excerptsData:`, excerptsData);
  const answerRef = useRef(null);
  const [graphDialogOpen, setGraphDialogOpen] = useState(false);
  const [snackbarOpen, setSnackbarOpen] = useState(false);
  const [hoveredSourceInfo, setHoveredSourceInfo] = useState(null);
  const popupTimeoutRef = useRef(null);

  // Get the graph action from the actions prop.
  const graphAction = actions && actions.find(a => a.name === "graph");

  // Handler for copying answer to clipboard.
  const handleCopy = () => {
    if (answerRef.current) {
      const textToCopy = answerRef.current.innerText || answerRef.current.textContent;
      navigator.clipboard.writeText(textToCopy)
        .then(() => {
          console.log('Copied to clipboard:', textToCopy);
          setSnackbarOpen(true);
        })
        .catch((err) => console.error('Failed to copy text:', err));
    }
  };

  // Snackbar close handler
  const handleSnackbarClose = (event, reason) => {
    if (reason === 'clickaway') return;
    setSnackbarOpen(false);
  };

  // Combine partial chunks (tokenChunks) if present; else fall back to the aiAnswer string.
  const combinedAnswer = (tokenChunks && tokenChunks.length > 0)
  ? tokenChunks.join("")
  : aiAnswer;
  const hasTokens = combinedAnswer && combinedAnswer.length > 0;
  // Assume streaming is in progress if thinkingTime is not set.
  const isStreaming = thinkingTime === null || thinkingTime === undefined;

  // Helper to render the thought label.
  const renderThoughtLabel = () => {
    if (!hasTokens) {
      return thoughtLabel;
    } else {
      if (thoughtLabel && thoughtLabel.startsWith("Thought and searched for")) {
        return thoughtLabel;
      }
      return null;
    }
  };

  // Helper to render sources read.
  const renderSourcesRead = () => {
    if (!sourcesRead && sourcesRead !== 0) return null;
    return sourcesRead;
  };

  // When tasks first appear, automatically open the sidebar.
  const prevTasksRef = useRef(tasks);
  useEffect(() => {
    if (prevTasksRef.current.length === 0 && tasks && tasks.length > 0) {
      openRightSidebar("tasks", blockId);
    }
    prevTasksRef.current = tasks;
  }, [tasks, blockId, openRightSidebar]);

  // Handle getting the reference to the content for copy functionality
  const handleContentRef = (ref) => {
    answerRef.current = ref;
  };

  // Handle showing the source popup
  const showSourcePopup = useCallback((sourceIndex, targetElement, statementText) => {
    // Clear any existing timeout to prevent flickering
    if (popupTimeoutRef.current) {
      clearTimeout(popupTimeoutRef.current);
      popupTimeoutRef.current = null;
    }

    if (!finalSources || !finalSources[sourceIndex] || !targetElement) return;

    const rect = targetElement.getBoundingClientRect();
    const scrollY = window.scrollY || window.pageYOffset;
    const scrollX = window.scrollX || window.pageXOffset;

    const newHoverInfo = {
      index: sourceIndex,
      statementText,
      position: {
        top: rect.top + scrollY - 10, // Position above the reference
        left: rect.left + scrollX + rect.width / 2, // Center horizontally
      }
    };
    setHoveredSourceInfo(newHoverInfo);
  }, [finalSources]);

  const hideSourcePopup = useCallback(() => {
    if (popupTimeoutRef.current) {
        clearTimeout(popupTimeoutRef.current); // Clear existing timeout if mouse leaves quickly
    }
    popupTimeoutRef.current = setTimeout(() => {
      setHoveredSourceInfo(null);
      popupTimeoutRef.current = null;
    }, 15); // Delay allows moving mouse onto popup
  }, []);

  // Handle mouse enter on the popup to cancel the hide timeout
  const cancelHidePopup = useCallback(() => {
    // Clear the hide timeout if the mouse enters the popup itself
    if (popupTimeoutRef.current) {
      clearTimeout(popupTimeoutRef.current);
      popupTimeoutRef.current = null;
    }
  }, []);

  // Determine button state and appearance for excerpts icon
  const excerptsLoaded = !!excerptsData;  // True if excerptsData is not null/empty
  const canFetchExcerpts = finalSources && finalSources.length > 0 && 
                           !isError && !excerptsLoaded && !isLoadingExcerpts;
  const buttonDisabled = isLoadingExcerpts || excerptsLoaded; // Disable button if loading or loaded
  const buttonIcon = isLoadingExcerpts
                      ? <FaSpinner className="spin" style={{ fontSize: 20 }} />
                      : excerptsLoaded
                      ? <FaCheckCircle 
                          style={{ 
                            width: 22, 
                            height: 22, 
                            color: 'var(--secondary-color)',
                            filter: 'brightness(0.75)' 
                          }} 
                        />
                      : <img src={excerpts} alt="excerpts icon" />; 
  const buttonClassName = `excerpts-icon ${isLoadingExcerpts ? 'loading' : ''} ${excerptsLoaded ? 'loaded' : ''}`;

  return (
    <>
      { !hasTokens ? (
        // If no tokens, render pre-stream UI.
        (!isError && thoughtLabel) ? (
          <div className="answer-container">
            {/* User Message */}
            <div className="message-row user-message">
              <div className="message-bubble user-bubble">
                <p className="question">{userMessage}</p>
              </div>
              <div className="user-icon">
                <img src={user} alt="user icon" />
              </div>
            </div>
            {/* Bot Message (pre-stream with spinner) */}
            <div className="message-row bot-message pre-stream">
              <div className="bot-container">
                <div className="thinking-info">
                  <Box mt={1} display="flex" alignItems="center">
                    <Box className="custom-spinner" />
                    <Box ml={1}>
                      <span
                        className="thinking-time"
                        onClick={() => openRightSidebar("tasks", blockId)}
                      >
                        {thoughtLabel}
                      </span>
                    </Box>
                  </Box>
                </div>
              </div>
            </div>
          </div>
        ) : (
          // Render without spinner (user message only)
          <div className="answer-container">
            <div className="message-row user-message">
              <div className="message-bubble user-bubble">
                <p className="question">{userMessage}</p>
              </div>
              <div className="user-icon">
                <img src={user} alt="user icon" />
              </div>
            </div>
          </div>
        )
      ) : (
        // Render Full Chat Message
        <div className="answer-container">
          {/* User Message */}
          <div className="message-row user-message">
            <div className="message-bubble user-bubble">
              <p className="question">{userMessage}</p>
            </div>
            <div className="user-icon">
              <img src={user} alt="user icon" />
            </div>
          </div>
          {/* Bot Message */}
          <div className="message-row bot-message">
            <div className="bot-container">
              {!isError && renderThoughtLabel() && (
                <div className="thinking-info">
                  <span
                    className="thinking-time"
                    onClick={() => openRightSidebar("tasks", blockId)}
                  >
                    {renderThoughtLabel()}
                  </span>
                </div>
              )}
              {renderSourcesRead() !== null && (
                <div className="sources-read-container">
                  <p className="sources-read">
                    Sources Read: {renderSourcesRead()}
                  </p>
                </div>
              )}
              <div className="answer-block">
                <div className="bot-icon">
                  <img src={bot} alt="bot icon" />
                </div>
                <div className="message-bubble bot-bubble">
                  <div className="answer">
                    <Streaming
                      content={combinedAnswer}
                      isStreaming={isStreaming}
                      onContentRef={handleContentRef}
                      showSourcePopup={showSourcePopup} 
                      hideSourcePopup={hideSourcePopup} 
                    />
                  </div>
                </div>
                <div className="post-icons">
                  {!isStreaming && (
                    <div className="copy-icon" onClick={handleCopy}>
                      <img src={copy} alt="copy icon" />
                      <span className="tooltip">Copy</span>
                    </div>
                  )}
                  {actions && actions.some(a => a.name === "evaluate") && (
                    <div className="evaluate-icon" onClick={() => openRightSidebar("evaluate", blockId)}>
                      <img src={evaluate} alt="evaluate icon" />
                      <span className="tooltip">Evaluate</span>
                    </div>
                  )}
                  {actions && actions.some(a => a.name === "sources") && (
                    <div className="sources-icon" onClick={() => openRightSidebar("sources", blockId)}>
                      <img src={sourcesIcon} alt="sources icon" />
                      <span className="tooltip">Sources</span>
                    </div>
                  )}
                  {actions && actions.some(a => a.name === "graph") && (
                    <div className="graph-icon" onClick={() => setGraphDialogOpen(true)}>
                      <img src={graphIcon} alt="graph icon" />
                      <span className="tooltip">View Graph</span>
                    </div>
                  )}
                  {/* Show Excerpts Button - Conditionally Rendered */}
                  {finalSources && finalSources.length > 0 && !isError && (
                    <div
                      className={buttonClassName}
                      onClick={() => canFetchExcerpts && onFetchExcerpts(blockId)}
                      style={{
                        cursor: buttonDisabled ? 'default' : 'pointer',
                        opacity: excerptsLoaded ? 0.6 : 1
                      }}
                    >
                      {buttonIcon}
                      <span className="tooltip">
                        {excerptsLoaded ? 'Excerpts Loaded'
                          : isLoadingExcerpts ? 'Loading Excerpts…'
                          : 'Show Excerpts'}
                      </span>
                    </div>
                  )}
                </div>
              </div>
            </div>
          </div>
          {/* Render the GraphDialog when graphDialogOpen is true */}
          {graphDialogOpen && (
            <GraphDialog 
              open={graphDialogOpen} 
              onClose={() => setGraphDialogOpen(false)} 
              payload={graphAction ? graphAction.payload : { query: userMessage }} 
            />
          )}
        </div>
      )}
      {/* Render Source Popup */}
      {hoveredSourceInfo && finalSources && finalSources[hoveredSourceInfo.index] && (
        <SourcePopup
          sourceData={finalSources[hoveredSourceInfo.index]}
          excerptsData={excerptsData}
          position={hoveredSourceInfo.position}
          onMouseEnter={cancelHidePopup} // Keep popup open if mouse enters it
          onMouseLeave={hideSourcePopup}
          statementText={hoveredSourceInfo.statementText}
        />
      )}
      {/* Render error container if there's an error */}
      {isError && (
        <div className="error-block" style={{ marginTop: '1rem' }}>
          <h3>Error</h3>
          <p>{errorMessage}</p>
        </div>
      )}
      <Snackbar
        open={snackbarOpen}
        autoHideDuration={3000}
        onClose={handleSnackbarClose}
        message="Copied To Clipboard"
        anchorOrigin={{ vertical: 'bottom', horizontal: 'left' }}
        TransitionComponent={SlideTransition}
        ContentProps={{ classes: { root: 'custom-snackbar' } }}
        action={
          <IconButton
            size="small"
            aria-label="close"
            color="inherit"
            onClick={handleSnackbarClose}
          >
            <FaTimes />
          </IconButton>
        }
      />
    </>
  );
}

export default ChatWindow;