static/script.js

document.addEventListener('DOMContentLoaded', () => {
    const uploadBtn = document.getElementById('upload-btn');
    const indexBtn = document.getElementById('index-btn');
    const sendBtn = document.getElementById('send-btn');
    const promptInput = document.getElementById('prompt-input');
    const chatContainer = document.getElementById('chat-container');
    const topKInput = document.getElementById('top-k');
    const statusMsg = document.getElementById('index-status');

    // Auto-resize textarea
    promptInput.addEventListener('input', function () {
        this.style.height = 'auto';
        this.style.height = (this.scrollHeight) + 'px';
        if (this.value.trim().length > 0) {
            sendBtn.removeAttribute('disabled');
        } else {
            sendBtn.setAttribute('disabled', 'true');
        }
    });

    promptInput.addEventListener('keydown', (e) => {
        if (e.key === 'Enter' && !e.shiftKey) {
            e.preventDefault();
            handleSend();
        }
    });

    sendBtn.addEventListener('click', handleSend);

    async function fetchDocuments() {
        const docList = document.getElementById('documents-list');
        try {
            const res = await fetch('/api/documents');
            const data = await res.json();

            docList.innerHTML = '';
            if (data.documents.length === 0) {
                docList.innerHTML = '<div style="font-size:0.8rem; color:var(--text-secondary); padding: 8px;">No documents uploaded.</div>';
                return;
            }

            data.documents.forEach(doc => {
                const card = document.createElement('div');
                card.className = 'document-card';

                const sizeKb = (doc.size / 1024).toFixed(1);

                let iconHtml = '';
                if (doc.thumbnail) {
                    iconHtml = `<img src="${doc.thumbnail}" alt="${doc.name}" class="doc-thumbnail" onerror="this.style.display='none'">`;
                } else {
                    const iconClass = doc.type === 'pdf' ? 'icon-pdf' :
                        doc.type === 'docx' ? 'icon-docx' :
                            doc.type === 'txt' ? 'icon-txt' : 'icon-default';
                    iconHtml = `<div class="doc-icon ${iconClass}">${doc.type.toUpperCase()}</div>`;
                }

                card.innerHTML = `
                    ${iconHtml}
                    <div class="doc-name" title="${doc.name}">${doc.name}</div>
                    <div class="doc-size">${sizeKb} KB</div>
                `;

                // Add click handler to open embeddings viewer
                card.style.cursor = 'pointer';
                card.addEventListener('click', () => openChunksModal(doc.name));

                docList.appendChild(card);
            });
        } catch (e) {
            console.error("Failed to load documents", e);
        }
    }

    // Call once on load
    fetchDocuments();
    fetchStats();

    let isIndexing = false;
    let pollInterval = null;
    const overlay = document.getElementById('indexing-overlay');
    const overlayText = document.getElementById('overlay-text');

    async function checkIndexStatus() {
        try {
            const res = await fetch('/api/index/status');
            const data = await res.json();

            if (data.is_indexing) {
                if (data.progress) {
                    overlayText.textContent = data.progress;
                }
                if (!isIndexing) {
                    isIndexing = true;
                    statusMsg.innerHTML = '<span class="spinner"></span> Indexing in progress...';
                    statusMsg.style.color = 'var(--brand-color)';
                    indexBtn.disabled = true;
                    uploadBtn.disabled = true;
                    overlay.style.display = 'flex';
                    if (!pollInterval) {
                        pollInterval = setInterval(checkIndexStatus, 2000);
                    }
                }
            } else {
                if (isIndexing) {
                    isIndexing = false;
                    statusMsg.textContent = 'Ready.';
                    statusMsg.style.color = 'var(--text-secondary)';
                    indexBtn.disabled = false;
                    uploadBtn.disabled = false;
                    overlay.style.display = 'none';
                    if (pollInterval) {
                        clearInterval(pollInterval);
                        pollInterval = null;

                        // Wait a tiny bit and refresh so that thumbnails appear
                        setTimeout(fetchDocuments, 1000);
                    }
                }
            }
        } catch (e) {
            console.error("Failed to check index status", e);
        }
    }

    // Start polling if we are currently indexing
    checkIndexStatus();

    uploadBtn.addEventListener('click', async () => {
        const fileInput = document.getElementById('file-upload');
        if (fileInput.files.length === 0) return alert('Select files first.');

        const formData = new FormData();
        for (const file of fileInput.files) {
            formData.append('files', file);
        }

        uploadBtn.textContent = 'Uploading...';
        uploadBtn.disabled = true;

        try {
            const response = await fetch('/api/upload', {
                method: 'POST',
                body: formData
            });
            const result = await response.json();
            alert(result.message);
            fileInput.value = "";

            // Refresh documents list 
            fetchDocuments();
            // Force status check for auto-indexing
            checkIndexStatus();

        } catch (error) {
            console.error(error);
            alert('Upload failed.');
        } finally {
            uploadBtn.textContent = 'Upload';
            if (!isIndexing) uploadBtn.disabled = false;
        }
    });

    indexBtn.addEventListener('click', async () => {
        indexBtn.textContent = 'Starting...';
        indexBtn.disabled = true;
        statusMsg.innerHTML = '<span class="spinner"></span> Triggering index...';
        try {
            const response = await fetch('/api/index', { method: 'POST' });
            const result = await response.json();
            console.log(result.message);
            checkIndexStatus();
        } catch (error) {
            console.error(error);
            statusMsg.textContent = "Error triggering indexing.";
            indexBtn.disabled = false;
        } finally {
            indexBtn.textContent = 'Re-index All';
        }
    });

    const clearBtn = document.getElementById('clear-btn');
    clearBtn.addEventListener('click', async () => {
        if (!confirm('Are you sure you want to clear ALL indexed data? This action cannot be undone.')) {
            return;
        }

        clearBtn.textContent = 'Clearing...';
        clearBtn.disabled = true;
        try {
            const response = await fetch('/api/index/clear', { method: 'POST' });
            const result = await response.json();
            alert(result.message);
            fetchDocuments();
            fetchStats();
        } catch (error) {
            console.error(error);
            alert('Failed to clear database.');
        } finally {
            clearBtn.textContent = 'Clear All Data';
            clearBtn.disabled = false;
        }
    });

    async function fetchStats() {
        try {
            const response = await fetch('/api/stats');
            const data = await response.json();
            document.getElementById('stat-embed').textContent = data.tokens.embedding_tokens.toLocaleString();
            document.getElementById('stat-prompt').textContent = data.tokens.prompt_tokens.toLocaleString();
            document.getElementById('stat-response').textContent = data.tokens.completion_tokens.toLocaleString();
            document.getElementById('stat-toc').textContent = (data.tokens.toc_analysis_tokens || 0).toLocaleString();
            document.getElementById('stat-total').textContent = data.total_tokens.toLocaleString();
        } catch (error) {
            console.error('Error fetching stats:', error);
        }
    }

    function cleanMarkdown(text) {
        return marked.parse(text);
    }

    function createResponseCard(title) {
        const card = document.createElement('div');
        card.className = 'response-card';
        card.innerHTML = `
            <div class="response-header">
                <svg width="20" height="20" viewBox="0 0 24 24" fill="none" class="ai-icon">
                    <path d="M12 2L2 7L12 12L22 7L12 2Z" stroke="currentColor" stroke-width="2" stroke-linejoin="round"/>
                    <path d="M2 17L12 22L22 17" stroke="currentColor" stroke-width="2" stroke-linejoin="round"/>
                    <path d="M2 12L12 17L22 12" stroke="currentColor" stroke-width="2" stroke-linejoin="round"/>
                </svg>
                <span>${title}</span>
            </div>
            <div class="response-content"><span class="typing-indicator"></span></div>
            <div class="metrics">
                <div class="metrics-row">
                    <div class="token-badges">
                        <span class="badge prompt-badge" style="display:none;" title="Prompt Tokens">Prompt: --</span>
                        <span class="badge completion-badge" style="display:none;" title="Completion Tokens">Response: --</span>
                    </div>
                    <button class="sources-toggle">View Sources ▼</button>
                </div>
                <div class="sources-list" style="display: none;"></div>
            </div>
        `;
        return card;
    }

    async function handleSend() {
        const query = promptInput.value.trim();
        if (!query) return;

        promptInput.value = '';
        promptInput.style.height = 'auto';
        sendBtn.disabled = true;

        // Create User Message
        const userMsg = document.createElement('div');
        userMsg.className = 'message user';
        userMsg.innerHTML = `<div class="user-bubble">${query}</div>`;
        chatContainer.appendChild(userMsg);

        // Create Container for Responses
        const sysMsg = document.createElement('div');
        sysMsg.className = 'message system';
        const responsesGrid = document.createElement('div');
        responsesGrid.className = 'system-responses';

        // Create Cards
        const vectorCard = createResponseCard('Vector RAG');
        const pageCard = createResponseCard('Page Index RAG');
        const tocCard = createResponseCard('TOC Index RAG');

        responsesGrid.appendChild(vectorCard);
        responsesGrid.appendChild(pageCard);
        responsesGrid.appendChild(tocCard);
        sysMsg.appendChild(responsesGrid);
        chatContainer.appendChild(sysMsg);

        chatContainer.scrollTop = chatContainer.scrollHeight;

        const topK = parseInt(topKInput.value) || 5;

        // Trigger all three SSE calls
        fetchSSE('/api/chat/vector', { query, top_k: topK }, vectorCard);
        fetchSSE('/api/chat/page', { query, top_k: topK }, pageCard);
        fetchSSE('/api/chat/toc', { query, top_k: topK }, tocCard);
    }

    async function fetchSSE(url, payload, cardElement) {
        const contentDiv = cardElement.querySelector('.response-content');
        const tokenDisplay = cardElement.querySelector('.prompt-tokens');
        const sourcesBtn = cardElement.querySelector('.sources-toggle');
        const sourcesList = cardElement.querySelector('.sources-list');

        sourcesBtn.addEventListener('click', () => {
            if (sourcesList.style.display === 'none') {
                sourcesList.style.display = 'block';
                sourcesBtn.textContent = 'Hide Sources ▲';
            } else {
                sourcesList.style.display = 'none';
                sourcesBtn.textContent = 'View Sources ▼';
            }
        });

        try {
            const response = await fetch(url, {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify(payload)
            });

            const reader = response.body.getReader();
            const decoder = new TextDecoder("utf-8");
            let fullText = "";
            let buffer = "";

            while (true) {
                const { done, value } = await reader.read();
                if (done) break;

                buffer += decoder.decode(value, { stream: true });
                const parts = buffer.split('\n\n');
                buffer = parts.pop();

                for (const part of parts) {
                    if (part.startsWith('data: ')) {
                        try {
                            const dataStr = part.substring(6);
                            if (!dataStr.trim()) continue;
                            const data = JSON.parse(dataStr);

                            if (data.type === 'token') {
                                fullText += data.content;
                                contentDiv.innerHTML = cleanMarkdown(fullText) + '<span class="typing-indicator"></span>';
                            } else if (data.type === 'sources') {
                                // Render sources
                                let sourcesHtml = '';
                                data.sources.forEach((s, idx) => {
                                    let label = s.chunk_index !== undefined ? `Chunk #${s.chunk_index}` : s.node_id !== undefined ? `Node ${s.node_id}: ${s.title || ''}` : `Page ${s.page_num}`;
                                    sourcesHtml += `
                                        <div class="source-item">
                                            <span class="source-meta">[${idx + 1}] ${s.source} — ${label} (score: ${s.score})</span>
                                            <span class="source-text">${s.text.substring(0, 150)}...</span>
                                        </div>
                                    `;
                                });
                                sourcesList.innerHTML = sourcesHtml || '<div class="source-item">No sources found.</div>';
                            } else if (data.type === 'stats') {
                                const promptBadge = cardElement.querySelector('.prompt-badge');
                                const completionBadge = cardElement.querySelector('.completion-badge');
                                if (promptBadge) {
                                    promptBadge.textContent = 'Prompt: ' + data.prompt_eval_count;
                                    promptBadge.style.display = 'inline-block';
                                }
                                if (completionBadge) {
                                    completionBadge.textContent = 'Response: ' + data.eval_count;
                                    completionBadge.style.display = 'inline-block';
                                }
                                fetchStats();
                            } else if (data.type === 'error') {
                                fullText += `\n\n**Error:** ${data.content}`;
                                contentDiv.innerHTML = cleanMarkdown(fullText);
                            }
                        } catch (e) {
                            console.error('JSON parse error:', e, part);
                        }
                    }
                }
                chatContainer.scrollTop = chatContainer.scrollHeight;
            }

            // Remove typing indicator at the end
            contentDiv.innerHTML = cleanMarkdown(fullText);

        } catch (error) {
            contentDiv.innerHTML = "<em>Connection error parsing stream.</em>";
            console.error(error);
        }
    }

    // Modal Logic
    const chunksModal = document.getElementById('chunks-modal');
    const closeModalBtn = document.getElementById('close-modal-btn');
    const modalDocTitle = document.getElementById('modal-doc-title');
    const pageChunksList = document.getElementById('page-chunks-list');
    const vectorChunksList = document.getElementById('vector-chunks-list');
    const tocChunksList = document.getElementById('toc-chunks-list');

    closeModalBtn.addEventListener('click', () => {
        chunksModal.style.display = 'none';
    });

    // Close if clicking outside
    window.addEventListener('click', (e) => {
        if (e.target === chunksModal) {
            chunksModal.style.display = 'none';
        }
    });

    async function openChunksModal(filename) {
        modalDocTitle.textContent = filename;
        pageChunksList.innerHTML = '<div class="spinner large-spinner"></div>';
        vectorChunksList.innerHTML = '<div class="spinner large-spinner"></div>';
        tocChunksList.innerHTML = '<div class="spinner large-spinner"></div>';
        chunksModal.style.display = 'flex';

        try {
            const res = await fetch(`/api/documents/${encodeURIComponent(filename)}/chunks`);
            const data = await res.json();

            // Render Page Chunks
            pageChunksList.innerHTML = '';
            let totalPageTokens = 0;
            if (data.page_chunks && data.page_chunks.length > 0) {
                data.page_chunks.forEach(chunk => {
                    const el = document.createElement('div');
                    el.className = 'chunk-item';
                    el.innerHTML = `
                        <span class="chunk-meta" style="display:flex; justify-content:space-between;">
                            <span>Page ${chunk.page_num}</span>
                            <span style="color:var(--text-secondary); font-size:0.85em;">Tokens: ${chunk.tokens || 0}</span>
                        </span>
                        <div class="chunk-text">${chunk.text}</div>
                    `;
                    pageChunksList.appendChild(el);
                    totalPageTokens += (chunk.tokens || 0);
                });

                const totalEl = document.createElement('div');
                totalEl.style.marginTop = '16px';
                totalEl.style.padding = '12px';
                totalEl.style.backgroundColor = 'var(--bg-secondary)';
                totalEl.style.borderRadius = '6px';
                totalEl.style.fontWeight = 'bold';
                totalEl.style.display = 'flex';
                totalEl.style.justifyContent = 'space-between';
                totalEl.style.border = '1px solid var(--border-color)';
                totalEl.innerHTML = `<span>Total Page RAG Tokens:</span> <span>${totalPageTokens.toLocaleString()}</span>`;
                pageChunksList.appendChild(totalEl);
            } else {
                pageChunksList.innerHTML = '<div style="padding:16px;">No page chunks found.</div>';
            }

            // Render Vector Chunks
            vectorChunksList.innerHTML = '';
            let totalVectorTokens = 0;
            if (data.vector_chunks && data.vector_chunks.length > 0) {
                data.vector_chunks.forEach(chunk => {
                    const el = document.createElement('div');
                    el.className = 'chunk-item';
                    el.innerHTML = `
                        <span class="chunk-meta" style="display:flex; justify-content:space-between;">
                            <span>Chunk #${chunk.chunk_index}</span>
                            <span style="color:var(--text-secondary); font-size:0.85em;">Tokens: ${chunk.tokens || 0}</span>
                        </span>
                        <div class="chunk-text">${chunk.text}</div>
                    `;
                    vectorChunksList.appendChild(el);
                    totalVectorTokens += (chunk.tokens || 0);
                });

                const totalEl = document.createElement('div');
                totalEl.style.marginTop = '16px';
                totalEl.style.padding = '12px';
                totalEl.style.backgroundColor = 'var(--bg-secondary)';
                totalEl.style.borderRadius = '6px';
                totalEl.style.fontWeight = 'bold';
                totalEl.style.display = 'flex';
                totalEl.style.justifyContent = 'space-between';
                totalEl.style.border = '1px solid var(--border-color)';
                totalEl.innerHTML = `<span>Total Vector RAG Tokens:</span> <span>${totalVectorTokens.toLocaleString()}</span>`;
                vectorChunksList.appendChild(totalEl);
            } else {
                vectorChunksList.innerHTML = '<div style="padding:16px;">No vector chunks found.</div>';
            }

            // Render TOC Chunks
            tocChunksList.innerHTML = '';
            let totalTocTokens = 0;
            if (data.toc_chunks && data.toc_chunks.length > 0) {
                data.toc_chunks.forEach(chunk => {
                    const el = document.createElement('div');
                    el.className = 'chunk-item';
                    el.innerHTML = `
                        <span class="chunk-meta" style="display:flex; justify-content:space-between;">
                            <span>Node ${chunk.node_id}: ${chunk.title || ''}</span>
                            <span style="color:var(--text-secondary); font-size:0.85em;">Tokens: ${chunk.tokens || 0}</span>
                        </span>
                        <div class="chunk-text">${chunk.text}</div>
                    `;
                    tocChunksList.appendChild(el);
                    totalTocTokens += (chunk.tokens || 0);
                });

                const totalEl = document.createElement('div');
                totalEl.style.marginTop = '16px';
                totalEl.style.padding = '12px';
                totalEl.style.backgroundColor = 'var(--bg-secondary)';
                totalEl.style.borderRadius = '6px';
                totalEl.style.fontWeight = 'bold';
                totalEl.style.display = 'flex';
                totalEl.style.justifyContent = 'space-between';
                totalEl.style.border = '1px solid var(--border-color)';
                totalEl.innerHTML = `<span>Total TOC RAG Tokens:</span> <span>${totalTocTokens.toLocaleString()}</span>`;
                tocChunksList.appendChild(totalEl);
            } else {
                tocChunksList.innerHTML = '<div style="padding:16px;">No TOC nodes found.</div>';
            }

        } catch (error) {
            console.error(error);
            pageChunksList.innerHTML = '<div style="color:red; padding:16px;">Error loading data.</div>';
            vectorChunksList.innerHTML = '<div style="color:red; padding:16px;">Error loading data.</div>';
            tocChunksList.innerHTML = '<div style="color:red; padding:16px;">Error loading data.</div>';
        }
    }
});