frontend/src/Components/AiComponents/ChatWindow.js

import React, { useRef, useState, 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 } from 'react-icons/fa';
import GraphDialog from './ChatComponents/Graph';
import Streaming from './ChatComponents/Streaming';
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';

// 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,
  actions,
  tasks,
  openRightSidebar,
  // openLeftSidebar,
  isError,
  errorMessage
}) {
  const answerRef = useRef(null);
  const [graphDialogOpen, setGraphDialogOpen] = useState(false);
  const [snackbarOpen, setSnackbarOpen] = useState(false);

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

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