frontend/src/Components/AiPage.js

import React, { useState, useEffect, useRef, useCallback, useMemo } from 'react';
import { flushSync } from 'react-dom';
import Snackbar from '@mui/material/Snackbar';
import Alert from '@mui/material/Alert';
import { FaCog, FaPaperPlane, FaStop } from 'react-icons/fa';
import IntialSetting from './IntialSetting';
import ChatWindow from './AiComponents/ChatWindow';
import RightSidebar from './AiComponents/ChatComponents/RightSidebar';
import './AiPage.css';

function AiPage() {
  // Sidebar and other states
  const [isRightSidebarOpen, setRightSidebarOpen] = useState(
    localStorage.getItem("rightSidebarState") === "true"
  );
  const [rightSidebarWidth, setRightSidebarWidth] = useState(300);
  const [sidebarContent, setSidebarContent] = useState("default");

  const [searchText, setSearchText] = useState("");
  const textAreaRef = useRef(null);
  const [showSettingsModal, setShowSettingsModal] = useState(false);

  const [showChatWindow, setShowChatWindow] = useState(false);
  const [chatBlocks, setChatBlocks] = useState([]);
  const [selectedChatBlockId, setSelectedChatBlockId] = useState(null);

  const [defaultChatHeight, setDefaultChatHeight] = useState(null);
  const [chatBottomPadding, setChatBottomPadding] = useState("60px");

  // States/refs for streaming
  const [isProcessing, setIsProcessing] = useState(false);
  const [activeBlockId, setActiveBlockId] = useState(null);
  const activeEventSourceRef = useRef(null);

  // Snackbar state
  const [snackbar, setSnackbar] = useState({
    open: false,
    message: "",
    severity: "success",
  });

  // Function to open the snackbar
  const openSnackbar = (message, severity = "success") => {
    setSnackbar({ open: true, message, severity });
  };

  // Function to close the snackbar
  const closeSnackbar = (event, reason) => {
    if (reason === 'clickaway') return;
    setSnackbar(prev => ({ ...prev, open: false }));
  };

  useEffect(() => {
    localStorage.setItem("rightSidebarState", isRightSidebarOpen);
  }, [isRightSidebarOpen]);

  useEffect(() => {
    document.documentElement.style.setProperty('--right-sidebar-width', rightSidebarWidth + 'px');
  }, [rightSidebarWidth]);

  // Dynamically increase height of chat input field based on newlines entered
  useEffect(() => {
    if (textAreaRef.current) {
      if (!defaultChatHeight) {
        setDefaultChatHeight(textAreaRef.current.scrollHeight);
      }
      textAreaRef.current.style.height = "auto";
      textAreaRef.current.style.overflowY = "hidden";

      const newHeight = textAreaRef.current.scrollHeight;
      let finalHeight = newHeight;
      if (newHeight > 200) {
        finalHeight = 200;
        textAreaRef.current.style.overflowY = "auto";
      }
      textAreaRef.current.style.height = `${finalHeight}px`;

      const minPaddingPx = 0;
      const maxPaddingPx = 59;
      let newPaddingPx = minPaddingPx;
      if (defaultChatHeight && finalHeight > defaultChatHeight) {
        newPaddingPx =
          minPaddingPx +
          ((finalHeight - defaultChatHeight) / (200 - defaultChatHeight)) *
            (maxPaddingPx - minPaddingPx);
        if (newPaddingPx > maxPaddingPx) newPaddingPx = maxPaddingPx;
      }
      setChatBottomPadding(`${newPaddingPx}px`);
    }
  }, [searchText, defaultChatHeight]);

  const handleOpenRightSidebar = (content, chatBlockId = null) => {
    flushSync(() => {
      if (chatBlockId) {
        setSelectedChatBlockId(chatBlockId);
      }
      setSidebarContent(content ? content : "default");
      setRightSidebarOpen(true);
    });
  };

  const handleEvaluationError = useCallback((blockId, errorMsg) => {
    setChatBlocks(prev =>
      prev.map(block =>
        block.id === blockId
          ? { ...block, isError: true, errorMessage: errorMsg }
          : block
      )
    );
  }, []);

  // Initiate the SSE
  const initiateSSE = (query, blockId) => {
    const startTime = Date.now();
    const sseUrl = `/message-sse?user_message=${encodeURIComponent(query)}`;
    const eventSource = new EventSource(sseUrl);
    activeEventSourceRef.current = eventSource;

    eventSource.addEventListener("token", (e) => {
      const { chunk, index } = JSON.parse(e.data);
      console.log("[SSE token chunk]", JSON.stringify(chunk));
      console.log("[SSE token index]", JSON.stringify(index));
    
      setChatBlocks(prevBlocks => {
        return prevBlocks.map(block => {
          if (block.id === blockId) {
            const newTokenArray = block.tokenChunks ? [...block.tokenChunks] : [];
            newTokenArray[index] = chunk;
    
            return {
              ...block,
              tokenChunks: newTokenArray
            };
          }
          return block;
        });
      });
    });    

    eventSource.addEventListener("final_message", (e) => {
      const endTime = Date.now();
      const thinkingTime = ((endTime - startTime) / 1000).toFixed(1);
      // Only update thinkingTime so the streaming flag turns false and the cursor disappears
      setChatBlocks(prev => prev.map(block =>
        block.id === blockId
          ? { ...block, thinkingTime }
          : block
      ));
    });    
  
    // Listen for the "complete" event to know when to close the connection.
    eventSource.addEventListener("complete", (e) => {
      console.log("Complete event received:", e.data);
      eventSource.close();
      activeEventSourceRef.current = null;
      setIsProcessing(false);
      setActiveBlockId(null);
    });

    // Update actions for only this chat block.
    eventSource.addEventListener("action", (e) => {
      try {
        const actionData = JSON.parse(e.data);
        console.log("Action event received:", actionData);
        setChatBlocks(prev => prev.map(block => {
          if (block.id === blockId) {
            let updatedBlock = { ...block, actions: [...(block.actions || []), actionData] };
            if (actionData.name === "sources") {
              updatedBlock.sources = actionData.payload;
            }
            if (actionData.name === "graph") {
              updatedBlock.graph = actionData.payload;
            }
            return updatedBlock;
          }
          return block;
        }));
      } catch (err) {
        console.error("Error parsing action event:", err);
      }
    });
    
    // Update the error for this chat block.
    eventSource.addEventListener("error", (e) => {
      console.error("Error from SSE:", e.data);
      setChatBlocks(prev => prev.map(block =>
        block.id === blockId
          ? { 
              ...block, 
              isError: true, 
              errorMessage: e.data, 
              aiAnswer: "",
              tasks: [] 
            }
          : block
      ));
      eventSource.close();
      activeEventSourceRef.current = null;
      setIsProcessing(false);
      setActiveBlockId(null);
    });

    eventSource.addEventListener("step", (e) => {
      console.log("Step event received:", e.data);
      setChatBlocks(prev => prev.map(block =>
        block.id === blockId
          ? { ...block, thoughtLabel: e.data }
          : block
      ));
    });

    eventSource.addEventListener("sources_read", (e) => {
      console.log("Sources read event received:", e.data);
      try {
        const parsed = JSON.parse(e.data);
        let count;
        if (typeof parsed === 'number') {
          count = parsed;
        } else if (parsed && typeof parsed.count === 'number') {
          count = parsed.count;
        }
        if (typeof count === 'number') {
          setChatBlocks(prev => prev.map(block =>
            block.id === blockId
              ? { ...block, sourcesRead: count, sources: parsed.sources || [] }
              : block
          ));
        }
      } catch(err) {
        if (e.data.trim() !== "") {
          setChatBlocks(prev => prev.map(block =>
            block.id === blockId
              ? { ...block, sourcesRead: e.data }
              : block
          ));
        }
      }
    });

    eventSource.addEventListener("task", (e) => {
      console.log("Task event received:", e.data);
      try {
        const taskData = JSON.parse(e.data);
        setChatBlocks(prev => prev.map(block => {
          if (block.id === blockId) {
            const existingTaskIndex = (block.tasks || []).findIndex(t => t.task === taskData.task);
            if (existingTaskIndex !== -1) {
              const updatedTasks = [...block.tasks];
              updatedTasks[existingTaskIndex] = { ...updatedTasks[existingTaskIndex], status: taskData.status };
              return { ...block, tasks: updatedTasks };
            } else {
              return { ...block, tasks: [...(block.tasks || []), taskData] };
            }
          }
          return block;
        }));
      } catch (error) {
        console.error("Error parsing task event:", error);
      }
    });
  };

  // Create a new chat block and initiate the SSE
  const handleSend = () => {
    if (!searchText.trim()) return;
    const blockId = new Date().getTime();
    setActiveBlockId(blockId);
    setIsProcessing(true);
    setChatBlocks(prev => [
      ...prev,
      { 
        id: blockId, 
        userMessage: searchText, 
        tokenChunks: [],
        aiAnswer: "", 
        thinkingTime: null, 
        thoughtLabel: "", 
        sourcesRead: "",
        tasks: [], 
        sources: [],
        actions: []
      }
    ]);
    setShowChatWindow(true);
    const query = searchText;
    setSearchText("");
    initiateSSE(query, blockId);
  };

  const handleKeyDown = (e) => {
    if (e.key === "Enter" && !e.shiftKey) {
      e.preventDefault();
      if (!isProcessing) {
        handleSend();
      }
    }
  };

  // Stop the user request and close the active SSE connection
  const handleStop = async () => {
    // Close the active SSE connection if it exists
    if (activeEventSourceRef.current) {
      activeEventSourceRef.current.close();
      activeEventSourceRef.current = null;
    }
    // Send POST request to /stop and update the chat block with the returned message
    try {
      const response = await fetch('/stop', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({})
      });
      const data = await response.json();

      if (activeBlockId) {
        setChatBlocks(prev => prev.map(block =>
          block.id === activeBlockId
            ? { ...block, aiAnswer: data.message, thinkingTime: 0, tasks: [] }
            : block
        ));
      }
    } catch (error) {
      console.error("Error stopping the request:", error);
      if (activeBlockId) {
        setChatBlocks(prev => prev.map(block =>
          block.id === activeBlockId
            ? { ...block, aiAnswer: "Error stopping task", thinkingTime: 0, tasks: [] }
            : block
        ));
      }
    }
    setIsProcessing(false);
    setActiveBlockId(null);
  };

  const handleSendButtonClick = () => {
    if (searchText.trim()) handleSend();
  };

  // Get the chat block whose details should be shown in the sidebar.
  const selectedBlock = chatBlocks.find(block => block.id === selectedChatBlockId);
  const evaluateAction = selectedBlock && selectedBlock.actions
    ? selectedBlock.actions.find(a => a.name === "evaluate")
    : null;

  // Memoized evaluation object
  const evaluation = useMemo(() => {
    if (!evaluateAction) return null;
    return {
      ...evaluateAction.payload,
      blockId: selectedBlock?.id,
      onError: handleEvaluationError,
    };
  }, [evaluateAction, selectedBlock?.id, handleEvaluationError]);

  return (
    <div 
      className="app-container"
      style={{
        paddingRight: isRightSidebarOpen
          ? Math.max(0, rightSidebarWidth - 250) + 'px'
          : 0,
      }}
    >
      {showChatWindow && selectedBlock && (sidebarContent !== "default" || (selectedBlock.tasks && selectedBlock.tasks.length > 0) || (selectedBlock.sources && selectedBlock.sources.length > 0)) && (
        <div className="floating-sidebar">
          <RightSidebar 
            isOpen={isRightSidebarOpen}
            rightSidebarWidth={rightSidebarWidth}
            setRightSidebarWidth={setRightSidebarWidth}
            toggleRightSidebar={() => setRightSidebarOpen(!isRightSidebarOpen)}
            sidebarContent={sidebarContent}
            tasks={selectedBlock.tasks || []}
            tasksLoading={false}
            sources={selectedBlock.sources || []}
            sourcesLoading={false}
            onSourceClick={(source) => {
              if (!source || !source.link) return;
              window.open(source.link, '_blank');
            }}
            evaluation={evaluation}            
          />
        </div>
      )}
      
      <main className="main-content">
        {showChatWindow ? (
          <>
            <div className="chat-container">
              {chatBlocks.map((block) => (
                <ChatWindow
                  key={block.id}
                  blockId={block.id}
                  userMessage={block.userMessage}
                  tokenChunks={block.tokenChunks}
                  aiAnswer={block.aiAnswer}
                  thinkingTime={block.thinkingTime}
                  thoughtLabel={block.thoughtLabel}
                  sourcesRead={block.sourcesRead}
                  actions={block.actions}        
                  tasks={block.tasks}            
                  openRightSidebar={handleOpenRightSidebar}
                  openLeftSidebar={() => { /* if needed */ }}
                  isError={block.isError}
                  errorMessage={block.errorMessage}
                />
              ))}
            </div>
            <div 
              className="floating-chat-search-bar"
              style={{
                transform: isRightSidebarOpen
                  ? `translateX(calc(-50% - ${Math.max(0, (rightSidebarWidth - 250) / 2)}px))`
                  : 'translateX(-50%)'               
              }}
            >
              <div className="chat-search-input-wrapper" style={{ paddingBottom: chatBottomPadding }}>
                <textarea
                  rows="1"
                  className="chat-search-input"
                  placeholder="Message..."
                  value={searchText}
                  onChange={(e) => setSearchText(e.target.value)}
                  onKeyDown={handleKeyDown}
                  ref={textAreaRef}
                />
              </div>
              <div className="chat-icon-container">
                <button
                  className="chat-settings-btn"
                  onClick={() => setShowSettingsModal(true)}
                >
                  <FaCog />
                </button>
                {/* Conditionally render Stop or Send button */}
                <button 
                  className={`chat-send-btn ${isProcessing ? 'stop-btn' : ''}`}
                  onClick={isProcessing ? handleStop : handleSendButtonClick}
                >
                  {isProcessing ? <FaStop size={12} color="black" /> : <FaPaperPlane />}
                </button>
              </div>
            </div>
          </>
        ) : (
          <div className="search-area">
            <h1>How can I help you today?</h1>
            <div className="search-bar">
              <div className="search-input-wrapper">
                <textarea
                  rows="1"
                  className="search-input"
                  placeholder="Message..."
                  value={searchText}
                  onChange={(e) => setSearchText(e.target.value)}
                  onKeyDown={handleKeyDown}
                  ref={textAreaRef}
                />
              </div>
              <div className="icon-container">
                <button
                  className="settings-btn"
                  onClick={() => setShowSettingsModal(true)}
                >
                  <FaCog />
                </button>
                <button 
                  className={`send-btn ${isProcessing ? 'stop-btn' : ''}`}
                  onClick={isProcessing ? handleStop : handleSendButtonClick}
                >
                  {isProcessing ? <FaStop /> : <FaPaperPlane />}
                </button>
              </div>
            </div>
          </div>
        )}
      </main>

      {showSettingsModal && (
        <IntialSetting
          trigger={true}
          setTrigger={() => setShowSettingsModal(false)}
          fromAiPage={true}
          openSnackbar={openSnackbar}
          closeSnackbar={closeSnackbar}
        />
      )}
      <Snackbar
        open={snackbar.open}
        autoHideDuration={snackbar.severity === 'success' ? 3000 : null}
        onClose={closeSnackbar}
        anchorOrigin={{ vertical: 'top', horizontal: 'center' }}
      >
        <Alert onClose={closeSnackbar} severity={snackbar.severity} variant="filled" sx={{ width: '100%' }}>
          {snackbar.message}
        </Alert>
      </Snackbar>
    </div>
  );
}

export default AiPage;