frontend/script.js

// API Configuration - Use relative URL for HuggingFace
const API_BASE_URL = window.location.origin;

// Session Management
let currentSessionId = null;
let conversationHistory = []; // Store all Q&A pairs
let queryHistory = [];

// DOM Elements
const queryInput = document.getElementById('queryInput');
const submitBtn = document.getElementById('submitBtn');
const btnText = submitBtn.querySelector('.btn-text');
const loader = submitBtn.querySelector('.loader');
const responseSection = document.getElementById('responseSection');
const initialState = document.getElementById('initialState');
const loadingState = document.getElementById('loadingState');
const answerDiv = document.getElementById('answer');
const sourcesDiv = document.getElementById('sources');
const errorToast = document.getElementById('errorToast');
const errorMessage = document.getElementById('errorMessage');
const expandedQueriesDiv = document.getElementById('expandedQueries');
const queriesList = document.getElementById('queriesList');
const exportBtn = document.getElementById('exportBtn');
const githubBtn = document.getElementById('githubBtn');
const historySelect = document.getElementById('historySelect');

// Initialize
function init() {
    currentSessionId = sessionStorage.getItem('sessionId') || generateSessionId();
    sessionStorage.setItem('sessionId', currentSessionId);
    
    // Load history
    const saved = sessionStorage.getItem('queryHistory');
    if (saved) {
        try {
            queryHistory = JSON.parse(saved);
            updateHistoryDropdown();
        } catch (e) {
            queryHistory = [];
        }
    }
    
    setupEventListeners();
    checkHealth();
}

