document.addEventListener('DOMContentLoaded', function() {
    // DOM Elements
    const conversationEl = document.getElementById('conversation');
    const promptInputEl = document.getElementById('promptInput');
    const actionBtnEl = document.getElementById('actionBtn');
    const downloadBtnEl = document.getElementById('downloadBtn');
    const notebookPreviewEl = document.getElementById('notebookPreview');
    const notebookTitleEl = document.getElementById('notebookTitle');
    const apiKeyInputEl = document.getElementById('apiKeyInput');
    const saveApiKeyBtnEl = document.getElementById('saveApiKeyBtn');
    const apiKeyFeedbackEl = document.getElementById('apiKeyFeedback');
    const modelSelectEl = document.getElementById('modelSelect');
    const generateModeBtnEl = document.getElementById('generateModeBtn');
    const editModeBtnEl = document.getElementById('editModeBtn');
    
    // Create scroll-to-bottom button
    const scrollToBottomBtn = document.createElement('button');
    scrollToBottomBtn.className = 'scroll-to-bottom';
    scrollToBottomBtn.innerHTML = '<i class="bi bi-arrow-down"></i>';
    scrollToBottomBtn.title = 'Scroll to bottom';
    document.body.appendChild(scrollToBottomBtn);
    
    // Scroll to bottom button event listener
    scrollToBottomBtn.addEventListener('click', function() {
        notebookPreviewEl.scrollTo({
            top: notebookPreviewEl.scrollHeight,
            behavior: 'smooth'
        });
    });
    
    // Show/hide scroll to bottom button based on scroll position
    notebookPreviewEl.addEventListener('scroll', function() {
        const scrollPosition = notebookPreviewEl.scrollTop + notebookPreviewEl.clientHeight;
        const scrollThreshold = notebookPreviewEl.scrollHeight - 100;
        
        if (scrollPosition < scrollThreshold) {
            scrollToBottomBtn.classList.add('visible');
        } else {
            scrollToBottomBtn.classList.remove('visible');
        }
    });
    
    // Global state
    let currentNotebook = null;
    let eventSource = null;
    let aiResponseText = '';
    let currentMode = 'generate'; // 'generate' or 'edit'
    
    // Add a debounce flag to prevent button spamming
    let isActionInProgress = false;
    
    // Initialize tooltips
    const tooltipTriggerList = [].slice.call(document.querySelectorAll('[data-bs-toggle="tooltip"]'));
    tooltipTriggerList.map(function(tooltipTriggerEl) {
        return new bootstrap.Tooltip(tooltipTriggerEl);
    });
    
    // Check if we have an API key in localStorage on page load
    const savedApiKey = localStorage.getItem('notegenie_api_key');
    if (savedApiKey) {
        // Try to set the API key automatically
        apiKeyInputEl.value = savedApiKey;
        checkAndSetApiKey();
    }
    
    // Event Listeners
    saveApiKeyBtnEl.addEventListener('click', saveApiKey);
    actionBtnEl.addEventListener('click', handleAction);
    downloadBtnEl.addEventListener('click', downloadNotebook);
    generateModeBtnEl.addEventListener('click', () => setMode('generate'));
    editModeBtnEl.addEventListener('click', () => setMode('edit'));
    
    // Allow Enter key to submit prompt (Shift+Enter for new line)
    promptInputEl.addEventListener('keydown', function(e) {
        if (e.key === 'Enter' && !e.shiftKey) {
            e.preventDefault();
            handleAction();
        }
    });
    
    // Set mode (generate or edit)
    function setMode(mode) {
        currentMode = mode;
        
        if (mode === 'generate') {
            generateModeBtnEl.classList.add('active');
            generateModeBtnEl.classList.remove('btn-outline-primary');
            generateModeBtnEl.classList.add('btn-primary');
            
            editModeBtnEl.classList.remove('active');
            editModeBtnEl.classList.remove('btn-primary');
            editModeBtnEl.classList.add('btn-outline-primary');
            
            promptInputEl.placeholder = 'Describe the notebook you want...';
        } else { // edit mode
            editModeBtnEl.classList.add('active');
            editModeBtnEl.classList.remove('btn-outline-primary');
            editModeBtnEl.classList.add('btn-primary');
            
            generateModeBtnEl.classList.remove('active');
            generateModeBtnEl.classList.remove('btn-primary');
            generateModeBtnEl.classList.add('btn-outline-primary');
            
            promptInputEl.placeholder = 'Describe what changes you want to make to the notebook...';
            promptInputEl.focus();
        }
    }
    
    // Handle action button click based on current mode
    function handleAction() {
        // Prevent multiple rapid clicks
        if (isActionInProgress) return;
        
        // Check if API key is set
        const apiKey = localStorage.getItem('notegenie_api_key');
        if (!apiKey) {
            addSystemMessage('<i class="bi bi-exclamation-triangle"></i> API key is not set. Please set your API key first.');
            showApiKeyModal();
            return;
        }
        
        // Get the prompt text and check if it's empty
        const prompt = promptInputEl.value.trim();
        
        if (currentMode === 'generate') {
            if (!prompt) {
                addSystemMessage('<i class="bi bi-exclamation-triangle"></i> Please enter a prompt to generate a notebook.');
                return;
            }
            
            // Set the debounce flag and disable the button
            isActionInProgress = true;
            actionBtnEl.disabled = true;
            
            generateNotebook();
        } else {
            if (!prompt || !currentNotebook) {
                if (!currentNotebook) {
                    addSystemMessage('<i class="bi bi-exclamation-triangle"></i> No notebook to edit. Generate a notebook first.');
                } else if (!prompt) {
                    addSystemMessage('<i class="bi bi-exclamation-triangle"></i> Please describe the changes you want to make.');
                }
                return;
            }
            
            // Set the debounce flag and disable the button
            isActionInProgress = true;
            actionBtnEl.disabled = true;
            
            editNotebook();
        }
    }
    
    // API Key handling
    function checkAndSetApiKey() {
        const apiKey = apiKeyInputEl.value.trim();
        
        if (!apiKey) {
            return;
        }
        
        fetch('/set_api_key', {
            method: 'POST',
            headers: {
                'Content-Type': 'application/x-www-form-urlencoded',
            },
            body: new URLSearchParams({
                'api_key': apiKey
            })
        })
        .then(response => response.json())
        .then(data => {
            if (data.success) {
                // Store API key in localStorage as backup
                localStorage.setItem('notegenie_api_key', apiKey);
                
                if (document.querySelector('#apiKeyModal.show')) {
                    showApiKeyFeedback('API key saved successfully!', 'success');
                    setTimeout(() => {
                        document.querySelector('#apiKeyModal .btn-close').click();
                    }, 1500);
                }
            } else {
                showApiKeyFeedback(`Error: ${data.message}`, 'danger');
                localStorage.removeItem('notegenie_api_key'); // Remove invalid key
            }
        })
        .catch(error => {
            showApiKeyFeedback(`Error: ${error.message}`, 'danger');
            localStorage.removeItem('notegenie_api_key'); // Remove key on error
        });
    }
    
    function saveApiKey() {
        checkAndSetApiKey();
    }
    
    function showApiKeyFeedback(message, type) {
        apiKeyFeedbackEl.innerHTML = `<div class="alert alert-${type} mb-0">${message}</div>`;
        setTimeout(() => {
            apiKeyFeedbackEl.innerHTML = '';
        }, 5000);
    }
    
    // Notebook generation
    function generateNotebook() {
        const prompt = promptInputEl.value.trim();
        
        if (!prompt) {
            addSystemMessage('<i class="bi bi-exclamation-triangle"></i> Please enter a prompt to generate a notebook.');
            return;
        }
        
        // Clear any existing notebook
        notebookPreviewEl.innerHTML = '';
        downloadBtnEl.disabled = true;
        editModeBtnEl.disabled = true;
        currentNotebook = null;
        
        // Add user message to conversation
        addUserMessage(prompt);
        
        // Disable input during generation
        setGeneratingState(true);
        
        const modelName = modelSelectEl.value;
        
        // Add AI typing indicator
        const aiMessageId = 'ai-typing-' + Date.now();
        addTypingIndicator(aiMessageId);
        
        // Always use streaming for better UX
        handleStreamingResponse(prompt, modelName, aiMessageId);
        
        // Clear the input field after sending
        promptInputEl.value = '';
    }
    
    // Notebook editing
    function editNotebook() {
        const editPrompt = promptInputEl.value.trim();
        
        if (!editPrompt || !currentNotebook) {
            if (!currentNotebook) {
                addSystemMessage('<i class="bi bi-exclamation-triangle"></i> No notebook to edit. Generate a notebook first.');
            } else {
                addSystemMessage('<i class="bi bi-exclamation-triangle"></i> Please describe the changes you want to make.');
            }
            return;
        }
        
        // Add user message to conversation
        addUserMessage('Edit request: ' + editPrompt);
        
        // Disable input during editing
        setGeneratingState(true);
        
        const modelName = modelSelectEl.value;
        
        // Add AI typing indicator
        const aiMessageId = 'ai-typing-' + Date.now();
        addTypingIndicator(aiMessageId);
        
        // Update AI message to show we're starting the edit
        updateAiMessage(aiMessageId, "**NoteGenie:** Starting to edit your notebook...");
        
        handleEditStreamingResponse(editPrompt, currentNotebook, modelName, aiMessageId);
        
        // Clear the input field after sending
        promptInputEl.value = '';
    }
    
    // Create a custom EventSource that supports headers 
    function createEventSourceWithHeaders(url, headers) {
        if (typeof EventSourcePolyfill !== 'undefined') {
            return new EventSourcePolyfill(url, { headers: headers });
        }
        
        // If native EventSource is all we have but we need headers, 
        // we fall back to using URL parameters for authentication
        return new EventSource(url);
    }

    function handleStreamingResponse(prompt, modelName, aiMessageId) {
        aiResponseText = '';
        
        // Clear any existing notebook preview and add loading indicator
        notebookPreviewEl.innerHTML = '<div class="loading-preview-message text-center p-5"><div class="spinner-border text-primary" role="status"></div><p class="mt-3">Building notebook preview...</p></div>';
        
        // Close any existing event source
        if (eventSource) {
            eventSource.close();
        }
        
        // Keep track of when we should update the preview (not on every tiny chunk)
        let lastPreviewUpdate = 0;
        const PREVIEW_UPDATE_INTERVAL = 1000; // Update preview every 1 second during streaming
        
        // Track the last time we got a chunk - for timeout detection
        let lastChunkTime = Date.now();
        const CONNECTION_TIMEOUT = 30000; // 30 seconds without data = timeout
        
        // Start a monitoring timer to detect stalled connections
        const connectionTimer = setInterval(() => {
            if (Date.now() - lastChunkTime > CONNECTION_TIMEOUT) {
                clearInterval(connectionTimer);
                if (eventSource) {
                    console.log("Connection timed out - closing event source");
                    eventSource.close();
                    eventSource = null;
                    
                    // Process what we've got so far
                    updateAiMessage(aiMessageId, "**NoteGenie:** Generation timed out, but I'll process what I've received so far.");
                    processNotebookResponse(aiResponseText);
                    setGeneratingState(false);
                }
            }
        }, 5000); // Check every 5 seconds
        
        // Get API key from localStorage as a backup
        const backupApiKey = localStorage.getItem('notegenie_api_key');
        
        // Create URL parameters including API key as fallback
        const urlParams = new URLSearchParams({
            prompt: prompt,
            model: modelName,
            stream: true
        });
        
        // Add API key to URL params if available from localStorage (as a fallback)
        if (backupApiKey) {
            urlParams.append('api_key', backupApiKey);
        }
        
        // Set headers with API key if available
        const headers = {};
        if (backupApiKey) {
            headers['X-API-Key'] = backupApiKey;
        }
        
        // Create a new event source with API key included in both URL and headers
        eventSource = createEventSourceWithHeaders(`/generate_notebook?${urlParams.toString()}`, headers);
        
        eventSource.onmessage = function(event) {
            // Update our last-activity timestamp
            lastChunkTime = Date.now();
            
            try {
                const data = JSON.parse(event.data);
                
                // Handle errors sent from the server
                if (data.error) {
                    console.error("Server error:", data.error);
                    updateAiMessage(aiMessageId, `**Error:** ${data.error}`);
                    
                    // Try to salvage what we have so far
                    if (aiResponseText && (aiResponseText.includes('MARKDOWN CELL') || aiResponseText.includes('CODE CELL'))) {
                        processNotebookResponse(aiResponseText);
                    }
                    
                    eventSource.close();
                    eventSource = null;
                    setGeneratingState(false);
                    clearInterval(connectionTimer);
                    return;
                }
                
                if (data.chunk) {
                    aiResponseText += data.chunk;
                    
                    // Extract notebook info as soon as it's available
                    const nameMatch = aiResponseText.match(/NOTEBOOK_NAME:?\s*(.+?)(?:\n|$)/);
                    const descMatch = aiResponseText.match(/NOTEBOOK_DESCRIPTION:?\s*(.+?)(?:\n|$)/);
                    
                    if (nameMatch && nameMatch[1].trim()) {
                        // Update notebook title immediately when found
                        notebookTitleEl.textContent = nameMatch[1].trim();
                    }
                    
                    if (descMatch && descMatch[1].trim()) {
                        // Update AI message to only display the description, not the full response
                        updateAiMessage(aiMessageId, `**NoteGenie:** ${descMatch[1].trim()}`);
                    } else {
                        // Show a simple generating message while waiting for description
                        updateAiMessage(aiMessageId, "**NoteGenie:** Generating your notebook...");
                    }
                    
                    // Update preview periodically during streaming
                    const now = Date.now();
                    if (now - lastPreviewUpdate > PREVIEW_UPDATE_INTERVAL) {
                        lastPreviewUpdate = now;
                        
                        // Only try to update preview if we have meaningful content
                        if (aiResponseText.includes('MARKDOWN CELL') || aiResponseText.includes('CODE CELL')) {
                            updateNotebookPreviewDuringStream(aiResponseText);
                        }
                    }
                }
                
                if (data.done) {
                    eventSource.close();
                    eventSource = null;
                    clearInterval(connectionTimer);
                    
                    // Process the complete response for final rendering
                    processNotebookResponse(aiResponseText);
                    setGeneratingState(false);
                }
            } catch (error) {
                console.error('Error parsing event data:', error, event.data);
            }
        };
        
        eventSource.onerror = function(err) {
            console.error('EventSource error:', err);
            eventSource.close();
            eventSource = null;
            clearInterval(connectionTimer);
            
            // Check if it's an auth error (most likely API key not set)
            if (err.status === 401) {
                updateAiMessage(aiMessageId, '**Error: API key not set or invalid.** \n\nPlease click the API Key button in the top right corner to set your Google Gemini API key.');
                showApiKeyModal();
            } else {
                // Try to salvage what we have so far
                if (aiResponseText && (aiResponseText.includes('MARKDOWN CELL') || aiResponseText.includes('CODE CELL'))) {
                    updateAiMessage(aiMessageId, '**Warning:** Connection issue occurred but I\'ll try to process what I received so far.');
                    processNotebookResponse(aiResponseText);
                } else {
                    updateAiMessage(aiMessageId, '**Error:** Failed to generate notebook. Please try again.');
                }
            }
            
            setGeneratingState(false);
        };
    }
    
    function handleEditStreamingResponse(editPrompt, notebook, modelName, aiMessageId) {
        aiResponseText = '';
        
        // Clear any existing notebook preview and add loading indicator
        notebookPreviewEl.innerHTML = '<div class="loading-preview-message text-center p-5"><div class="spinner-border text-primary" role="status"></div><p class="mt-3">Updating notebook based on your edit request...</p></div>';
        
        // Close any existing event source
        if (eventSource) {
            eventSource.close();
            eventSource = null;
        }
        
        // Keep track of when we should update the preview (not on every tiny chunk)
        let lastPreviewUpdate = 0;
        const PREVIEW_UPDATE_INTERVAL = 1000; // Update preview every 1 second during streaming, same as generate
        
        // Get API key from localStorage as a backup
        const backupApiKey = localStorage.getItem('notegenie_api_key');
        
        // Make sure notebook is in proper format before sending
        const notebookToSend = typeof notebook === 'string' ? JSON.parse(notebook) : notebook;
        
        // Debug logging to help diagnose issues
        console.log('Notebook edit request:', {
            cellCount: notebookToSend.cells.length,
            promptLength: editPrompt.length,
            model: modelName
        });
        
        // First, send the notebook data to the server so it's available in the session
        fetch('/prepare_edit_notebook', {
            method: 'POST',
            headers: {
                'Content-Type': 'application/json',
                // Add API key as header if available
                ...(backupApiKey && {'X-API-Key': backupApiKey})
            },
            body: JSON.stringify({
                notebook: notebookToSend
            })
        })
        .then(response => {
            if (!response.ok) {
                throw new Error(`Server error (${response.status}): ${response.statusText}`);
            }
            return response.json();
        })
        .then(data => {
            if (data.success) {
                updateAiMessage(aiMessageId, "**NoteGenie:** Processing your edit request...");
                
                // Create URL parameters including API key as fallback
                const urlParams = new URLSearchParams({
                    edit_prompt: editPrompt,
                    model: modelName,
                    stream: true
                });
                
                // Add API key to URL params if available from localStorage (as a fallback)
                if (backupApiKey) {
                    urlParams.append('api_key', backupApiKey);
                }
                
                // Set headers with API key if available
                const headers = {};
                if (backupApiKey) {
                    headers['X-API-Key'] = backupApiKey;
                }
                
                console.log('Starting EventSource for edit with headers:', Object.keys(headers).length > 0 ? 'Headers present' : 'No headers');
                
                // Now create a new event source for editing with the API key included in both URL and headers
                try {
                    eventSource = createEventSourceWithHeaders(`/edit_notebook?${urlParams.toString()}`, headers);
                    
                    // Track the last time we got a chunk - for timeout detection
                    let lastChunkTime = Date.now();
                    const CONNECTION_TIMEOUT = 30000; // 30 seconds without data = timeout
                    
                    // Start a monitoring timer to detect stalled connections
                    const connectionTimer = setInterval(() => {
                        if (Date.now() - lastChunkTime > CONNECTION_TIMEOUT) {
                            clearInterval(connectionTimer);
                            if (eventSource) {
                                console.log("Connection timed out - closing event source");
                                eventSource.close();
                                eventSource = null;
                                
                                // Process what we've got so far
                                updateAiMessage(aiMessageId, "**NoteGenie:** Edit processing timed out, but I'll use what I've received so far.");
                                processNotebookResponse(aiResponseText);
                                setGeneratingState(false);
                                
                                // Switch back to generate mode
                                setMode('generate');
                            }
                        }
                    }, 5000); // Check every 5 seconds
                    
                    eventSource.onopen = function(event) {
                        console.log('EventSource connection opened for edit');
                    };
                    
                    eventSource.onmessage = function(event) {
                        // Update our last-activity timestamp
                        lastChunkTime = Date.now();
                        
                        try {
                            const data = JSON.parse(event.data);
                            
                            // Handle errors sent from the server
                            if (data.error) {
                                console.error("Server error during edit:", data.error);
                                updateAiMessage(aiMessageId, `**Error:** ${data.error}`);
                                
                                // Try to salvage what we have so far
                                if (aiResponseText && (aiResponseText.includes('MARKDOWN CELL') || aiResponseText.includes('CODE CELL'))) {
                                    processNotebookResponse(aiResponseText);
                                }
                                
                                eventSource.close();
                                eventSource = null;
                                setGeneratingState(false);
                                clearInterval(connectionTimer);
                                
                                // Switch back to generate mode
                                setMode('generate');
                                return;
                            }
                            
                            if (data.chunk) {
                                aiResponseText += data.chunk;
                                
                                // Extract notebook info as soon as it's available, similar to generate function
                                const nameMatch = aiResponseText.match(/NOTEBOOK_NAME:?\s*(.+?)(?:\n|$)/);
                                const descMatch = aiResponseText.match(/NOTEBOOK_DESCRIPTION:?\s*(.+?)(?:\n|$)/);
                                
                                if (nameMatch && nameMatch[1].trim()) {
                                    // Update notebook title immediately when found
                                    notebookTitleEl.textContent = nameMatch[1].trim();
                                }
                                
                                if (descMatch && descMatch[1].trim()) {
                                    // Update AI message with the actual description from the edited notebook
                                    updateAiMessage(aiMessageId, `**NoteGenie:** ${descMatch[1].trim()}`);
                                } else {
                                    // Show a simple editing message while waiting for description
                                    updateAiMessage(aiMessageId, "**NoteGenie:** Updating notebook based on your edit request...");
                                }
                                
                                // Update preview periodically during streaming, just like in generate function
                                const now = Date.now();
                                if (now - lastPreviewUpdate > PREVIEW_UPDATE_INTERVAL) {
                                    lastPreviewUpdate = now;
                                    
                                    // Only try to update preview if we have meaningful content
                                    if (aiResponseText.includes('MARKDOWN CELL') || aiResponseText.includes('CODE CELL')) {
                                        updateNotebookPreviewDuringStream(aiResponseText);
                                    }
                                }
                            }
                            
                            if (data.done) {
                                console.log('Edit completed successfully');
                                eventSource.close();
                                eventSource = null;
                                clearInterval(connectionTimer);
                                
                                // Process the complete response for final rendering
                                processNotebookResponse(aiResponseText);
                                setGeneratingState(false);
                                
                                // Get the final description for the AI message
                                const finalDescMatch = aiResponseText.match(/NOTEBOOK_DESCRIPTION:?\s*(.+?)(?=\n\s*---|\n\s*$|$)/);
                                const finalDesc = finalDescMatch ? finalDescMatch[1].trim() : "I've updated the notebook based on your edit request.";
                                
                                // Update AI message with final description
                                updateAiMessage(aiMessageId, `**NoteGenie:** ${finalDesc}`);
                                
                                // Switch back to generate mode
                                setMode('generate');
                            }
                        } catch (error) {
                            console.error('Error parsing event data:', error, event.data);
                        }
                    };
                    
                    eventSource.onerror = function(err) {
                        // More detailed error logging
                        console.error('EventSource error during edit:', err);
                        
                        // Get any status text or error message if available
                        let errorMessage = 'Failed to connect to edit service';
                        if (err instanceof Event && err.target && err.target.status) {
                            errorMessage += ` (Status: ${err.target.status})`;
                        }
                        
                        // Close and clean up event source
                        if (eventSource) {
                            eventSource.close();
                            eventSource = null;
                        }
                        
                        // Clear any timers
                        clearInterval(connectionTimer);
                        
                        // Try to salvage what we have so far
                        if (aiResponseText && (aiResponseText.includes('MARKDOWN CELL') || aiResponseText.includes('CODE CELL'))) {
                            updateAiMessage(aiMessageId, '**Warning:** Connection issue occurred but I\'ll try to process what I received so far.');
                            processNotebookResponse(aiResponseText);
                        } else {
                            updateAiMessage(aiMessageId, `**Error:** Failed to edit notebook. ${errorMessage}`);
                        }
                        
                        setGeneratingState(false);
                        
                        // Switch back to generate mode
                        setMode('generate');
                    };
                } catch (e) {
                    console.error('Error creating EventSource:', e);
                    updateAiMessage(aiMessageId, `**Error:** Could not establish connection for editing: ${e.message}`);
                    setGeneratingState(false);
                }
            } else {
                throw new Error(data.message || 'Failed to prepare notebook for editing');
            }
        })
        .catch(error => {
            console.error('Error preparing notebook for edit:', error);
            updateAiMessage(aiMessageId, '**Error:** ' + error.message);
            notebookPreviewEl.innerHTML = '<div class="alert alert-danger">Error: Failed to prepare notebook for editing. Please try regenerating the notebook.</div>';
            setGeneratingState(false);
        });
    }
    
    // Show API key modal
    function showApiKeyModal() {
        const apiKeyModal = new bootstrap.Modal(document.getElementById('apiKeyModal'));
        apiKeyModal.show();
    }
    
    // Helper functions
    function updateNotebookPreviewDuringStream(text) {
        try {
            // Extract notebook info for title
            const nameMatch = text.match(/NOTEBOOK_NAME:?\s*(.+?)(?:\n|$)/);
            const name = nameMatch ? nameMatch[1].trim() : 'Generating Notebook...';
            
            // Update notebook title
            notebookTitleEl.textContent = name;
            
            // Find all cell markers in order of appearance
            const cellMarkers = [...text.matchAll(/---\s*(MARKDOWN|CODE)\s*CELL\s*---/g)];
            
            // If no markers found, don't update preview yet
            if (cellMarkers.length === 0) return;
            
            // Check if user is already near the bottom before deciding to auto-scroll
            const isNearBottom = notebookPreviewEl.scrollHeight - notebookPreviewEl.scrollTop - notebookPreviewEl.clientHeight < 100;
            
            // Initialize cell counters
            let markdownCellCount = 0;
            let codeCellCount = 0;
            
            // Clear preview and prepare for new content
            notebookPreviewEl.innerHTML = '';
            
            // Process each cell in order
            for (let i = 0; i < cellMarkers.length; i++) {
                const currentMarker = cellMarkers[i];
                const nextMarker = cellMarkers[i + 1];
                const cellType = currentMarker[1]; // MARKDOWN or CODE
                
                // Extract cell content between current marker and next marker (or end of text)
                const markerEndIndex = currentMarker.index + currentMarker[0].length;
                const contentEndIndex = nextMarker ? nextMarker.index : text.length;
                let cellContent = text.substring(markerEndIndex, contentEndIndex).trim();
                
                if (cellType === 'MARKDOWN') {
                    markdownCellCount++;
                    createMarkdownCell(cellContent, markdownCellCount, notebookPreviewEl);
                } else if (cellType === 'CODE') {
                    // Extract Python code from between ```python and ```
                    const codeMatch = cellContent.match(/```python\s*([\s\S]*?)```/);
                    if (codeMatch) {
                        codeCellCount++;
                        createCodeCell(codeMatch[1].trim(), codeCellCount, notebookPreviewEl);
                    }
                }
            }
            
            // Only auto-scroll if user was already near the bottom
            if (isNearBottom) {
                notebookPreviewEl.scrollTop = notebookPreviewEl.scrollHeight;
            } else {
                // Show the scroll to bottom button
                scrollToBottomBtn.classList.add('visible');
            }
        } catch (error) {
            console.error('Error updating preview during stream:', error);
            // Don't update if there's an error - wait for complete response
        }
    }
    
    function createMarkdownCell(content, cellNumber, container) {
        const cellDiv = document.createElement('div');
        cellDiv.className = 'notebook-cell cell-markdown';
        
        // Add header for markdown cell
        const headerDiv = document.createElement('div');
        headerDiv.className = 'cell-header markdown-header';
        headerDiv.innerHTML = `<span class="cell-type">Markdown [${cellNumber}]</span>`;
        cellDiv.appendChild(headerDiv);
        
        // Add markdown content
        const contentDiv = document.createElement('div');
        contentDiv.className = 'cell-content';
        contentDiv.innerHTML = marked.parse(content);
        cellDiv.appendChild(contentDiv);
        
        container.appendChild(cellDiv);
    }
    
    function createCodeCell(content, cellNumber, container) {
        const cellDiv = document.createElement('div');
        cellDiv.className = 'notebook-cell cell-code';
        
        // Add header for code cell
        const headerDiv = document.createElement('div');
        headerDiv.className = 'cell-header code-header';
        headerDiv.innerHTML = `<span class="cell-type">Code [${cellNumber}]</span>`;
        cellDiv.appendChild(headerDiv);
        
        // Add code content
        const contentDiv = document.createElement('div');
        contentDiv.className = 'cell-content';
        const preEl = document.createElement('pre');
        const codeEl = document.createElement('code');
        codeEl.className = 'language-python';
        codeEl.textContent = content;
        preEl.appendChild(codeEl);
        contentDiv.appendChild(preEl);
        cellDiv.appendChild(contentDiv);
        
        container.appendChild(cellDiv);
        
        // Highlight syntax
        Prism.highlightElement(codeEl);
    }
    
    function processNotebookResponse(text) {
        // Extract notebook info using improved regex patterns that handle markdown better
        const nameMatch = text.match(/NOTEBOOK_NAME:?\s*(.+?)(?=\n\s*NOTEBOOK_DESCRIPTION|\n\s*---|\n\s*$|$)/);
        const descMatch = text.match(/NOTEBOOK_DESCRIPTION:?\s*(.+?)(?=\n\s*---|\n\s*$|$)/);
        
        // Clean up the name and description
        let name = nameMatch ? nameMatch[1].trim() : 'Generated Notebook';
        let description = descMatch ? descMatch[1].trim() : '';
        
        // Remove markdown formatting
        name = name.replace(/\*\*/g, '').replace(/\*/g, '');
        description = description.replace(/\*\*/g, '').replace(/\*/g, '');
        
        // Update notebook title
        notebookTitleEl.textContent = name;
        
        // Update the AI message with just the description
        const chatMessageId = document.querySelector('.ai-message').id;
        updateAiMessage(chatMessageId, `**NoteGenie:** ${description}`);
        
        // OPTIMIZATION: Process the notebook client-side if possible for small to medium notebooks
        if (text.length < 50000) { // Only try client-side processing for reasonably sized notebooks
            try {
                // Try to process notebook directly in the client
                const notebookJson = clientSideFormatNotebook(text);
                renderNotebook(notebookJson);
                currentNotebook = notebookJson;
                downloadBtnEl.disabled = false;
                editModeBtnEl.disabled = false;
                return; // Exit if successful
            } catch (error) {
                console.log("Client-side formatting failed, falling back to server:", error);
                // Fall through to server-side processing if client-side fails
            }
        }
        
        // OPTIMIZATION: Show loading indicator with progress message
        notebookPreviewEl.innerHTML = '<div class="loading-preview-message text-center p-5"><div class="spinner-border text-primary" role="status"></div><p class="mt-3">Finalizing notebook format...</p></div>';
        
        // Make API call to convert text to notebook format with a timeout to prevent UI freezing
        setTimeout(() => {
            fetch('/generate_notebook', {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                },
                body: JSON.stringify({
                    prompt: text,  // Send the full AI response as a prompt
                    model: 'gemini-2.0-flash',  // Use a faster model for formatting
                    stream: false,
                    format_only: true
                })
            })
            .then(response => response.json())
            .then(data => {
                if (data.success) {
                    renderNotebook(data.notebook);
                    currentNotebook = data.notebook;
                    downloadBtnEl.disabled = false;
                    editModeBtnEl.disabled = false;
                }
            })
            .catch(error => {
                console.error('Error formatting notebook:', error);
                notebookPreviewEl.innerHTML = `<div class="alert alert-danger">Error formatting notebook: ${error.message}</div>`;
            });
        }, 10); // Small delay to allow UI to update
    }
    
    // ADDED: Client-side notebook formatting function to reduce server dependency
    function clientSideFormatNotebook(content) {
        // Extract cells with improved regex that captures end of file properly
        const markdownCells = content.match(/---\s*MARKDOWN\s*CELL\s*---\s*([\s\S]*?)(?=---\s*(?:MARKDOWN|CODE)\s*CELL\s*---|$)/g) || [];
        const codeCells = content.match(/---\s*CODE\s*CELL\s*---\s*```python\s*([\s\S]*?)```/gs) || [];
        
        // Get cell order with improved markers
        const cellMarkers = [...content.matchAll(/---\s*(MARKDOWN|CODE)\s*CELL\s*---/g)];
        const cellTypes = cellMarkers.map(match => match[1]);
        
        // Initialize notebook structure
        const cells = [];
        let mdIdx = 0;
        let codeIdx = 0;
        
        // Populate cells in the correct order
        for (let i = 0; i < cellTypes.length; i++) {
            const cellType = cellTypes[i];
            
            if (cellType === 'MARKDOWN' && mdIdx < markdownCells.length) {
                // Extract markdown content with improved regex
                let md = markdownCells[mdIdx].replace(/---\s*MARKDOWN\s*CELL\s*---/, '').trim();
                cells.push({
                    cell_type: 'markdown',
                    metadata: {},
                    source: md.split('\n')
                });
                mdIdx++;
            } else if (cellType === 'CODE' && codeIdx < codeCells.length) {
                // Extract code content with improved regex
                let codeMatch = codeCells[codeIdx].match(/```python\s*([\s\S]*?)```/s);
                let code = codeMatch ? codeMatch[1].trim() : '';
                cells.push({
                    cell_type: 'code',
                    execution_count: null,
                    metadata: {},
                    outputs: [],
                    source: code.split('\n')
                });
                codeIdx++;
            }
        }
        
        // Process the final cell if it wasn't captured above
        if (cellMarkers.length > 0) {
            const lastMarker = cellMarkers[cellMarkers.length - 1];
            const lastMarkerEndIndex = lastMarker.index + lastMarker[0].length;
            
            // If there's content after the last marker, process it
            if (lastMarkerEndIndex < content.length) {
                const lastCellType = lastMarker[1]; // MARKDOWN or CODE
                const lastCellContent = content.substring(lastMarkerEndIndex).trim();
                
                // Only add if not already captured
                const isAlreadyCaptured = 
                    (lastCellType === 'MARKDOWN' && mdIdx >= markdownCells.length) ||
                    (lastCellType === 'CODE' && codeIdx >= codeCells.length);
                    
                if (!isAlreadyCaptured && lastCellContent) {
                    if (lastCellType === 'MARKDOWN') {
                        cells.push({
                            cell_type: 'markdown',
                            metadata: {},
                            source: lastCellContent.split('\n')
                        });
                    } else if (lastCellType === 'CODE') {
                        // Extract code if present
                        const codeMatch = lastCellContent.match(/```python\s*([\s\S]*?)```/s);
                        if (codeMatch) {
                            cells.push({
                                cell_type: 'code',
                                execution_count: null,
                                metadata: {},
                                outputs: [],
                                source: codeMatch[1].trim().split('\n')
                            });
                        }
                    }
                }
            }
        }
        
        // Fallback approach if no cells were found or if the extraction didn't work
        // ...existing code...
        
        return {
            cells: cells,
            metadata: {
                kernelspec: {
                    display_name: 'Python 3',
                    language: 'python',
                    name: 'python3'
                },
                language_info: {
                    name: 'python',
                    version: '3.8.0'
                }
            },
            nbformat: 4,
            nbformat_minor: 4
        };
    }
    
    function downloadNotebook() {
        if (!currentNotebook) {
            return;
        }
        
        // Improved filename cleaning - remove markdown formatting and unsafe filename characters
        let cleanName = notebookTitleEl.textContent || 'generated_notebook';
        // First remove markdown formatting characters
        cleanName = cleanName.replace(/\*\*/g, '').replace(/\*/g, '').replace(/_/g, ' ');
        // Then remove any characters that aren't safe for filenames
        cleanName = cleanName.replace(/[^\w\s-]/gi, '').trim();
        // Replace multiple spaces with single space
        cleanName = cleanName.replace(/\s+/g, ' ');
        
        const filename = cleanName || 'generated_notebook';
        
        fetch('/download_notebook', {
            method: 'POST',
            headers: {
                'Content-Type': 'application/json',
            },
            body: JSON.stringify({
                notebook: currentNotebook,
                filename: `${filename}.ipynb`
            })
        })
        .then(response => response.blob())
        .then(blob => {
            const url = window.URL.createObjectURL(blob);
            const a = document.createElement('a');
            a.style.display = 'none';
            a.href = url;
            a.download = `${filename}.ipynb`;
            document.body.appendChild(a);
            a.click();
            window.URL.revokeObjectURL(url);
        })
        .catch(error => {
            console.error('Error downloading notebook:', error);
            addSystemMessage('Error downloading notebook: ' + error.message);
        });
    }
    
    function renderNotebook(notebook) {
        notebookPreviewEl.innerHTML = '';
        
        if (!notebook || !notebook.cells || !notebook.cells.length) {
            notebookPreviewEl.innerHTML = '<div class="alert alert-warning">No notebook content available</div>';
            return;
        }
        
        let cellCount = 0;
        
        notebook.cells.forEach((cell, index) => {
            cellCount++;
            
            if (cell.cell_type === 'markdown') {
                const content = Array.isArray(cell.source) ? cell.source.join('\n') : cell.source;
                createMarkdownCell(content, cellCount, notebookPreviewEl);
            } else if (cell.cell_type === 'code') {
                const codeContent = Array.isArray(cell.source) ? cell.source.join('\n') : cell.source;
                createCodeCell(codeContent, cellCount, notebookPreviewEl);
            }
        });
        
        // Check if we should scroll to the bottom after rendering
        const isNearBottom = notebookPreviewEl.scrollHeight - notebookPreviewEl.scrollTop - notebookPreviewEl.clientHeight < 100;
        if (isNearBottom) {
            notebookPreviewEl.scrollTop = notebookPreviewEl.scrollHeight;
        } else {
            // Show the scroll to bottom button
            scrollToBottomBtn.classList.add('visible');
        }
        
        // Enable edit and download buttons when notebook is available
        downloadBtnEl.disabled = false;
        editModeBtnEl.disabled = false;
    }
    
    // UI Message handling
    function addUserMessage(text) {
        // Remove welcome message if present
        const welcomeMessage = document.querySelector('.welcome-message');
        if (welcomeMessage) {
            welcomeMessage.remove();
        }
        
        const messageDiv = document.createElement('div');
        messageDiv.className = 'message user-message';
        
        // Fix: Create a span for the "You:" prefix and handle the text separately
        // to prevent marked from adding paragraph tags that cause line breaks
        const userPrefix = document.createElement('span');
        userPrefix.className = 'user-prefix';
        userPrefix.innerHTML = '<strong>You:</strong> ';
        
        const messageContent = document.createElement('span');
        messageContent.className = 'message-content';
        messageContent.innerHTML = marked.parse(text);
        
        // Remove any leading <p> and trailing </p> tags that marked adds
        messageContent.innerHTML = messageContent.innerHTML
            .replace(/^<p>/, '')
            .replace(/<\/p>$/, '');
        
        messageDiv.appendChild(userPrefix);
        messageDiv.appendChild(messageContent);
        
        conversationEl.appendChild(messageDiv);
        conversationEl.scrollTop = conversationEl.scrollHeight;
    }
    
    function addTypingIndicator(id) {
        const messageDiv = document.createElement('div');
        messageDiv.className = 'message ai-message';
        messageDiv.id = id;
        
        const typingDiv = document.createElement('div');
        typingDiv.className = 'typing-indicator';
        for (let i = 0; i < 3; i++) {
            const dot = document.createElement('span');
            dot.className = 'typing-dot';
            typingDiv.appendChild(dot);
        }
        
        messageDiv.appendChild(typingDiv);
        conversationEl.appendChild(messageDiv);
        conversationEl.scrollTop = conversationEl.scrollHeight;
    }
    
    function updateAiMessage(id, text) {
        const messageDiv = document.getElementById(id);
        if (messageDiv) {
            if (text.startsWith('**NoteGenie:**')) {
                text = text.replace('**NoteGenie:**', '**NoteGenie:**');
            } else if (text.startsWith('**Error:**')) {
                text = text.replace('**Error:**', '**<i class="bi bi-exclamation-triangle"></i> Error:**');
            }
            
            messageDiv.innerHTML = marked.parse(text);
            // Ensure auto-scrolling in the chat panel
            conversationEl.scrollTop = conversationEl.scrollHeight;
        }
    }
    
    function addSystemMessage(text) {
        const messageDiv = document.createElement('div');
        messageDiv.className = 'info-card mb-3';
        messageDiv.innerHTML = text;
        conversationEl.appendChild(messageDiv);
        conversationEl.scrollTop = conversationEl.scrollHeight;
    }
    
    // UI state management
    function setGeneratingState(isGenerating) {
        if (isGenerating) {
            actionBtnEl.disabled = true;
            // Update loading state for the send button
            actionBtnEl.innerHTML = '<span class="material-icons spinning">autorenew</span>';
            promptInputEl.disabled = true;
            
            // Disable buttons but preserve their active/inactive visual state
            generateModeBtnEl.disabled = true;
            editModeBtnEl.disabled = true;
            
            // Make sure mode toggle visual state is preserved
            if (currentMode === 'generate') {
                generateModeBtnEl.classList.add('active');
                editModeBtnEl.classList.remove('active');
            } else {
                editModeBtnEl.classList.add('active');
                generateModeBtnEl.classList.remove('active');
            }
        } else {
            // Reset the debounce flag when generation completes
            isActionInProgress = false;
            
            actionBtnEl.disabled = false;
            // Always use the send icon for the inline button
            actionBtnEl.innerHTML = '<span class="material-icons">send</span>';
            promptInputEl.disabled = false;
            promptInputEl.focus();
            
            generateModeBtnEl.disabled = false;
            
            // Only enable edit mode button if we have a current notebook
            editModeBtnEl.disabled = !currentNotebook;
        }
    }
});