script.js

const API_BASE_URL = 'https://perspicacity-backend-a0tq.onrender.com';

if ('serviceWorker' in navigator) {
    window.addEventListener('load', () => {
        navigator.serviceWorker.register('/sw.js', { scope: '/' })
            .then(registration => {
                console.log('ServiceWorker registration successful with scope: ', registration.scope);
            })
            .catch(err => {
                console.log('ServiceWorker registration failed: ', err);
            });
    });
}

document.addEventListener('DOMContentLoaded', () => {
    // Elements
    const settingsModal = document.getElementById('settings-modal');
    const openSettingsBtn = document.getElementById('open-settings');
    const closeSettingsBtn = document.getElementById('close-settings');
    const saveSettingsBtn = document.getElementById('save-settings');
    const apiKeyInput = document.getElementById('api-key-input');
    const googleApiKeyInput = document.getElementById('google-api-key-input');
    const chatForm = document.getElementById('chat-form');
    const queryInput = document.getElementById('query-input');
    const chatMessages = document.getElementById('chat-messages');
    const workflowBtn = document.getElementById('workflow-btn');
    const workflowDropdown = document.getElementById('workflow-dropdown');
    const workflowIcon = workflowBtn.querySelector('.workflow-icon');
    const workflowOptions = document.querySelectorAll('.workflow-option');
    const sendBtn = document.getElementById('send-btn');
    const scrollToBottomBtn = document.getElementById('scroll-to-bottom');
    const inputBox = document.getElementById('input-box');

    // --- REFACTOR 3: STATE OBJECTS ---
    const appState = {
        currentMode: 'search',
        abortController: null,
        isGenerating: false,
        isUserScrolledUp: false,
        currentStreamingDiv: null,
        isStreamingStopped: false,
        conversationCount: 0
    };
    
    const responseData = {
        sourcesMap: {},
        allSourcesList: [],
        allArticlesList: [],
        allBooksList: [],
        researchStepsList: [],
        wordCount: 0
    };
    // --- END REFACTOR 3 ---

    // Set initial state: welcome screen visible, no scroll
    chatMessages.classList.add('has-welcome');

    loadSettings();

    chatMessages.addEventListener('scroll', () => {
        const scrollTop = chatMessages.scrollTop;
        const scrollHeight = chatMessages.scrollHeight;
        const clientHeight = chatMessages.clientHeight;

        if (scrollHeight - scrollTop - clientHeight > 100) {
            appState.isUserScrolledUp = true;
        } else {
            appState.isUserScrolledUp = false;
        }

        // Only show if content is longer than screen
        if (scrollHeight <= clientHeight) {
            scrollToBottomBtn.classList.add('hidden');
            return;
        }

        // Find the last user message (last prompt)
        const lastUserMessage = Array.from(chatMessages.querySelectorAll('.message-user')).pop();

        if (lastUserMessage) {
            const headerHeight = document.querySelector('.chat-header').offsetHeight;
            const lastPromptTop = lastUserMessage.offsetTop;
            const userScrollPosition = scrollTop + headerHeight;

            // Show button only if user is above the last prompt
            if (userScrollPosition < lastPromptTop - 100) {
                scrollToBottomBtn.classList.remove('hidden');
            } else {
                scrollToBottomBtn.classList.add('hidden');
            }
        } else {
            // No messages yet, hide button
            scrollToBottomBtn.classList.add('hidden');
        }
    });


    scrollToBottomBtn.addEventListener('click', () => {
        const lastUserMessage = Array.from(chatMessages.querySelectorAll('.message-user')).pop();

        if (lastUserMessage) {
            const headerHeight = document.querySelector('.chat-header').offsetHeight;
            const targetPosition = lastUserMessage.offsetTop - headerHeight - 20;

            chatMessages.scrollTo({
                top: targetPosition,
                behavior: 'smooth'
            });
        }
    });

    function scrollToMessage(messageElement) {
        setTimeout(() => {
            const headerHeight = document.querySelector('.chat-header').offsetHeight;
            const elementTop = messageElement.offsetTop;
            chatMessages.scrollTo({
                top: elementTop - headerHeight - 20,
                behavior: 'smooth'
            });
        }, 100);
    }

    function extractContextFromDOM(currentMessageDiv) {
        const context = [];
        let previousMessage = currentMessageDiv ? currentMessageDiv.previousElementSibling : null;

        while (previousMessage) {
            if (previousMessage.classList.contains('welcome-screen')) {
                break;
            }

            if (previousMessage.classList.contains('message')) {
                let role = 'user';
                let content = '';

                if (previousMessage.classList.contains('message-user')) {
                    role = 'user';
                    const contentDiv = previousMessage.querySelector('.message-content');
                    if (contentDiv) {
                        content = contentDiv.textContent.trim();
                    }
                // --- REFACTOR 2: CONTEXT PARSING ---
                } else if (previousMessage.classList.contains('message-assistant')) {
                    role = 'assistant';
                    // Read the clean content directly from the dataset.
                    // This is 100% reliable and ignores all UI elements.
                    content = previousMessage.dataset.cleanContent || '';
                }
                // --- END REFACTOR 2 ---

                if (content) {
                    context.unshift({ role, content });
                }
            }

            previousMessage = previousMessage.previousElementSibling;
        }

        return context;
    }

    function updatePlaceholder() {
        if (appState.conversationCount >= 1) {
            queryInput.placeholder = queryInput.dataset.followupPlaceholder;
        } else {
            queryInput.placeholder = queryInput.dataset.initialPlaceholder;
        }
    }

    updatePlaceholder();

    function updateSendButtonState() {
        if (appState.isGenerating) return;
        const query = queryInput.value.trim();
        sendBtn.disabled = query.length === 0;
    }

    queryInput.addEventListener('input', updateSendButtonState);
    updateSendButtonState();

    queryInput.addEventListener('input', () => {
        queryInput.style.height = 'auto';
        queryInput.style.height = Math.min(queryInput.scrollHeight, 150) + 'px';
    });

    queryInput.addEventListener('keydown', (e) => {
        if (e.key === 'Enter' && !e.shiftKey) {
            const isMobile = window.innerWidth < 768;
            if (!isMobile) {
                e.preventDefault();
                if (appState.isGenerating || queryInput.value.trim().length > 0) {
                    chatForm.requestSubmit();
                }
            }
        }
    });

    document.addEventListener('click', (e) => {
        const userMessage = e.target.closest('.message-user');
        if (userMessage && !userMessage.classList.contains('editing')) {
            if (window.innerWidth < 768) {
                document.querySelectorAll('.message-user.active').forEach(msg => {
                    if (msg !== userMessage) msg.classList.remove('active');
                });
                userMessage.classList.toggle('active');
            }
        } else if (!e.target.closest('.user-message-actions')) {
            document.querySelectorAll('.message-user.active').forEach(msg => {
                msg.classList.remove('active');
            });
        }
    });

    workflowBtn.addEventListener('click', (e) => {
        e.stopPropagation();
        workflowDropdown.classList.toggle('hidden');
    });

    document.addEventListener('click', () => {
        workflowDropdown.classList.add('hidden');
    });

    workflowDropdown.addEventListener('click', (e) => {
        e.stopPropagation();
    });

    workflowOptions.forEach(option => {
        option.addEventListener('click', () => {
            const mode = option.dataset.mode;
            appState.currentMode = mode;
            const icon = option.querySelector('svg').cloneNode(true);
            workflowIcon.innerHTML = icon.innerHTML;
            workflowDropdown.classList.add('hidden');
            workflowOptions.forEach(opt => opt.classList.remove('active'));
            option.classList.add('active');
        });
    });

    workflowOptions[0].classList.add('active');

    openSettingsBtn.addEventListener('click', () => {
        settingsModal.classList.remove('hidden');
    });

    closeSettingsBtn.addEventListener('click', () => {
        settingsModal.classList.add('hidden');
    });

    settingsModal.addEventListener('click', (e) => {
        if (e.target === settingsModal) {
            settingsModal.classList.add('hidden');
        }
    });

    saveSettingsBtn.addEventListener('click', async () => {
        const apiKey = apiKeyInput.value.trim();
        const googleApiKey = googleApiKeyInput.value.trim();
        const originalBtnHTML = saveSettingsBtn.innerHTML;

        if (!apiKey || !googleApiKey) {
            showNotification('Please fill in both API keys.', 'error');
            return;
        }

        // 1. Show loading state
        saveSettingsBtn.disabled = true;
        saveSettingsBtn.innerHTML = 'Validating...';

        // 2. --- NEW: Create AbortController for timeout ---
        const controller = new AbortController();
        const timeoutId = setTimeout(() => {
            controller.abort();
        }, 8000); // 8-second timeout, just like your streaming endpoints

        try {
            // 3. Call the new validation endpoint, passing the signal
            const response = await fetch(`${API_BASE_URL}/validate-keys`, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json'
                },
                body: JSON.stringify({
                    api_key: apiKey,
                    google_api_key: googleApiKey
                }),
                signal: controller.signal // --- NEW: Pass the signal ---
            });

            // 4. --- NEW: Clear timeout if response is successful ---
            clearTimeout(timeoutId);

            const data = await response.json();

            if (!response.ok) {
                // 5. Handle validation failure (e.g., "Invalid Gemini API key")
                throw new Error(data.error || 'Validation failed. Please check your keys.');
            }

            // 6. Handle validation success
            localStorage.setItem('gemini_api_key', apiKey);
            localStorage.setItem('google_api_key', googleApiKey);
            showNotification('Settings saved successfully!');
            settingsModal.classList.add('hidden');

        } catch (error) {
            // 7. --- NEW: Clear timeout in case of non-abort error ---
            clearTimeout(timeoutId);

            // 8. --- NEW: Updated catch block to handle AbortError ---
            let errorMessage;
            if (error.name === 'AbortError') {
                // This is our custom 8-second timeout firing
                errorMessage = "Server is waking up. Please try again in a moment.";
            } else if (error.name === 'TypeError') {
                // This is a genuine network/CORS error (less likely but good to keep)
                errorMessage = "Failed to connect to server. Check your network.";
            } else {
                // This is a validation error from the server (e.g., "Invalid Gemini Key")
                errorMessage = error.message || 'An unknown error occurred.';
            }
            showNotification(errorMessage, 'error');
        
        } finally {
            // 9. Always restore the button
            saveSettingsBtn.disabled = false;
            saveSettingsBtn.innerHTML = originalBtnHTML;
        }
    });

    function setSendButtonState(generating) {
        appState.isGenerating = generating;
        if (generating) {
            sendBtn.innerHTML = `<svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><rect x="7" y="7" width="10" height="10" rx="1"></rect></svg>`;
            sendBtn.classList.add('stop-btn');
            sendBtn.disabled = false;
            sendBtn.style.display = 'flex';
        } else {
            sendBtn.innerHTML = `<svg width="20" height="20" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><line x1="22" y1="2" x2="11" y2="13"></line><polygon points="22 2 15 22 11 13 2 9 22 2"></polygon></svg>`;
            sendBtn.classList.remove('stop-btn');
            updateSendButtonState();
        }
    }

    function stopGeneration() {
        console.log('Stop generation called');
        appState.isStreamingStopped = true;
        if (appState.abortController) {
            appState.abortController.abort();
            appState.abortController = null;
        }
    
        // NEW: If there's a current streaming div without action buttons, add them
        if (appState.currentStreamingDiv) {
            const contentDiv = appState.currentStreamingDiv.querySelector('.message-content');
            let answerTab = contentDiv?.querySelector('.response-content');
        
            // If response-content doesn't exist (stopped during spinner), create it
            if (!answerTab) {
                const query = appState.currentStreamingDiv.dataset.query || '';
                const mode = appState.currentStreamingDiv.dataset.mode || 'search';
            
                // Remove any status spinners
                const statusDiv = contentDiv.querySelector('.search-status');
                if (statusDiv) statusDiv.remove();
                
                // Remove any research steps
                const stepsDiv = contentDiv.querySelector('.research-steps');
                if (stepsDiv) stepsDiv.remove();
                
                // Create the tab structure
                const tabsContainer = document.createElement('div');
                tabsContainer.className = 'response-tabs';
            
                // --- CHANGE HERE ---
                // Set tab based on the mode that was running
                const activeTabName = (mode === 'research') ? 'research' : 'answer';
                const activeTabLabel = (mode === 'research') ? 'Research' : 'Answer';
                tabsContainer.innerHTML = `<button class="response-tab active" data-tab="${activeTabName}">${activeTabLabel}</button>`;
                // --- END CHANGE ---
            
                contentDiv.appendChild(tabsContainer);
            
                // Create the answer tab content
                answerTab = document.createElement('div');
                answerTab.className = 'tab-content response-content active';
                
                // --- CHANGE HERE ---
                answerTab.dataset.tab = activeTabName;
                // --- END CHANGE ---
            
                answerTab.innerHTML = '<p style="color: var(--text-secondary);">Generation stopped.</p>';
                contentDiv.appendChild(answerTab);
            
                // Add action buttons
                addActionButtons(answerTab, 'Generation stopped.', [], query, false);
            } else if (!answerTab.querySelector('.message-actions')) {
                // Response-content exists but no buttons yet
                const query = appState.currentStreamingDiv.dataset.query || '';
            
                // Remove any status spinners
                const statusDiv = contentDiv.querySelector('.search-status');
                if (statusDiv) statusDiv.remove();
            
                // Add a stopped message if answer tab is empty
                if (!answerTab.textContent.trim()) {
                    answerTab.innerHTML = '<p style="color: var(--text-secondary);">Generation stopped.</p>';
                }
            
                // Add action buttons
                addActionButtons(answerTab, answerTab.textContent, [], query, false);
            }
        }
    
        setSendButtonState(false);
        showNotification('Generation stopped', 'error');
    }
    
    chatForm.addEventListener('submit', async (e) => {
        e.preventDefault();
        if (appState.isGenerating) {
            stopGeneration();
            return;
        }
        const query = queryInput.value.trim();
        if (!query) return;
        if (appState.abortController) appState.abortController.abort();
        const apiKey = localStorage.getItem('gemini_api_key');
        const googleApiKey = localStorage.getItem('google_api_key');
        if (!apiKey) {
            showNotification('Please set your Gemini API key in settings', 'error');
            settingsModal.classList.remove('hidden');
            return;
        }
        if (!googleApiKey) {
            showNotification('Please set your Google API key in settings', 'error');
            settingsModal.classList.remove('hidden');
            return;
        }

        const welcomeScreen = document.querySelector('.welcome-screen');
        if (welcomeScreen) {
            welcomeScreen.remove();
            // Enable scrolling after welcome screen is removed
            chatMessages.classList.remove('has-welcome');
        }

        // IMMEDIATELY: Switch to compact layout and change placeholder
        if (appState.conversationCount === 0) {
            inputBox.classList.add('compact');
            appState.conversationCount++;
            updatePlaceholder();
        }

        const usedMode = appState.currentMode;
        const userMessageDiv = addMessage(query, 'user', usedMode);

        scrollToMessage(userMessageDiv);

        appState.isUserScrolledUp = false;
        queryInput.value = '';
        queryInput.style.height = 'auto';
        updateSendButtonState();
        responseData.sourcesMap = {};
        responseData.allSourcesList = [];
        responseData.allArticlesList = [];
        responseData.allBooksList = [];
        responseData.researchStepsList = [];
        responseData.wordCount = 0;
        appState.isStreamingStopped = false;
        setSendButtonState(true);

        const conversationHistory = extractContextFromDOM(userMessageDiv);

        if (usedMode === 'search') {
            handleSearchMode(apiKey, googleApiKey, query, conversationHistory);
        } else {
            handleResearchMode(apiKey, googleApiKey, query, conversationHistory);
        }
    });

    // --- REFACTOR 1: DRY ---
    // This new function contains the logic previously duplicated
    // across handleSearchMode and handleSearchModeRegenerate.
    async function handleStreamError(error, messageDiv, query, mode, didTimeout) {
        // Clean up any UI spinners
        const contentDiv = messageDiv.querySelector('.message-content');
        const statusDiv = contentDiv.querySelector('.search-status');
        if (statusDiv) statusDiv.remove();
        const stepsDiv = contentDiv.querySelector('.research-steps');
        if (stepsDiv) stepsDiv.remove();
    
        let errorMessage;
        if (error.name === 'AbortError' && didTimeout) {
            // This was our 8-second timeout
            errorMessage = "I couldn't get a response from the server. It might be busy or just waking up. Please try sending your message again in a moment.";
        } else if (error.name === 'TypeError') {
            // This is a network or CORS error
            errorMessage = "I'm having trouble connecting. Please check your network connection and try again.";
        } else if (error.name === 'AbortError') {
            // This was a user "Stop" click.
            // The stopGeneration() function handles the UI update, so we do nothing.
            console.log('Stream stopped by user.');
            return; 
        } else {
            // Any other general error
            errorMessage = error.message || `I encountered an issue during the ${mode}. Please try again.`;
        }
    
        // Stream the error message to the UI
        await streamResponse(messageDiv, errorMessage, [], query, mode);
    }
    
    // --- REFACTOR 1: DRY ---
    // This new function contains the core logic for a search request.
    async function runSearchRequest(apiKey, googleApiKey, query, messageDiv, conversationHistory) {
        const contentDiv = messageDiv.querySelector('.message-content');
        let searchStatusDiv = null;
        let timeoutId = null;
        let didTimeout = false;
    
        // --- REFACTOR 3: STATE ---
        // Create a local controller and assign it to the global state
        const localAbortController = new AbortController();
        appState.abortController = localAbortController;
        // --- END REFACTOR 3 ---
    
        timeoutId = setTimeout(() => {
            didTimeout = true;
            // --- REFACTOR 3: STATE ---
            if (appState.abortController === localAbortController) {
                appState.abortController.abort();
            }
            // --- END REFACTOR 3 ---
        }, 8000);
    
        const headers = { 'Content-Type': 'application/json' };
        const body = JSON.stringify({
            api_key: apiKey,
            google_api_key: googleApiKey,
            query: query,
            conversation_history: conversationHistory
        });
    
        try {
            const response = await fetch(`${API_BASE_URL}/search`, {
                method: 'POST',
                headers,
                body,
                // --- REFACTOR 3: STATE ---
                signal: localAbortController.signal
                // --- END REFACTOR 3 ---
            });
    
            const contentType = response.headers.get('content-type');
    
            if (contentType && contentType.includes('text/event-stream')) {
                const reader = response.body.getReader();
                const decoder = new TextDecoder();
                let buffer = '';
    
                while (true) {
                    const { done, value } = await reader.read();
                    clearTimeout(timeoutId); // Server is alive
    
                    if (done) break;
    
                    buffer += decoder.decode(value, { stream: true });
                    const lines = buffer.split('\n');
                    buffer = lines.pop() || '';
    
                    for (const line of lines) {
                        if (line.startsWith('data: ')) {
                            const data = line.slice(6);
                            if (data === 'Stream finished.') continue;
    
                            try {
                                const result = JSON.parse(data);
    
                                if (result.status) {
                                    if (!searchStatusDiv) {
                                        searchStatusDiv = document.createElement('div');
                                        searchStatusDiv.className = 'search-status';
                                        searchStatusDiv.innerHTML = '<div class="status-spinner"></div><span class="status-text"></span>';
                                        contentDiv.appendChild(searchStatusDiv);
                                    }
                                    const statusText = searchStatusDiv.querySelector('.status-text');
                                    if (statusText) {
                                        statusText.textContent = result.status;
                                    }
                                }
    
                                if (result.final_answer) {
                                    if (searchStatusDiv) searchStatusDiv.remove();
    
                                    if (result.sources) {
                                        responseData.allSourcesList = result.sources;
                                        result.sources.forEach((source, index) => {
                                            responseData.sourcesMap[index + 1] = source;
                                        });
                                    }
    
                                    if (result.articles) {
                                        responseData.allArticlesList = result.articles;
                                    }
    
                                    if (result.books) {
                                        responseData.allBooksList = result.books;
                                    }
    
                                    await streamResponse(messageDiv, result.final_answer, result.sources, query, 'search', result.show_ask_scholar_button);
                                }
                            } catch (e) {
                                console.error('Error parsing JSON:', e);
                            }
                        }
                    }
                }
            }
        } catch (error) {
            clearTimeout(timeoutId);
            await handleStreamError(error, messageDiv, query, 'search', didTimeout);
        } finally {
            clearTimeout(timeoutId);
            
            // --- REFACTOR 3: STATE ---
            if (appState.currentStreamingDiv === messageDiv) {
                 appState.currentStreamingDiv = null;
            }
            // Only clear the controller and button if THIS stream was the last one
            if (appState.abortController === localAbortController) { 
                 setSendButtonState(false);
                 appState.abortController = null;
            }
            // --- END REFACTOR 3 ---
        }
    }

    // --- REFACTOR 1: DRY ---
    // This new function contains the core logic for a research request.
    async function runResearchRequest(apiKey, googleApiKey, query, messageDiv, conversationHistory) {
        const contentDiv = messageDiv.querySelector('.message-content');
        const stepsContainer = document.createElement('div');
        stepsContainer.className = 'research-steps';
        contentDiv.appendChild(stepsContainer);
        let currentStepDiv = null;
    
        let timeoutId = null;
        let didTimeout = false;
        
        // --- REFACTOR 3: STATE ---
        const localAbortController = new AbortController();
        appState.abortController = localAbortController;
        // --- END REFACTOR 3 ---
    
        timeoutId = setTimeout(() => {
            didTimeout = true;
            // --- REFACTOR 3: STATE ---
            if (appState.abortController === localAbortController) {
                appState.abortController.abort();
            }
            // --- END REFACTOR 3 ---
        }, 8000);
    
        try {
            const response = await fetch(`${API_BASE_URL}/research`, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json'
                },
                body: JSON.stringify({
                    query: query,
                    api_key: apiKey,
                    google_api_key: googleApiKey,
                    conversation_history: conversationHistory
                }),
                // --- REFACTOR 3: STATE ---
                signal: localAbortController.signal
                // --- END REFACTOR 3 ---
            });
    
            const reader = response.body.getReader();
            const decoder = new TextDecoder();
            let buffer = '';
    
            while (true) {
                const { done, value } = await reader.read();
                clearTimeout(timeoutId); // Server is alive
    
                if (done) break;
    
                buffer += decoder.decode(value, { stream: true });
                const lines = buffer.split('\n');
                buffer = lines.pop() || '';
    
                for (const line of lines) {
                    if (line.startsWith('data: ')) {
                        const data = line.slice(6);
                        if (data === 'Stream finished.') continue;
    
                        try {
                            const result = JSON.parse(data);
                            if (result.status) {
                                if (currentStepDiv) currentStepDiv.classList.add('completed');
                                currentStepDiv = addResearchStep(stepsContainer, result.status);
                                responseData.researchStepsList.push(result.status);
                                periodicScroll();
                            } else if (result.final_answer) {
                                stepsContainer.remove();
                                if (result.sources) {
                                    responseData.allSourcesList = result.sources;
                                    result.sources.forEach((source, index) => {
                                        responseData.sourcesMap[index + 1] = source;
                                    });
                                }
                                if (result.articles) {
                                    responseData.allArticlesList = result.articles;
                                }
                                if (result.books) {
                                    responseData.allBooksList = result.books;
                                }
                                await streamResponse(messageDiv, result.final_answer, result.sources, query, 'research', result.show_ask_scholar_button);
                            } else if (result.error) {
                                // This is a streamed error from the backend *after* connection
                                stepsContainer.remove();
                                const errorMessage = result.message || result.error || 'An error occurred during research.';
                                await streamResponse(messageDiv, errorMessage, [], query, 'research');
                            }
                        } catch (e) {
                            console.error('Error parsing JSON:', e);
                        }
                    }
                }
            }
        } catch (error) {
            clearTimeout(timeoutId);
            await handleStreamError(error, messageDiv, query, 'research', didTimeout);
        } finally {
            clearTimeout(timeoutId);

            // --- REFACTOR 3: STATE ---
            if (appState.currentStreamingDiv === messageDiv) {
                 appState.currentStreamingDiv = null;
            }
            // Only clear the controller and button if THIS stream was the last one
            if (appState.abortController === localAbortController) { 
                 setSendButtonState(false);
                 appState.abortController = null;
            }
            // --- END REFACTOR 3 ---
        }
    }
    
    // --- REFACTOR 1: DRY ---
    // The original handleSearchMode is now just a simple wrapper.
    async function handleSearchMode(apiKey, googleApiKey, query, conversationHistory) {
        const messageDiv = addMessage('', 'assistant', 'search');
        messageDiv.dataset.query = query;
        messageDiv.dataset.mode = 'search';
        appState.currentStreamingDiv = messageDiv;
    
        // Just call the new core function
        await runSearchRequest(apiKey, googleApiKey, query, messageDiv, conversationHistory);
    }
    
    // --- REFACTOR 1: DRY ---
    // The original handleResearchMode is now just a simple wrapper.
    async function handleResearchMode(apiKey, googleApiKey, query, conversationHistory) {
        const messageDiv = addMessage('', 'assistant', 'research');
        messageDiv.dataset.query = query;
        messageDiv.dataset.mode = 'research';
        appState.currentStreamingDiv = messageDiv;
    
        // Just call the new core function
        await runResearchRequest(apiKey, googleApiKey, query, messageDiv, conversationHistory);
    }
    // --- END REFACTOR 1 ---


    function addMessage(content, role, modeOrLoadingType = null) {
        const messageDiv = document.createElement('div');
        messageDiv.className = `message message-${role}`;
        const contentDiv = document.createElement('div');
        contentDiv.className = 'message-content';
        if (role === 'user') {
            messageDiv.dataset.workflow = modeOrLoadingType || 'search';
            messageDiv.dataset.query = content;
            messageDiv.dataset.originalQuery = content;
            contentDiv.textContent = content;
            addUserActionButtons(messageDiv, content, modeOrLoadingType || 'search');
        } else if (modeOrLoadingType === 'search') {
            contentDiv.innerHTML = '';
        } else if (modeOrLoadingType === 'research') {
            contentDiv.innerHTML = '';
        } else {
            contentDiv.textContent = content;
        }
        messageDiv.appendChild(contentDiv);
        chatMessages.appendChild(messageDiv);
        return messageDiv;
    }

    function addUserActionButtons(messageDiv, query, workflow) {
        const actionsDiv = document.createElement('div');
        actionsDiv.className = 'user-message-actions';

        const copyBtn = document.createElement('button');
        copyBtn.className = 'user-action-btn';
        copyBtn.innerHTML = `<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><rect x="9" y="9" width="13" height="13" rx="2" ry="2"></rect><path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1"></path></svg>`;
        copyBtn.title = 'Copy';
        
        copyBtn.addEventListener('click', (e) => {
            e.stopPropagation();
            e.preventDefault();
            // Get the parent message and read its current query
            const messageDiv = e.target.closest('.message-user');
            const currentQuery = messageDiv.dataset.query;
            copyUserQuery(currentQuery, copyBtn);
        });

        copyBtn.addEventListener('touchend', (e) => {
            e.stopPropagation();
            e.preventDefault();
            // Get the parent message and read its current query
            const messageDiv = e.target.closest('.message-user');
            const currentQuery = messageDiv.dataset.query;
            copyUserQuery(currentQuery, copyBtn);
        });

        
        const editBtn = document.createElement('button');
        editBtn.className = 'user-action-btn';
        editBtn.innerHTML = `<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M11 4H4a2 2 0 0 0-2 2v14a2 2 0 0 0 2 2h14a2 2 0 0 0 2-2v-7"></path><path d="M18.5 2.5a2.121 2.121 0 0 1 3 3L12 15l-4 1 1-4 9.5-9.5z"></path></svg>`;
        editBtn.title = 'Edit';
        editBtn.addEventListener('click', (e) => {
            e.stopPropagation();
            e.preventDefault();
            enterEditMode(messageDiv);
        });

        editBtn.addEventListener('touchend', (e) => {
            e.stopPropagation();
            e.preventDefault();
            enterEditMode(messageDiv);
        });
        
        actionsDiv.appendChild(copyBtn);
        actionsDiv.appendChild(editBtn);
        messageDiv.appendChild(actionsDiv);
    }

    function copyUserQuery(query, button) {
        navigator.clipboard.writeText(query).then(() => {
            const originalHTML = button.innerHTML;
            button.innerHTML = `<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><polyline points="20 6 9 17 4 12"></polyline></svg>`;
            button.classList.add('copied');
            setTimeout(() => {
                button.innerHTML = originalHTML;
                button.classList.remove('copied');
            }, 2000);
        }).catch(err => {
            showNotification('Failed to copy', 'error');
        });
    }

    function enterEditMode(messageDiv) {
        const contentDiv = messageDiv.querySelector('.message-content');
        const currentQuery = messageDiv.dataset.query;
        const actionsDiv = messageDiv.querySelector('.user-message-actions');
        if (actionsDiv) actionsDiv.style.display = 'none';

        const editWrapper = document.createElement('div');
        editWrapper.className = 'edit-wrapper';
        const textarea = document.createElement('textarea');
        textarea.className = 'edit-textarea';
        textarea.value = currentQuery;
        textarea.rows = 1;

        textarea.addEventListener('keydown', (e) => {
            if (e.key === 'Enter' && !e.shiftKey) {
                const isMobile = window.innerWidth < 768;
                if (!isMobile) {
                    e.preventDefault();
                    const newQuery = textarea.value.trim();
                    const isEmpty = newQuery.length === 0;
                    if (!isEmpty) {
                        confirmEdit(messageDiv, newQuery);
                    }
                }
            }
        });

        const editActionsDiv = document.createElement('div');
        editActionsDiv.className = 'edit-actions';

        const cancelBtn = document.createElement('button');
        cancelBtn.className = 'edit-action-btn cancel-btn';
        cancelBtn.textContent = 'Cancel';
        cancelBtn.onclick = () => exitEditMode(messageDiv, currentQuery);

        const confirmBtn = document.createElement('button');
        confirmBtn.className = 'edit-action-btn confirm-btn';
        confirmBtn.textContent = 'Done';
        confirmBtn.onclick = () => confirmEdit(messageDiv, textarea.value.trim());

        textarea.addEventListener('input', () => {
            textarea.style.height = 'auto';
            textarea.style.height = Math.min(textarea.scrollHeight, 150) + 'px';
        });

        editActionsDiv.appendChild(cancelBtn);
        editActionsDiv.appendChild(confirmBtn);
        editWrapper.appendChild(textarea);
        editWrapper.appendChild(editActionsDiv);
        contentDiv.innerHTML = '';
        contentDiv.appendChild(editWrapper);
        textarea.focus();
        textarea.style.height = 'auto';
        textarea.style.height = Math.min(textarea.scrollHeight, 150) + 'px';
        messageDiv.classList.add('editing');
    }

    function exitEditMode(messageDiv, currentQuery) {
        const contentDiv = messageDiv.querySelector('.message-content');
        contentDiv.textContent = currentQuery;
        const actionsDiv = messageDiv.querySelector('.user-message-actions');
        if (actionsDiv) actionsDiv.style.display = 'flex';
        messageDiv.classList.remove('editing');
    }

    function confirmEdit(messageDiv, newQuery) {
        if (!newQuery) return;
        const currentQuery = messageDiv.dataset.query;

        if (newQuery === currentQuery) {
            exitEditMode(messageDiv, currentQuery);
            return;
        }

        messageDiv.dataset.query = newQuery;
        const contentDiv = messageDiv.querySelector('.message-content');
        contentDiv.textContent = newQuery;
        const actionsDiv = messageDiv.querySelector('.user-message-actions');
        if (actionsDiv) actionsDiv.style.display = 'flex';
        messageDiv.classList.remove('editing');

        const responseDiv = messageDiv.nextElementSibling;
        if (!responseDiv || !responseDiv.classList.contains('message-assistant')) {
            showNotification('Could not find response to regenerate', 'error');
            return;
        }

        scrollToMessage(messageDiv);

        const workflow = responseDiv.dataset.mode || messageDiv.dataset.workflow;
        regenerateWithNewQuery(responseDiv, newQuery, workflow);
    }

    function addResearchStep(container, text) {
        let cleanText = text;
        cleanText = cleanText.replace(/\*\*\*([^*]+)\*\*\*/g, '$1');
        cleanText = cleanText.replace(/\*\*([^*]+)\*\*\*/g, '$1');
        cleanText = cleanText.replace(/\*([^*]+)\*/g, '$1');
        cleanText = cleanText.replace(/_([^_]+)_/g, '$1');
        cleanText = cleanText.replace(/`([^`]+)`/g, '$1');
        cleanText = cleanText.replace(/["']/g, '');
        cleanText = cleanText.replace(/[#]+\s*/g, '');
        cleanText = cleanText.replace(/^\d+\.\s*/g, '');
        cleanText = cleanText.replace(/^[-*+]\s*/g, '');
        cleanText = cleanText.replace(/\[\d+\]/g, '');
        cleanText = cleanText.replace(/\s+/g, ' ').trim();

        const stepDiv = document.createElement('div');
        stepDiv.className = 'research-step';
        stepDiv.innerHTML = `<div class="step-icon"></div><div class="step-text">${cleanText}</div>`;
        container.appendChild(stepDiv);
        return stepDiv;
    }

    function periodicScroll() {
        // 1. If the user has scrolled up, do nothing. This gives them freedom.
        if (appState.isUserScrolledUp) return; 
    
        // 2. If they are at the bottom, scroll to the new bottom.
        // This is called periodically (every 50 words or every research step)
        // ensuring the scroll follows the streaming text.
        chatMessages.scrollTo({
            top: chatMessages.scrollHeight,
            behavior: 'auto' // 'auto' is instant, 'smooth' can lag behind streaming
        });
    }

    async function streamResponse(messageDiv, text, sources, originalQuery, mode, showAskButton = false) {
        const contentDiv = messageDiv.querySelector('.message-content');
        contentDiv.innerHTML = '';

        const tabsContainer = document.createElement('div');
        tabsContainer.className = 'response-tabs';

        // Build the tabs HTML dynamically
        let tabsHTML = '';
        if (mode === 'search') {
            tabsHTML += `<button class="response-tab active" data-tab="answer">Answer</button>`;
            if (responseData.allSourcesList && responseData.allSourcesList.length > 0) {
                tabsHTML += `<button class="response-tab" data-tab="sources">Sources</button>`;
            }
            if (responseData.allArticlesList && responseData.allArticlesList.length > 0) {
                tabsHTML += `<button class="response-tab" data-tab="articles">Articles</button>`;
            }
            if (responseData.allBooksList && responseData.allBooksList.length > 0) {
                tabsHTML += `<button class="response-tab" data-tab="books">Books</button>`;
            }
        } else { // research mode
            tabsHTML += `<button class="response-tab active" data-tab="research">Research</button>`;
            if (responseData.allSourcesList && responseData.allSourcesList.length > 0) {
                tabsHTML += `<button class="response-tab" data-tab="sources">Sources</button>`;
            }
            if (responseData.allArticlesList && responseData.allArticlesList.length > 0) {
                tabsHTML += `<button class="response-tab" data-tab="articles">Articles</button>`;
            }
            if (responseData.allBooksList && responseData.allBooksList.length > 0) {
                tabsHTML += `<button class="response-tab" data-tab="books">Books</button>`;
            }
            if (responseData.researchStepsList && responseData.researchStepsList.length > 0) {
                tabsHTML += `<button class="response-tab" data-tab="steps">Steps</button>`;
            }
        }
        tabsContainer.innerHTML = tabsHTML;

        const answerTabContent = document.createElement('div');
        answerTabContent.className = 'tab-content response-content active';
        answerTabContent.dataset.tab = mode === 'research' ? 'research' : 'answer';

        contentDiv.appendChild(tabsContainer);
        contentDiv.appendChild(answerTabContent);

        let sourcesTabContent = null;
        if (responseData.allSourcesList && responseData.allSourcesList.length > 0) {
            sourcesTabContent = document.createElement('div');
            sourcesTabContent.className = 'tab-content';
            sourcesTabContent.dataset.tab = 'sources';
            contentDiv.appendChild(sourcesTabContent);
        }
    
        // ADDED: Create Articles and Books tab content
        let articlesTabContent = null;
        if (responseData.allArticlesList && responseData.allArticlesList.length > 0) {
            articlesTabContent = document.createElement('div');
            articlesTabContent.className = 'tab-content';
            articlesTabContent.dataset.tab = 'articles';
            contentDiv.appendChild(articlesTabContent);
        }
    
        let booksTabContent = null;
        if (responseData.allBooksList && responseData.allBooksList.length > 0) {
            booksTabContent = document.createElement('div');
            booksTabContent.className = 'tab-content';
            booksTabContent.dataset.tab = 'books';
            contentDiv.appendChild(booksTabContent);
        }

        let stepsTabContent = null;
        if (mode === 'research' && responseData.researchStepsList && responseData.researchStepsList.length > 0) {
            stepsTabContent = document.createElement('div');
            stepsTabContent.className = 'tab-content';
            stepsTabContent.dataset.tab = 'steps';
            contentDiv.appendChild(stepsTabContent);
        }

        tabsContainer.querySelectorAll('.response-tab').forEach(tab => {
            tab.addEventListener('click', () => switchTab(messageDiv, tab.dataset.tab));
        });

        let processedText = text.replace(/\n\nSources:\s*\[.*?\]\(.*?\).*$/s, '')
            .replace(/\n\n---\n\*This answer is AI-generated.*$/s, '')
            .replace(/This answer is AI-generated.*$/s, '');

        const htmlContent = marked.parse(processedText);
        const parsedHTML = htmlContent.replace(/\[(\d+)\]/g, (match, num) => {
            const source = responseData.sourcesMap[num];
            if (source) {
                return `<a href="${source.url}" target="_blank" rel="noopener" class="inline-citation" title="${source.title}">[${num}]</a>`;
            }
            return match;
        });
        const words = parsedHTML.split(' ');
        let currentHTML = '';
        responseData.wordCount = 0;
        for (let i = 0; i < words.length; i++) {
            if (appState.isStreamingStopped) break;
            currentHTML += words[i] + ' ';
            answerTabContent.innerHTML = currentHTML;
            responseData.wordCount++;
            if (responseData.wordCount % 50 === 0) periodicScroll();
            await new Promise(resolve => setTimeout(resolve, 20));
        }

        // --- REFACTOR 2: CONTEXT ---
        // Set the clean content on the dataset for reliable history
        messageDiv.dataset.cleanContent = processedText;
        // --- END REFACTOR 2 ---

        if (sourcesTabContent) {
            sourcesTabContent.innerHTML = renderSourcesGrid(responseData.allSourcesList);
        }
        // ADDED: Populate Articles and Books tabs
        if (articlesTabContent) {
            articlesTabContent.innerHTML = renderSourcesGrid(responseData.allArticlesList);
        }
        if (booksTabContent) {
            booksTabContent.innerHTML = renderSourcesGrid(responseData.allBooksList);
        }
        if (mode === 'research' && stepsTabContent) {
            stepsTabContent.innerHTML = renderStepsDisplay(responseData.researchStepsList);
        }
    
        addActionButtons(answerTabContent, processedText, sources, originalQuery, showAskButton);

        appState.currentStreamingDiv = null;
        setSendButtonState(false);
    }

    function switchTab(messageDiv, tabName) {
        const contentDiv = messageDiv.querySelector('.message-content');

        contentDiv.querySelectorAll('.response-tab').forEach(tab => {
            if (tab.dataset.tab === tabName) {
                tab.classList.add('active');
            } else {
                tab.classList.remove('active');
            }
        });

        contentDiv.querySelectorAll('.tab-content').forEach(content => {
            if (content.dataset.tab === tabName) {
                content.classList.add('active');
            } else {
                content.classList.remove('active');
            }
        });
    }

    // Replace showTabsForSearch
    function showTabsForSearch(messageDiv) {
        const contentDiv = messageDiv.querySelector('.message-content');
        const sourcesTabContent = contentDiv.querySelector('[data-tab="sources"]');
        if (sourcesTabContent) {
            sourcesTabContent.innerHTML = renderSourcesGrid(responseData.allSourcesList);
        }
        const articlesTabContent = contentDiv.querySelector('[data-tab="articles"]');
        if (articlesTabContent) {
            articlesTabContent.innerHTML = renderSourcesGrid(responseData.allArticlesList);
        }
        const booksTabContent = contentDiv.querySelector('[data-tab="books"]');
        if (booksTabContent) {
            booksTabContent.innerHTML = renderSourcesGrid(responseData.allBooksList);
        }
    }
    
    // Replace showTabsForResearch
    function showTabsForResearch(messageDiv) {
        const contentDiv = messageDiv.querySelector('.message-content');
        const sourcesTabContent = contentDiv.querySelector('[data-tab="sources"]');
        const stepsTabContent = contentDiv.querySelector('[data-tab="steps"]');
        const articlesTabContent = contentDiv.querySelector('[data-tab="articles"]');
        const booksTabContent = contentDiv.querySelector('[data-tab="books"]');
    
        if (sourcesTabContent) {
            sourcesTabContent.innerHTML = renderSourcesGrid(responseData.allSourcesList);
        }
        if (articlesTabContent) {
            articlesTabContent.innerHTML = renderSourcesGrid(responseData.allArticlesList);
        }
        if (booksTabContent) {
            booksTabContent.innerHTML = renderSourcesGrid(responseData.allBooksList);
        }
        if (stepsTabContent) {
            stepsTabContent.innerHTML = renderStepsDisplay(responseData.researchStepsList);
        }
    }
    
    function renderSourcesGrid(sources) {
        if (!sources || sources.length === 0) {
            return '<p style="color: var(--text-secondary); padding: 1rem;">No sources available.</p>';
        }

        const INITIAL_SHOW = 5;
        const showMore = sources.length > INITIAL_SHOW;
        const initialSources = sources.slice(0, INITIAL_SHOW);
        const remainingSources = sources.slice(INITIAL_SHOW);

        let html = '<div class="sources-grid">';

        initialSources.forEach((source, index) => {
            html += renderSourceCard(source, index);
        });

        if (showMore) {
            html += '<div class="remaining-sources hidden">';
            remainingSources.forEach((source, index) => {
                html += renderSourceCard(source, INITIAL_SHOW + index);
            });
            html += '</div>';
            html += `<button class="show-more-sources" onclick="toggleRemainingSourcesHandler(this)">Show more</button>`;
        }

        html += '</div>';
        return html;
    }

    function renderSourceCard(source, index) {
        const url = new URL(source.url);
        const domain = url.hostname.replace('www.', '');
        const favicon = `https://www.google.com/s2/favicons?domain=${domain}&sz=32`;

        const title = source.title.length > 100 ? source.title.substring(0, 100) + '...' : source.title;
        const description = (source.text || source.description || '').substring(0, 150).trim() + '...';

        return `
            <a href="${source.url}" target="_blank" rel="noopener" class="source-card">
                <div class="source-card-header">
                    <img src="${favicon}" alt="${domain}" class="source-favicon" onerror="this.style.display='none'">
                    <span class="source-domain">${domain}</span>
                </div>
                <div class="source-title">${title}</div>
                ${description ? `<div class="source-description">${description}</div>` : ''}
            </a>
        `;
    }

    function renderStepsDisplay(steps) {
        if (!steps || steps.length === 0) {
            return '<p style="color: var(--text-secondary); padding: 1rem;">No research steps available.</p>';
        }

        let html = '<div class="research-steps-display">';
        steps.forEach((step, index) => {
            let cleanText = step.replace(/\*\*\*([^*]+)\*\*\*/g, '$1')
                .replace(/\*\*([^*]+)\*\*\*/g, '$1')
                .replace(/\*([^*]+)\*/g, '$1')
                .replace(/_([^_]+)_/g, '$1')
                .replace(/`([^`]+)`/g, '$1')
                .replace(/["']/g, '')
                .replace(/[#]+\s*/g, '')
                .replace(/^\d+\.\s*/g, '')
                .replace(/^[-*+]\s*/g, '')
                .replace(/\[\d+\]/g, '')
                .replace(/\s+/g, ' ').trim();

            html += `
                <div class="research-step-card">
                    <div class="step-text">${cleanText}</div>
                </div>
            `;
        });
        html += '</div>';
        return html;
    }

    window.toggleRemainingSourcesHandler = function(button) {
        const remainingDiv = button.previousElementSibling;
        remainingDiv.classList.remove('hidden');
        button.remove();
    };

    function addActionButtons(answerTabContent, markdownText, sources, originalQuery, showAskButton = false) {
        const existingActions = answerTabContent.querySelector('.message-actions');
        if (existingActions) return;

        const actionsDiv = document.createElement('div');
        actionsDiv.className = 'message-actions';

        // Conditionally create and add the "Ask a Scholar" button
        if (showAskButton) {
            const askScholarBtn = document.createElement('a');
            askScholarBtn.className = 'action-btn ask-scholar-btn';
            askScholarBtn.title = 'Ask a Scholar at IslamQA.info';
            askScholarBtn.href = 'https://islamqa.info/en/ask';
            askScholarBtn.target = '_blank';
            askScholarBtn.rel = 'noopener noreferrer';
            askScholarBtn.innerHTML = `<svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M21 11.5a8.38 8.38 0 0 1-.9 3.8 8.5 8.5 0 0 1-7.6 4.7 8.38 8.38 0 0 1-3.8-.9L3 21l1.9-5.7a8.38 8.38 0 0 1-.9-3.8 8.5 8.5 0 0 1 4.7-7.6 8.38 8.38 0 0 1 3.8-.9h.5a8.48 8.48 0 0 1 8 8v.5z"></path></svg>`;
            actionsDiv.appendChild(askScholarBtn);
        }

        const copyBtn = document.createElement('button');
        copyBtn.className = 'action-btn copy-btn';
        copyBtn.title = 'Copy';
        copyBtn.innerHTML = `<svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><rect x="9" y="9" width="13" height="13" rx="2" ry="2"></rect><path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1"></path></svg>`;
        copyBtn.onclick = () => copyResponse(markdownText, sources, originalQuery, copyBtn);

        const regenBtn = document.createElement('button');
        regenBtn.className = 'action-btn regenerate-btn';
        regenBtn.title = 'Regenerate';
        regenBtn.innerHTML = `<svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><polyline points="23 4 23 10 17 10"></polyline><path d="M20.49 15a9 9 0 1 1-2.12-9.36L23 10"></path></svg>`;
        regenBtn.onclick = () => showRegenerateDropdown(regenBtn, answerTabContent.closest('.message-content'));
        
        actionsDiv.appendChild(copyBtn);
        actionsDiv.appendChild(regenBtn);
        answerTabContent.appendChild(actionsDiv);
    }

    function copyResponse(markdownText, sources, originalQuery, button) {
        // Use the passed-in markdownText directly
        let cleanedMarkdown = markdownText.trim();

        let copyText = `${originalQuery}\n\n${cleanedMarkdown}\n\n`;

        if (sources && sources.length > 0) {
            copyText += 'Citations:\n';
            // Create a numbered list of sources
            sources.forEach((source, index) => {
                copyText += `[${index + 1}] ${source.url}\n`;
            });
        }

        navigator.clipboard.writeText(copyText).then(() => {
            const originalHTML = button.innerHTML;
            button.innerHTML = `<svg width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><polyline points="20 6 9 17 4 12"></polyline></svg>`;
            button.classList.add('copied');
            setTimeout(() => {
                button.innerHTML = originalHTML;
                button.classList.remove('copied');
            }, 2000);
        }).catch(err => {
            showNotification('Failed to copy', 'error');
        });
    }

    function showRegenerateDropdown(button, messageContent) {
        const messageDiv = messageContent.closest('.message-assistant');
        const existingDropdown = document.querySelector('.regenerate-dropdown');
        if (existingDropdown) existingDropdown.remove();
        const dropdown = document.createElement('div');
        dropdown.className = 'regenerate-dropdown';
        dropdown.innerHTML = `
            <div class="regenerate-option" data-mode="search">
                <svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
                    <circle cx="11" cy="11" r="8"></circle>
                    <path d="m21 21-4.35-4.35"></path>
                </svg>
                <span>Search</span>
            </div>
            <div class="regenerate-option" data-mode="research">
                <svg width="18" height="18" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2">
                    <path d="M2 3h6a4 4 0 0 1 4 4v14a3 3 0 0 0-3-3H2z"></path>
                    <path d="M22 3h-6a4 4 0 0 0-4 4v14a3 3 0 0 1 3-3h7z"></path>
                </svg>
                <span>Research</span>
            </div>
        `;
        button.parentElement.style.position = 'relative';
        button.parentElement.appendChild(dropdown);
        const options = dropdown.querySelectorAll('.regenerate-option');
        options.forEach(option => {
            option.onclick = () => {
                const mode = option.dataset.mode;
                regenerateResponse(messageDiv, mode);
                dropdown.remove();
            };
        });
        setTimeout(() => {
            document.addEventListener('click', function closeDropdown(e) {
                if (!dropdown.contains(e.target) && e.target !== button) {
                    dropdown.remove();
                    document.removeEventListener('click', closeDropdown);
                }
            });
        }, 0);
    }

    function regenerateResponse(messageDiv, newMode) {
        const query = messageDiv.dataset.query;
        if (!query) return;
        if (appState.isGenerating) {
            if (appState.abortController) {
                appState.abortController.abort();
                appState.abortController = null;
            }
        }

        scrollToMessage(messageDiv);

        const contentDiv = messageDiv.querySelector('.message-content');
        while (contentDiv.firstChild) {
            contentDiv.removeChild(contentDiv.firstChild);
        }

        messageDiv.dataset.mode = newMode;
        responseData.sourcesMap = {};
        responseData.allSourcesList = [];
        responseData.allArticlesList = [];
        responseData.allBooksList = [];
        responseData.researchStepsList = [];
        responseData.wordCount = 0;
        appState.isUserScrolledUp = false;
        appState.currentStreamingDiv = messageDiv;
        appState.isStreamingStopped = false;
        const apiKey = localStorage.getItem('gemini_api_key');
        const googleApiKey = localStorage.getItem('google_api_key');

        const userMessageDiv = messageDiv.previousElementSibling; // Get the user message before the assistant one
        const conversationHistory = extractContextFromDOM(userMessageDiv); // Pass the user message
        
        setSendButtonState(true);
        if (newMode === 'search') {
            handleSearchModeRegenerate(apiKey, googleApiKey, query, messageDiv, conversationHistory);
        } else {
            handleResearchModeRegenerate(apiKey, googleApiKey, query, messageDiv, conversationHistory);
        }
    }

    function regenerateWithNewQuery(responseDiv, newQuery, workflow) {
        if (appState.isGenerating) {
            if (appState.abortController) {
                appState.abortController.abort();
                appState.abortController = null;
            }
        }
        const contentDiv = responseDiv.querySelector('.message-content');
        while (contentDiv.firstChild) {
            contentDiv.removeChild(contentDiv.firstChild);
        }
        responseDiv.dataset.query = newQuery;
        responseDiv.dataset.mode = workflow;
        responseData.sourcesMap = {};
        responseData.allSourcesList = [];
        responseData.allArticlesList = [];
        responseData.allBooksList = [];
        responseData.researchStepsList = [];
        responseData.wordCount = 0;
        appState.isUserScrolledUp = false;
        appState.currentStreamingDiv = responseDiv;
        appState.isStreamingStopped = false;
        const apiKey = localStorage.getItem('gemini_api_key');
        const googleApiKey = localStorage.getItem('google_api_key');

        const userMessageDiv = responseDiv.previousElementSibling; // Get the user message before the assistant one
        const conversationHistory = extractContextFromDOM(userMessageDiv); // Pass the user message
        setSendButtonState(true);
        if (workflow === 'search') {
            handleSearchModeRegenerate(apiKey, googleApiKey, newQuery, responseDiv, conversationHistory);
        } else {
            handleResearchModeRegenerate(apiKey, googleApiKey, newQuery, responseDiv, conversationHistory);
        }
    }

    // --- REFACTOR 1: DRY ---
    // The original handleSearchModeRegenerate is now a simple wrapper.
    async function handleSearchModeRegenerate(apiKey, googleApiKey, query, messageDiv, conversationHistory) {
        // Just call the new core function
        await runSearchRequest(apiKey, googleApiKey, query, messageDiv, conversationHistory);
    }
    
    // --- REFACTOR 1: DRY ---
    // The original handleResearchModeRegenerate is now a simple wrapper.
    async function handleResearchModeRegenerate(apiKey, googleApiKey, query, messageDiv, conversationHistory) {
        // Just call the new core function
        await runResearchRequest(apiKey, googleApiKey, query, messageDiv, conversationHistory);
    }
    // --- END REFACTOR 1 ---


    function loadSettings() {
        const apiKey = localStorage.getItem('gemini_api_key');
        const googleApiKey = localStorage.getItem('google_api_key');
        if (apiKey) apiKeyInput.value = apiKey;
        if (googleApiKey) googleApiKeyInput.value = googleApiKey;
    }

    function showNotification(message, type = 'success') {
        const notification = document.createElement('div');
        notification.style.cssText = `position: fixed; top: 20px; right: 20px; padding: 1rem 1.5rem; border-radius: 12px; z-index: 9999; animation: slideUp 0.3s ease; background: ${type === 'error' ? 'var(--error-bg)' : 'var(--success-bg)'}; color: ${type === 'error' ? 'var(--error-text)' : 'var(--success-text)'}; border: 1px solid ${type === 'error' ? 'var(--error-text)' : 'var(--success-text)'};`;
        notification.textContent = message;
        document.body.appendChild(notification);
        setTimeout(() => {
            notification.remove();
        }, 3000);
    }
});