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, FaPlus, FaGoogle, FaMicrosoft, FaSlack } from 'react-icons/fa';
import IntialSetting from './IntialSetting';
import AddContentDropdown from './AiComponents/Dropdowns/AddContentDropdown';
import AddFilesDialog from './AiComponents/Dropdowns/AddFilesDialog';
import ChatWindow from './AiComponents/ChatWindow';
import RightSidebar from './AiComponents/Sidebars/RightSidebar';
import Notification from '../Components/AiComponents/Notifications/Notification';
import { useNotification } from '../Components/AiComponents/Notifications/useNotification';
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 addBtnRef = useRef(null);
  const chatAddBtnRef = useRef(null);
  const [isAddContentOpen, setAddContentOpen] = useState(false);
  const [isTooltipSuppressed, setIsTooltipSuppressed] = useState(false);

  const [isAddFilesDialogOpen, setIsAddFilesDialogOpen] = useState(false);

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

  const [sessionContent, setSessionContent] = useState({ files: [], links: [] });

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

  // State to track if we should auto-scroll to the bottom
  const [autoScrollEnabled, setAutoScrollEnabled] = useState(false);

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

  // State for tracking selected services
  const [selectedServices, setSelectedServices] = useState({
    google: [],
    microsoft: [],
    slack: false
  });

  // State for Slack authentication modal
  const [showSlackAuthModal, setShowSlackAuthModal] = useState(false);
  const [pendingSlackAuth, setPendingSlackAuth] = useState(null);

  // Notifications
  const { 
    notifications, 
    addNotification, 
    removeNotification,
    updateNotification 
  } = useNotification();

  // Token management
  const tokenExpiryTimersRef = useRef({});
  const notificationIdsRef = useRef({});

  // Function to check if we are near the bottom of the page
  const checkIfNearBottom = (threshold = 400) => {
    const scrollTop = window.pageYOffset || document.documentElement.scrollTop;
    const scrollHeight = document.documentElement.scrollHeight;
    const clientHeight = window.innerHeight;
    const distanceFromBottom = scrollHeight - (scrollTop + clientHeight);
    
    return distanceFromBottom <= threshold;
  };

  // Helper to scroll to bottom
  const scrollToBottom = (smooth = true) => {
    window.scrollTo({
      top: document.documentElement.scrollHeight,
      behavior: smooth ? 'smooth' : 'auto'
    });
  };

  // Function to open the snackbar
  const openSnackbar = useCallback((message, severity = "success", duration) => {
    let finalDuration;

    if (duration !== undefined) {
      // If a specific duration is provided (e.g., 5000 or null), use it.
      finalDuration = duration;
    } else {
      // Otherwise, use the default logic.
      finalDuration = severity === 'success' ? 3000 : null; // Success auto-hides, others are persistent by default.
    }

    setSnackbar({ open: true, message, severity, duration: finalDuration });
  }, []);

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

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

  // Add cleanup handler for when the user closes the tab/browser
  useEffect(() => {
    const handleCleanup = () => {
      navigator.sendBeacon('/cleanup');
    };
    window.addEventListener('beforeunload', handleCleanup);
    return () => window.removeEventListener('beforeunload', handleCleanup);
  }, []);

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

  // Update backend whenever selected services change
  useEffect(() => {
    const updateSelectedServices = async () => {
      try {
        await fetch('/api/selected-services', {
          method: 'POST',
          headers: { 'Content-Type': 'application/json' },
          body: JSON.stringify({
            services: selectedServices
          })
        });
      } catch (error) {
        console.error('Failed to update selected services:', error);
      }
    };
    
    updateSelectedServices();
  }, [selectedServices]);

  // Clear all tokens on page load
  useEffect(() => {
    // Clear all provider tokens on new tab/page load
    ['google', 'microsoft', 'slack'].forEach(provider => {
      sessionStorage.removeItem(`${provider}_token`);
      sessionStorage.removeItem(`${provider}_token_expiry`);
    });
    
    // Clear any existing timers
    Object.values(tokenExpiryTimersRef.current).forEach(timer => clearTimeout(timer));
    tokenExpiryTimersRef.current = {};
    
    console.log('Cleared all tokens for new session');
  }, []);

  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
      )
    );
  }, []);

  // Function to store token with expiry
  const storeTokenWithExpiry = (provider, token) => {
    const expiryTime = Date.now() + (60 * 60 * 1000); // 1 hour from now
    sessionStorage.setItem(`${provider}_token`, token);
    sessionStorage.setItem(`${provider}_token_expiry`, expiryTime.toString());
    
    // Set up expiry timer
    setupTokenExpiryTimer(provider, expiryTime);
  };

  // Function to check if token is valid
  const isTokenValid = (provider) => {
    const token = sessionStorage.getItem(`${provider}_token`);
    const expiry = sessionStorage.getItem(`${provider}_token_expiry`);
    
    if (!token || !expiry) return false;
    
    const expiryTime = parseInt(expiry);
    return Date.now() < expiryTime;
  };

  // Function to get provider icon
  const getProviderIcon = useCallback((provider) => {
    switch (provider.toLowerCase()) {
      case 'google':
        return <FaGoogle />;
      case 'microsoft':
        return <FaMicrosoft />;
      case 'slack':
        return <FaSlack />;
      default:
        return null;
    }
  }, []);

  // Function to get provider color
  const getProviderColor = useCallback((provider) => {
    switch (provider.toLowerCase()) {
      case 'google':
        return '#4285F4';
      case 'microsoft':
        return '#00A4EF';
      case 'slack':
        return '#4A154B';
      default:
        return '#666';
    }
  }, []);

  // Function to set up timer for token expiry notification
  const setupTokenExpiryTimer = useCallback((provider, expiryTime) => {
    // Clear existing timer if any
    if (tokenExpiryTimersRef.current[provider]) {
      clearTimeout(tokenExpiryTimersRef.current[provider]);
    }
    
    // Remove any existing notification for this provider
    if (notificationIdsRef.current[provider]) {
      removeNotification(notificationIdsRef.current[provider]);
      delete notificationIdsRef.current[provider];
    }
    
    const timeUntilExpiry = expiryTime - Date.now();
    
    if (timeUntilExpiry > 0) {
      tokenExpiryTimersRef.current[provider] = setTimeout(() => {
        const providerName = provider.charAt(0).toUpperCase() + provider.slice(1);
        const providerColor = getProviderColor(provider);
        
        // Add notification
        const notificationId = addNotification({
          type: 'warning',
          title: `${providerName} Authentication Expired`,
          message: `Your ${providerName} authentication has expired. Please reconnect to continue using ${providerName} services.`,
          icon: getProviderIcon(provider),
          dismissible: true,
          autoDismiss: false,
          actions: [
            {
              id: 'reconnect',
              label: `Reconnect ${providerName}`,
              style: {
                background: providerColor,
                color: 'white',
                border: 'none'
              },
              data: { provider }
            }
          ],
          style: {
            borderLeftColor: providerColor
          }
        });
        
        // Store notification ID
        notificationIdsRef.current[provider] = notificationId;
        
        // Clear token data
        sessionStorage.removeItem(`${provider}_token`);
        sessionStorage.removeItem(`${provider}_token_expiry`);
        
        // Update selected services to reflect disconnection
        if (provider === 'slack') {
          setSelectedServices(prev => ({ ...prev, slack: false }));
        } else {
          setSelectedServices(prev => ({ ...prev, [provider]: [] }));
        }
        
      }, timeUntilExpiry);
    }
  }, [addNotification, getProviderColor, getProviderIcon, removeNotification, setSelectedServices]);

  // Check existing tokens on component mount and set up timers
  useEffect(() => {
    ['google', 'microsoft', 'slack'].forEach(provider => {
      const expiry = sessionStorage.getItem(`${provider}_token_expiry`);
      if (expiry) {
        const expiryTime = parseInt(expiry);
        if (Date.now() < expiryTime) {
          setupTokenExpiryTimer(provider, expiryTime);
        } else {
          // Token already expired, clear it
          sessionStorage.removeItem(`${provider}_token`);
          sessionStorage.removeItem(`${provider}_token_expiry`);
        }
      }
    });
    
    // Cleanup timers on unmount
    return () => {
      Object.values(tokenExpiryTimersRef.current).forEach(timer => clearTimeout(timer));
    };
  }, [setupTokenExpiryTimer]);

  // 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) => {
      console.log("[SSE final message]", e.data);
      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 "final_sources" event to update sources in AI answer of this chat block.
    eventSource.addEventListener("final_sources", (e) => {
      try {
        const sources = JSON.parse(e.data);
        console.log("Final sources received:", sources);
        setChatBlocks(prev => prev.map(block =>
          block.id === blockId ? { ...block, finalSources: sources } : block
        ));
      } catch (err) {
        console.error("Error parsing final_sources event:", err);
      }
    });
  
    // 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;

    // Check if this is the first prompt (no existing chat blocks)
    const isFirstPrompt = chatBlocks.length === 0;

    // Only check if user is near bottom if this is NOT the first prompt before adding a new block
    const shouldScroll = !isFirstPrompt && checkIfNearBottom(1000); // 1000px threshold
    setAutoScrollEnabled(shouldScroll);

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

  // Auto-scroll when chat block is added
  useEffect(() => {
    if (autoScrollEnabled && isProcessing) {
      scrollToBottom();
    }
  }, [isProcessing, autoScrollEnabled]);

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

  // Toggle the Add Content dropdown
  const handleToggleAddContent = (event) => {
    event.stopPropagation(); // Prevents the click from closing the menu immediately
    // If we are about to close the dropdown, suppress the tooltip.
    if (isAddContentOpen) {
      setIsTooltipSuppressed(true);
    }
     setAddContentOpen(prev => !prev);
  };

  // Handle mouse enter on the Add Content button to suppress tooltip
  const handleMouseLeaveAddBtn = () => {
    setIsTooltipSuppressed(false);
  };

  // Close the Add Content dropdown
  const closeAddContentDropdown = () => {
    setAddContentOpen(false);
  };

  // Open the Add Files dialog
  const handleOpenAddFilesDialog = () => {
    setAddContentOpen(false); // Close the dropdown when opening the dialog
    setIsAddFilesDialogOpen(true);
  };

  // Fetch excerpts for a specific block
  const handleFetchExcerpts = useCallback(async (blockId) => {
    let blockIndex = -1;
    let currentBlock = null;
  
    // Find the block to check its current state
    setChatBlocks(prev => {
        blockIndex = prev.findIndex(b => b.id === blockId);
        if (blockIndex !== -1) {
            currentBlock = prev[blockIndex];
        }
        // No state change here, just reading the state
        return prev;
    });
  
    // Prevent fetching if already loaded or currently loading
    if (blockIndex === -1 || !currentBlock || currentBlock.excerptsData || currentBlock.isLoadingExcerpts) return;
  
    // Set loading state for the specific block
    setChatBlocks(prev => prev.map(b =>
      b.id === blockId ? { ...b, isLoadingExcerpts: true } : b
    ));
  
    try {
      // Call the backend endpoint to get excerpts
      const response = await fetch('/action/excerpts', {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ blockId: blockId })
      });
  
      if (!response.ok) {
        const errorData = await response.json();
        throw new Error(errorData.detail || `HTTP error! status: ${response.status}`);
      }
  
      const data = await response.json();
      console.log("Fetched excerpts data from backend:", data.result);
  
      // Update the specific block with the fetched excerptsData
      setChatBlocks(prev => prev.map(b =>
        b.id === blockId
          ? {
              ...b,
              excerptsData: data.result, // Store the fetched data
              isLoadingExcerpts: false,    // Turn off loading
            }
          : b
      ));
      openSnackbar("Excerpts loaded successfully!", "success");
  
    } catch (error) {
      console.error("Error requesting excerpts:", error);
      // Reset loading state on error
      setChatBlocks(prev => prev.map(b =>
        b.id === blockId ? { ...b, isLoadingExcerpts: false } : b
      ));
      openSnackbar(`Failed to load excerpts`, "error");
    }
  }, [openSnackbar]);

  // Function to handle notification actions
  const handleNotificationAction = (notificationId, actionId, actionData) => {
    console.log('Notification action triggered:', { notificationId, actionId, actionData });
    
    // Handle both 'reconnect' and 'connect' actions
    if ((actionId === 'reconnect' || actionId === 'connect') && actionData?.provider) {
      // Remove the notification
      removeNotification(notificationId);
      
      // Clean up stored notification ID if it exists
      if (notificationIdsRef.current[actionData.provider] === notificationId) {
        delete notificationIdsRef.current[actionData.provider];
      }
      
      // Trigger authentication
      initiateOAuth(actionData.provider);
    }
  };

  // Slack Auth Modal Component
  const SlackAuthModal = () => {
    if (!showSlackAuthModal) return null;

    return (
      <div className="slack-auth-modal" onClick={() => setShowSlackAuthModal(false)}>
        <div className="slack-auth-modal-content" onClick={(e) => e.stopPropagation()}>
          <div className="slack-auth-modal-header">
            <div className="slack-auth-modal-icon">
              <FaSlack />
            </div>
            <h2 className="slack-auth-modal-title">Connect to Slack</h2>
          </div>
          
          <div className="slack-auth-modal-body">
            <p>
              Slack will use your currently logged-in workspace. If you're logged into multiple workspaces, 
              you'll be able to select which one to use.
            </p>
            
            <div className="slack-auth-modal-tips">
              <div className="slack-auth-modal-tips-title">
                To use a different workspace:
              </div>
              <ul>
                <li>Sign out of Slack in your browser first</li>
                <li>Use an incognito/private browser window</li>
                <li>Or switch workspaces in Slack before continuing</li>
              </ul>
            </div>
            
            <p style={{ marginTop: '1rem', fontSize: '0.9rem', color: '#888' }}>
              You'll be redirected to Slack to authorize access. This allows the app to read messages 
              and channels from your selected workspace.
            </p>
          </div>
          
          <div className="slack-auth-modal-buttons">
            <button 
              className="slack-auth-modal-button slack-auth-modal-button-secondary"
              onClick={() => {
                setShowSlackAuthModal(false);
                if (pendingSlackAuth?.notificationId) {
                  removeNotification(pendingSlackAuth.notificationId);
                }
                setPendingSlackAuth(null);
              }}
            >
              Cancel
            </button>
            <button 
              className="slack-auth-modal-button slack-auth-modal-button-primary"
              onClick={() => {
                setShowSlackAuthModal(false);
                if (pendingSlackAuth) {
                  // Continue with the actual OAuth flow
                  proceedWithSlackAuth(pendingSlackAuth);
                }
              }}
            >
              Continue to Slack
            </button>
          </div>
        </div>
      </div>
    );
  };

  // Function to proceed with Slack authentication
  const proceedWithSlackAuth = (pendingAuth) => {
    const { provider, notificationId } = pendingAuth;
    setPendingSlackAuth(null);
    
    // Update the notification
    if (notificationId) {
      updateNotification(notificationId, {
        message: 'Redirecting to Slack for authentication...'
      });
    }
    
    // Now proceed with OAuth
    proceedWithOAuth(provider, notificationId);
  };

  // Function to initiate OAuth
  const initiateOAuth = (provider) => {
    // Special handling for Slack - show modal first
    if (provider === 'slack') {
      // Show the modal first
      const connectingNotificationId = addNotification({
        type: 'info',
        title: `Preparing to connect to Slack`,
        message: 'Please review the connection information...',
        icon: getProviderIcon(provider),
        dismissible: false,
        autoDismiss: false
      });
      
      setPendingSlackAuth({ provider, notificationId: connectingNotificationId });
      setShowSlackAuthModal(true);
      return;
    }
    
    // For Google and Microsoft, proceed normally
    proceedWithOAuth(provider);
  };

  // Function to proceed with OAuth (renamed from initiateOAuth)
  const proceedWithOAuth = (provider, existingNotificationId = null) => {
    const authUrls = {
      google: `https://accounts.google.com/o/oauth2/v2/auth?` +
        `client_id=${process.env.REACT_APP_GOOGLE_CLIENT_ID}&` +
        `response_type=token&` +
        `scope=email profile https://www.googleapis.com/auth/drive.readonly https://www.googleapis.com/auth/gmail.readonly https://www.googleapis.com/auth/calendar.readonly https://www.googleapis.com/auth/tasks.readonly&` +
        `redirect_uri=${window.location.origin}/auth-receiver.html&` +
        `state=google&` +
        `prompt=select_account`,
      
      microsoft: `https://login.microsoftonline.com/common/oauth2/v2.0/authorize?` +
        `client_id=${process.env.REACT_APP_MICROSOFT_CLIENT_ID}&` +
        `response_type=token&` +
        `scope=openid profile email Files.Read.All Mail.Read Calendars.Read Tasks.Read Notes.Read&` +
        `redirect_uri=${window.location.origin}/auth-receiver.html&` +
        `response_mode=fragment&` +
        `state=microsoft&` +
        `prompt=select_account`,
      
      slack: `https://slack.com/oauth/v2/authorize?` +
        `client_id=${process.env.REACT_APP_SLACK_CLIENT_ID}&` +
        `user_scope=channels:read,channels:history,files:read,groups:read,im:read,mpim:read,search:read,users:read&` +
        `redirect_uri=${window.location.origin}/auth-receiver.html&` +
        `state=slack`
    };
    
    const authWindow = window.open(
      authUrls[provider],
      'Connect Account',
      'width=600,height=700,left=200,top=100'
    );
    
    // Show connecting notification
    const connectingNotificationId = existingNotificationId || addNotification({
      type: 'info',
      title: `Connecting to ${provider.charAt(0).toUpperCase() + provider.slice(1)}`,
      message: 'Please complete the authentication in the popup window...',
      icon: getProviderIcon(provider),
      dismissible: false,
      autoDismiss: false
    });
    
    // Set up message listener
    const messageHandler = async (event) => {
      if (event.origin !== window.location.origin) return;
      
      if (event.data.type === 'auth-success') {
        const { token, code, team_id, team_name, provider: authProvider } = event.data;
        
        // Remove connecting notification
        removeNotification(connectingNotificationId);
        
        // Store token with expiry
        storeTokenWithExpiry(provider, token);
        
        // Store workspace info in sessionStorage for Slack
        if (provider === 'slack' && team_id) {
          sessionStorage.setItem('slack_workspace', JSON.stringify({
            team_id,
            team_name: team_name || 'Unknown Workspace'
          }));
        }
        
        // Build payload for backend
        const payload = {
          provider: authProvider || provider,  // Use provider from auth data if available
          token,
          origin: window.location.origin  // Add origin for backend
        };
        
        // Include code for Slack and Microsoft (if using code flow)
        if (code) {
          payload.code = code;
        }
        
        // Include workspace info for Slack
        if (provider === 'slack' || authProvider === 'slack') {
          if (team_id) payload.team_id = team_id;
          if (team_name) payload.team_name = team_name;
        }
        
        // Send token and workspace info to backend
        try {
          const response = await fetch('/api/session-token', {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify(payload)
          });
          
          if (response.ok) {
            // Show success notification with workspace info
            const workspaceInfo = provider === 'slack' && team_name 
              ? ` (Workspace: ${team_name})` 
              : '';
              
            addNotification({
              type: 'success',
              title: 'Connected Successfully',
              message: `Successfully connected to ${provider.charAt(0).toUpperCase() + provider.slice(1)}${workspaceInfo}!`,
              icon: getProviderIcon(provider),
              autoDismiss: true,
              duration: 3000,
              showProgress: true
            });
          }
        } catch (error) {
          console.error(`Failed to connect to ${provider}:`, error);
          addNotification({
            type: 'error',
            title: 'Connection Failed',
            message: `Failed to connect to ${provider.charAt(0).toUpperCase() + provider.slice(1)}. Please try again.`,
            autoDismiss: true,
            duration: 5000
          });
        }
        
        window.removeEventListener('message', messageHandler);
      } else if (event.data.type === 'auth-failed') {
        console.error(`Authentication failed for ${provider}:`, event.data.error);
        
        // Remove connecting notification
        removeNotification(connectingNotificationId);
        
        // Show error notification
        addNotification({
          type: 'error',
          title: 'Authentication Failed',
          message: `Failed to authenticate with ${provider.charAt(0).toUpperCase() + provider.slice(1)}. ${event.data.error_description || 'Please try again.'}`,
          autoDismiss: true,
          duration: 5000
        });
        
        window.removeEventListener('message', messageHandler);
      }
    };
    
    window.addEventListener('message', messageHandler);
    
    // Handle if user closes the popup without authenticating
    const checkInterval = setInterval(() => {
      if (authWindow.closed) {
        clearInterval(checkInterval);
        removeNotification(connectingNotificationId);
        window.removeEventListener('message', messageHandler);
      }
    }, 1000);
  };

  // Handle service selection from dropdown
  const handleServiceClick = useCallback((provider, service) => {
    // Toggle selection
    if (provider === 'slack') {
      setSelectedServices(prev => ({ ...prev, slack: !prev.slack }));
    } else {
      setSelectedServices(prev => ({
        ...prev,
        [provider]: prev[provider].includes(service) 
          ? prev[provider].filter(s => s !== service)
          : [...prev[provider], service]
      }));
    }
    
    // Check if token is valid
    if (!isTokenValid(provider)) {
      // Show notification prompting to authenticate
      const notificationId = addNotification({
        type: 'info',
        title: 'Authentication Required',
        message: `Please connect your ${provider.charAt(0).toUpperCase() + provider.slice(1)} account to use this service.`,
        icon: getProviderIcon(provider),
        actions: [
          {
            id: 'connect',
            label: `Connect ${provider.charAt(0).toUpperCase() + provider.slice(1)}`,
            style: {
              background: getProviderColor(provider),
              color: 'white',
              border: 'none'
            },
            data: { provider }
          }
        ],
        autoDismiss: true,
        duration: 5000,
        showProgress: true
      });
    }
  }, [addNotification, getProviderIcon, getProviderColor]);

  // 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,
      }}
    >
      <Notification
        notifications={notifications}
        position="top-right"
        animation="slide"
        stackDirection="down"
        maxNotifications={5}
        spacing={12}
        offset={{ x: 20, y: 20 }}
        onDismiss={removeNotification}
        onAction={handleNotificationAction}
        theme="light"
      />
      {showSlackAuthModal && <SlackAuthModal />}
      {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 ${isProcessing ? 'processing' : ''}`}>
              {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}
                  finalSources={block.finalSources}
                  excerptsData={block.excerptsData}
                  isLoadingExcerpts={block.isLoadingExcerpts}
                  onFetchExcerpts={handleFetchExcerpts}
                  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">
                <div className="chat-left-icons">
                  <div className="tooltip-wrapper">
                    <button
                      className="chat-settings-btn"
                      onClick={() => setShowSettingsModal(true)}
                    >
                      <FaCog />
                    </button>
                    <span className="tooltip">Settings</span>
                  </div>
                  <div 
                    className="tooltip-wrapper" 
                    onMouseLeave={handleMouseLeaveAddBtn}
                  >
                    <button className="chat-add-btn" onClick={handleToggleAddContent} ref={chatAddBtnRef}>
                      <FaPlus />
                    </button>
                    <span className={`tooltip ${isAddContentOpen || isTooltipSuppressed ? 'hidden' : ''}`}>Add Content</span>
                    <AddContentDropdown
                      isOpen={isAddContentOpen}
                      onClose={closeAddContentDropdown}
                      toggleButtonRef={chatAddBtnRef}
                      onAddFilesClick={handleOpenAddFilesDialog}
                      onServiceClick={handleServiceClick}
                      selectedServices={selectedServices}
                    />
                  </div>
                </div>
                {/* Conditionally render Stop or Send button */}
                <div className="tooltip-wrapper">
                  <button
                    className={`chat-send-btn ${isProcessing ? 'stop-btn' : ''}`}
                    onClick={isProcessing ? handleStop : handleSendButtonClick}
                  >
                    {isProcessing ? <FaStop size={12} color="black" /> : <FaPaperPlane />}
                  </button>
                  <span className="tooltip">{isProcessing ? 'Stop' : 'Send'}</span>
                </div>
              </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">
                <div className="left-icons">
                  <div className="tooltip-wrapper">
                    <button
                      className="settings-btn"
                      onClick={() => setShowSettingsModal(true)}
                    >
                      <FaCog />
                    </button>
                    <span className="tooltip">Settings</span>
                  </div>
                  <div 
                    className="tooltip-wrapper" 
                    onMouseLeave={handleMouseLeaveAddBtn}
                  >
                    <button className="add-btn" onClick={handleToggleAddContent} ref={addBtnRef}>
                      <FaPlus />
                    </button>
                    <span className={`tooltip ${isAddContentOpen || isTooltipSuppressed ? 'hidden' : ''}`}>Add Content</span>
                    <AddContentDropdown
                      isOpen={isAddContentOpen}
                      onClose={closeAddContentDropdown}
                      toggleButtonRef={addBtnRef}
                      onAddFilesClick={handleOpenAddFilesDialog}
                      onServiceClick={handleServiceClick}
                      selectedServices={selectedServices}
                    />
                  </div>
                </div>
                <div className="tooltip-wrapper">
                  <button
                    className={`send-btn ${isProcessing ? 'stop-btn' : ''}`}
                    onClick={isProcessing ? handleStop : handleSendButtonClick}
                  >
                    {isProcessing ? <FaStop /> : <FaPaperPlane />}
                  </button>
                  <span className="tooltip">{isProcessing ? 'Stop' : 'Send'}</span>
                </div>
              </div>
            </div>
          </div>
        )}
      </main>
      
      {showSettingsModal && (
        <IntialSetting
          trigger={true}
          setTrigger={() => setShowSettingsModal(false)}
          fromAiPage={true}
          openSnackbar={openSnackbar}
          closeSnackbar={closeSnackbar}
        />
      )}
      {isAddFilesDialogOpen && (
        <AddFilesDialog
          isOpen={isAddFilesDialogOpen}
          onClose={() => setIsAddFilesDialogOpen(false)}
          openSnackbar={openSnackbar}
          setSessionContent={setSessionContent}
        />
      )}
      <Snackbar
        open={snackbar.open}
        autoHideDuration={snackbar.duration}
        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;