function generateSessionId() {
    return `session_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
}

function setupEventListeners() {
    // Submit button
    if (submitBtn) {
        submitBtn.addEventListener('click', handleSubmit);
    }
    
    // Enter key (Ctrl+Enter)
    if (queryInput) {
        queryInput.addEventListener('keydown', (e) => {
            if (e.key === 'Enter' && e.ctrlKey) {
                handleSubmit();
            }
        });
    }
    
    // Example queries
    document.querySelectorAll('.example-item').forEach(item => {
        item.addEventListener('click', () => {
            const query = item.getAttribute('data-query');
            if (query && queryInput) {
                queryInput.value = query;
                handleSubmit();
            }
        });
    });
    
    // Export button
    if (exportBtn) {
        exportBtn.addEventListener('click', exportToPDF);
    }
    
    // GitHub button
    if (githubBtn) {
        githubBtn.addEventListener('click', () => {
            window.open('https://github.com/Adeyemi0/FinSight-RAG-Application-', '_blank');
        });
    }
    
    // History dropdown
    if (historySelect) {
        historySelect.addEventListener('change', (e) => {
            const query = e.target.value;
            if (query && queryInput) {
                queryInput.value = query;
            }
        });
    }
}

// Update history dropdown
function updateHistoryDropdown() {
    if (!historySelect) return;
    
    historySelect.innerHTML = '<option value="">Select a recent query</option>';
    
    queryHistory.slice().reverse().forEach((item, index) => {
        const option = document.createElement('option');
        option.value = item.query;
        option.textContent = item.query.substring(0, 60) + (item.query.length > 60 ? '...' : '');
        historySelect.appendChild(option);
    });
}

// Save to history
function saveToHistory(query, answer) {
    queryHistory.push({
        query,
        answer: answer.substring(0, 500),
        timestamp: new Date().toISOString()
    });
    
    // Keep last 20
    if (queryHistory.length > 20) {
        queryHistory = queryHistory.slice(-20);
    }
    
    try {
        sessionStorage.setItem('queryHistory', JSON.stringify(queryHistory));
        updateHistoryDropdown();
    } catch (e) {
        console.error('Failed to save history:', e);
    }
}

// Main submit handler
async function handleSubmit() {
    if (!queryInput) return;
    
    const query = queryInput.value.trim();
    
    if (!query) {
        showError('Please enter a question');
        return;
    }
    
    // Show loading
    setLoading(true);
    hideError();
    
    // Show loading indicator
    if (loadingState) loadingState.style.display = 'flex';
    
    // Hide initial state on first query
    if (initialState && conversationHistory.length === 0) {
        initialState.style.display = 'none';
    }
    
    // Show response section
    if (responseSection) responseSection.style.display = 'block';
    
    try {
        // Extract ticker from query if mentioned
        const tickerMatch = query.match(/\b([A-Z]{1,5})\b/);
        const detectedTicker = tickerMatch ? tickerMatch[1] : null;
        
        const requestData = {
            query,
            ticker: detectedTicker,
            doc_types: null,
            top_k: 10,
            session_id: currentSessionId
        };
        
        console.log('Request data:', requestData);
        
        const response = await fetch(`${API_BASE_URL}/query`, {
            method: 'POST',
            headers: { 'Content-Type': 'application/json' },
            body: JSON.stringify(requestData)
        });
        
        if (!response.ok) {
            const errorData = await response.json().catch(() => ({}));
            throw new Error(errorData.detail || `HTTP ${response.status}`);
        }
        
        const data = await response.json();
        
        if (!data) {
            throw new Error('Empty response from server');
        }
        
        // Add to conversation history
        conversationHistory.push({
            query,
            answer: data.answer,
            sources: data.sources,
            expanded_queries: data.expanded_queries,
            timestamp: new Date()
        });
        
        // Save to history
        saveToHistory(query, data.answer || '');
        
        // Display all conversation history
        displayConversationHistory();
        
        // Clear input
        queryInput.value = '';
        
        // Scroll to latest answer
        setTimeout(() => {
            const rightPanel = document.querySelector('.right-panel');
            if (rightPanel) {
                rightPanel.scrollTop = rightPanel.scrollHeight;
            }
        }, 100);
        
    } catch (error) {
        console.error('Error:', error);
        showError(error.message || 'Failed to process query');
        
        // Show initial state again if no history
        if (conversationHistory.length === 0) {
            if (loadingState) loadingState.style.display = 'none';
            if (initialState) initialState.style.display = 'flex';
        }
        
    } finally {
        setLoading(false);
        if (loadingState) loadingState.style.display = 'none';
    }
}

// Display entire conversation history
function displayConversationHistory() {
    if (!answerDiv || !sourcesDiv) return;
    
    // Clear existing content
    answerDiv.innerHTML = '';
    sourcesDiv.innerHTML = '';
    
    // Render each Q&A pair
    conversationHistory.forEach((item, index) => {
        // Create conversation item container
        const conversationItem = document.createElement('div');
        conversationItem.className = 'conversation-item';
        conversationItem.style.cssText = `
            margin-bottom: 2rem;
            padding-bottom: 2rem;
            border-bottom: 1px solid rgba(255, 255, 255, 0.1);
        `;
        
        // Add question
        const questionDiv = document.createElement('div');
        questionDiv.className = 'question-block';
        questionDiv.style.cssText = `
            background: rgba(59, 130, 246, 0.1);
            border-left: 3px solid #3b82f6;
            padding: 1rem 1.5rem;
            border-radius: 0.5rem;
            margin-bottom: 1.5rem;
        `;
        questionDiv.innerHTML = `
            <div style="color: #60a5fa; font-size: 0.875rem; margin-bottom: 0.5rem;">Question ${index + 1}</div>
            <div style="color: #e5e7eb; font-size: 1rem;">${escapeHtml(item.query)}</div>
        `;
        conversationItem.appendChild(questionDiv);
        
        // Add answer
        const answerBlock = document.createElement('div');
        answerBlock.className = 'answer-block';
        answerBlock.style.cssText = `
            background: rgba(16, 185, 129, 0.05);
            border-left: 3px solid #10b981;
            padding: 1rem 1.5rem;
            border-radius: 0.5rem;
            margin-bottom: 1.5rem;
        `;
        answerBlock.innerHTML = `
            <div style="color: #34d399; font-size: 0.875rem; margin-bottom: 0.5rem;">Answer</div>
            <div style="color: #e5e7eb; line-height: 1.7;">${formatAnswer(item.answer || 'No answer available')}</div>
        `;
        conversationItem.appendChild(answerBlock);
        
        // Add expanded queries if available
        if (item.expanded_queries && item.expanded_queries.length > 1) {
            const expandedDiv = document.createElement('div');
            expandedDiv.className = 'expanded-queries-block';
            expandedDiv.style.cssText = `
                background: rgba(139, 92, 246, 0.05);
                border-left: 3px solid #8b5cf6;
                padding: 1rem 1.5rem;
                border-radius: 0.5rem;
                margin-bottom: 1.5rem;
            `;
            expandedDiv.innerHTML = `
                <div style="color: #a78bfa; font-size: 0.875rem; margin-bottom: 0.5rem;">Expanded Queries</div>
                <ul style="margin: 0; padding-left: 1.5rem; color: #d1d5db;">
                    ${item.expanded_queries.map(q => `<li style="margin-bottom: 0.5rem;">${escapeHtml(q)}</li>`).join('')}
                </ul>
            `;
            conversationItem.appendChild(expandedDiv);
        }
        
        // Add sources
        if (item.sources && item.sources.length > 0) {
            const sourcesHeader = document.createElement('div');
            sourcesHeader.style.cssText = `
                color: #f59e0b;
                font-size: 0.875rem;
                margin-bottom: 1rem;
                font-weight: 600;
            `;
            sourcesHeader.textContent = `Sources (${item.sources.length})`;
            conversationItem.appendChild(sourcesHeader);
            
            const sourcesContainer = document.createElement('div');
            sourcesContainer.innerHTML = item.sources.map((source, sourceIndex) => {
                if (!source) return '';
                
                const docTypeLabel = source.doc_type === '10k' ? '10-K Report' : 
                                    source.doc_type.replace('_', ' ').replace(/\b\w/g, l => l.toUpperCase());
                
                return `
                    <div class="source-card" id="source-${index}-${sourceIndex}" style="
                        background: rgba(31, 41, 55, 0.5);
                        border: 1px solid rgba(75, 85, 99, 0.5);
                        border-radius: 0.5rem;
                        margin-bottom: 0.75rem;
                        overflow: hidden;
                        transition: all 0.3s ease;
                    ">
                        <div class="source-header" onclick="toggleSource(${index}, ${sourceIndex})" style="
                            padding: 0.75rem 1rem;
                            cursor: pointer;
                            display: flex;
                            justify-content: space-between;
                            align-items: center;
                            background: rgba(17, 24, 39, 0.5);
                        ">
                            <div class="source-title" style="display: flex; align-items: center; gap: 0.5rem;">
                                <span class="source-badge" style="
                                    background: rgba(59, 130, 246, 0.2);
                                    color: #60a5fa;
                                    padding: 0.25rem 0.5rem;
                                    border-radius: 0.25rem;
                                    font-size: 0.75rem;
                                    font-weight: 600;
                                ">${docTypeLabel}</span>
                                <span style="color: #e5e7eb; font-size: 0.875rem;">${escapeHtml(source.filename || 'Unknown')}</span>
                            </div>
                            <div class="source-similarity" style="color: #9ca3af; font-size: 0.875rem;">
                                Score: <strong style="color: #10b981;">${source.similarity_score ? (source.similarity_score * 100).toFixed(0) + '%' : 'N/A'}</strong>
                            </div>
                        </div>
                        <div class="source-content" style="
                            max-height: 0;
                            overflow: hidden;
                            transition: max-height 0.3s ease;
                            padding: 0 1rem;
                        ">
                            <div class="source-details" style="
                                padding: 0.75rem 0;
                                display: flex;
                                gap: 1rem;
                                flex-wrap: wrap;
                                border-bottom: 1px solid rgba(75, 85, 99, 0.3);
                            ">
                                ${source.ticker ? `<span style="color: #9ca3af; font-size: 0.875rem;"><strong style="color: #e5e7eb;">Ticker:</strong> ${escapeHtml(source.ticker)}</span>` : ''}
                                ${source.chunk_id ? `<span style="color: #9ca3af; font-size: 0.875rem;"><strong style="color: #e5e7eb;">Chunk:</strong> ${escapeHtml(source.chunk_id)}</span>` : ''}
                            </div>
                            <div class="source-preview" style="
                                padding: 1rem 0;
                                color: #d1d5db;
                                font-size: 0.875rem;
                                line-height: 1.6;
                                font-style: italic;
                            ">
                                "${escapeHtml(source.text_preview || 'No preview available')}"
                            </div>
                        </div>
                    </div>
                `;
            }).filter(Boolean).join('');
            
            conversationItem.appendChild(sourcesContainer);
        }
        
        // Add timestamp
        const timestamp = document.createElement('div');
        timestamp.style.cssText = `
            color: #6b7280;
            font-size: 0.75rem;
            text-align: right;
            margin-top: 1rem;
        `;
        timestamp.textContent = item.timestamp.toLocaleTimeString();
        conversationItem.appendChild(timestamp);
        
        answerDiv.appendChild(conversationItem);
    });
    
    // Show export button if there's content
    if (exportBtn && conversationHistory.length > 0) {
        exportBtn.style.display = 'flex';
    }
}

// Toggle source expansion
function toggleSource(conversationIndex, sourceIndex) {
    const card = document.getElementById(`source-${conversationIndex}-${sourceIndex}`);
    if (!card) return;
    
    const content = card.querySelector('.source-content');
    if (!content) return;
    
    if (content.style.maxHeight && content.style.maxHeight !== '0px') {
        content.style.maxHeight = '0px';
        content.style.padding = '0 1rem';
    } else {
        content.style.maxHeight = content.scrollHeight + 'px';
        content.style.padding = '0 1rem';
    }
}

// Format answer with enhanced styling - IMPROVED VERSION
function formatAnswer(answer) {
    if (!answer) return '';
    
    let formatted = escapeHtml(answer);
    
    // ============================================
    // 1. CONVERT LATEX MATH TO READABLE FORMAT
    // ============================================
    
    // Block math: \[ ... \] → Display as formatted equation
    formatted = formatted.replace(/\\\[(.*?)\\\]/gs, (match, equation) => {
        // Clean up the equation
        let cleanEq = equation.trim()
            .replace(/\\text\{([^}]+)\}/g, '$1')  // Remove \text{}
            .replace(/\\frac\{([^}]+)\}\{([^}]+)\}/g, '($1) ÷ ($2)')  // Convert fractions
            .replace(/\\_/g, '_')  // Fix escaped underscores
            .replace(/\\\\/g, '')  // Remove double backslashes
            .replace(/\\times/g, '×')  // Convert \times to ×
            .replace(/\\div/g, '÷')  // Convert \div to ÷
            .replace(/\\approx/g, '≈')  // Convert \approx to ≈
            .replace(/\\cdot/g, '·');  // Convert \cdot to ·
        
        return `<div class="math-formula">${cleanEq}</div>`;
    });
    
    // Inline math: \( ... \) → Convert to readable format
    formatted = formatted.replace(/\\\((.*?)\\\)/g, (match, equation) => {
        let cleanEq = equation.trim()
            .replace(/\\text\{([^}]+)\}/g, '$1')
            .replace(/\\frac\{([^}]+)\}\{([^}]+)\}/g, '($1) ÷ ($2)')
            .replace(/\\times/g, '×')
            .replace(/\\div/g, '÷')
            .replace(/\\approx/g, '≈');
        
        return `<span class="math-inline">${cleanEq}</span>`;
    });
    
    // ============================================
    // 2. CLEAN UP ASTERISKS AND FORMATTING
    // ============================================
    
    // Remove asterisks used for emphasis (markdown bold/italic)
    formatted = formatted.replace(/\*\*\*([^*]+)\*\*\*/g, '<strong><em>$1</em></strong>');  // ***bold italic***
    formatted = formatted.replace(/\*\*([^*]+)\*\*/g, '<strong>$1</strong>');  // **bold**
    formatted = formatted.replace(/\*([^*\n]+)\*/g, '<em>$1</em>');  // *italic*
    
    // Replace any remaining asterisks with multiplication symbol
    formatted = formatted.replace(/\s\*\s/g, ' × ');
    
    // ============================================
    // 3. FORMAT SECTIONS AND HEADERS
    // ============================================
    
    // Line breaks
    formatted = formatted.replace(/\n\n/g, '<br><br>');
    formatted = formatted.replace(/\n/g, '<br>');
    
    // Headers with ###
    formatted = formatted.replace(/###\s*([^<]+?)(<br>|$)/g, '<h4 class="answer-header">$1</h4>');
    
    // Section headers that end with colon (e.g., "Analysis:", "Sources:")
    formatted = formatted.replace(/<strong>([^<]+?):<\/strong>/g, '<div class="section-header">$1:</div>');
    
    // ============================================
    // 4. HIGHLIGHT FINANCIAL DATA
    // ============================================
    
    // Citations [Source N]
    formatted = formatted.replace(/\[Source (\d+)\]/g, 
        '<span class="citation">[Source $1]</span>');
    
    // Currency values (e.g., $2.54B, $1.26M, $1,234.56)
    formatted = formatted.replace(/\$[\d,]+\.?\d*[BMK]?/g, 
        match => `<span class="currency-value">${match}</span>`);
    
    // Percentages
    formatted = formatted.replace(/(\d+\.?\d*)%/g, 
        '<span class="percentage-value">$1%</span>');
    
    // Standalone ratios (e.g., 2.29, 1.5x, 3:1)
    formatted = formatted.replace(/\b(\d+\.?\d+)(x|:1)(\s|<|$)/g, 
        '<span class="ratio-value">$1$2</span>$3');
    
    // ============================================
    // 5. FORMAT LISTS
    // ============================================
    
    // Numbered lists with bold headers
    formatted = formatted.replace(/(\d+)\.\s+<strong>([^<]+)<\/strong>/g, 
        '<div class="list-item"><strong>$1. $2</strong></div>');
    
    // Bullet points
    formatted = formatted.replace(/(<br>|^)[\s]*[-•]\s+/g, '$1<span class="bullet-point">• </span>');
    
    // ============================================
    // 6. CLEAN UP DOUBLE SPACES
    // ============================================
    
    formatted = formatted.replace(/\s{2,}/g, ' ');
    
    return formatted;
}

// Export to PDF function
function exportToPDF() {
    if (conversationHistory.length === 0) {
        showError('No conversation to export. Please ask a question first.');
        return;
    }
    
    try {
        if (typeof window.jspdf === 'undefined') {
            showError('PDF library not loaded. Please refresh the page.');
            return;
        }
        
        const { jsPDF } = window.jspdf;
        const doc = new jsPDF();
        
        // Set title
        doc.setFontSize(18);
        doc.setFont(undefined, 'bold');
        doc.text('FinSage Analytics Report', 20, 20);
        
        doc.setFontSize(10);
        doc.setFont(undefined, 'normal');
        doc.setTextColor(100);
        doc.text(`Generated on ${new Date().toLocaleString()}`, 20, 28);
        doc.text(`Total Questions: ${conversationHistory.length}`, 20, 34);
        
        let yPos = 45;
        
        conversationHistory.forEach((item, index) => {
            // Check for page break
            if (yPos > 250) {
                doc.addPage();
                yPos = 20;
            }
            
            // Question
            doc.setFontSize(12);
            doc.setFont(undefined, 'bold');
            doc.setTextColor(0);
            doc.text(`Question ${index + 1}:`, 20, yPos);
            yPos += 8;
            
            doc.setFont(undefined, 'normal');
            doc.setFontSize(10);
            const queryLines = doc.splitTextToSize(item.query, 170);
            doc.text(queryLines, 20, yPos);
            yPos += queryLines.length * 7 + 5;
            
            // Answer
            if (yPos > 250) {
                doc.addPage();
                yPos = 20;
            }
            
            doc.setFontSize(11);
            doc.setFont(undefined, 'bold');
            doc.text('Answer:', 20, yPos);
            yPos += 8;
            
            doc.setFont(undefined, 'normal');
            doc.setFontSize(10);
            
            let cleanAnswer = item.answer
                .replace(/<[^>]*>/g, '')
                .replace(/\[Source \d+\]/g, '')
                .replace(/&nbsp;/g, ' ')
                .replace(/&lt;/g, '<')
                .replace(/&gt;/g, '>')
                .replace(/&amp;/g, '&')
                .replace(/\\/g, '');
            
            const answerLines = doc.splitTextToSize(cleanAnswer, 170);
            
            answerLines.forEach(line => {
                if (yPos > 270) {
                    doc.addPage();
                    yPos = 20;
                }
                doc.text(line, 20, yPos);
                yPos += 7;
            });
            
            yPos += 10;
        });
        
        const filename = `FinSage_Conversation_${Date.now()}.pdf`;
        doc.save(filename);
        
    } catch (error) {
        console.error('PDF Export Error:', error);
        showError('Failed to export PDF. Please try again.');
    }
}

// Loading state
function setLoading(isLoading) {
    if (!submitBtn) return;
    
    submitBtn.disabled = isLoading;
    
    if (btnText && loader) {
        btnText.style.display = isLoading ? 'none' : 'inline';
        loader.style.display = isLoading ? 'inline-block' : 'none';
    }
}

// Error handling
function showError(message) {
    if (!errorToast || !errorMessage) {
        alert(message);
        return;
    }
    
    errorMessage.textContent = message;
    errorToast.style.display = 'block';
    
    setTimeout(() => {
        errorToast.style.display = 'none';
    }, 5000);
}

function hideError() {
    if (errorToast) {
        errorToast.style.display = 'none';
    }
}

// Utility: Escape HTML
function escapeHtml(text) {
    if (!text) return '';
    const div = document.createElement('div');
    div.textContent = text;
    return div.innerHTML;
}

// Health check
async function checkHealth() {
    try {
        const response = await fetch(`${API_BASE_URL}/health`, {
            signal: AbortSignal.timeout(5000)
        });
        
        if (!response.ok) {
            console.warn('API health check failed');
        }
    } catch (error) {
        console.error('Cannot connect to API:', error);
        showError('Backend API is not responding. Please start the server.');
    }
}

// Make toggleSource available globally
window.toggleSource = toggleSource;

// Initialize on load
init();