static/script.js

document.addEventListener('DOMContentLoaded', () => {
  const queryInput = document.getElementById('queryInput');
  const askButton = document.getElementById('askButton');
  const chatMessages = document.getElementById('chatMessages');
  const attachFileBtn = document.getElementById('attachFileBtn');
  const fileInput = document.getElementById('fileInput');
  const ocrModal = document.getElementById('ocrModal');
  const closeOcrModal = document.getElementById('closeOcrModal');
  const useOcrText = document.getElementById('useOcrText');
  const queryWithOcr = document.getElementById('queryWithOcr');
  const cancelOcr = document.getElementById('cancelOcr');
  const ocrPreview = document.getElementById('ocrPreview');
  const ocrResults = document.getElementById('ocrResults');

  // 💬 Conversation state management
  let conversationHistory = [];
  let currentOcrText = '';
  let currentOcrImage = null;

  function addUserMessage(content) {
    conversationHistory.push({
      role: 'user',
      content: content,
      timestamp: Date.now()
    });
  }

  function addAssistantMessage(content, source) {
    conversationHistory.push({
      role: 'assistant',
      content: content,
      timestamp: Date.now(),
      source: source
    });
  }

  function clearConversation() {
    conversationHistory = [];
  }

  function getHistory() {
    return conversationHistory;
  }

  // 💬 Message rendering functions
  function renderUserMessage(content, timestamp = null) {
    const ts = timestamp || Date.now();
    return `
      <div class="message user" data-timestamp="${ts}">
        <div class="message-avatar">
          <i class="fas fa-user"></i>
        </div>
        <div class="message-content">
          <div class="message-bubble">
          <div class="message-text">${escapeHtml(content)}</div>
          </div>
        </div>
      </div>
    `;
  }

  function renderAssistantMessage(content, source) {
    const sourceBadge = source ? `<span class="source-badge">${source}</span>` : '';
    return `
      <div class="message assistant">
        <div class="message-avatar">
          <i class="fas fa-robot"></i>
        </div>
        <div class="message-content">
          ${sourceBadge}
          <div class="message-bubble">
          <div class="message-text">${formatAnswer(content)}</div>
          </div>
        </div>
      </div>
    `;
  }

  function renderLoadingMessage() {
    return `
      <div class="message assistant">
        <div class="message-avatar">
          <i class="fas fa-robot"></i>
        </div>
        <div class="message-content">
          <div class="message-bubble">
            <div class="loading-message">
              <span>Thinking...</span>
          <div class="typing-indicator">
            <span></span>
            <span></span>
            <span></span>
              </div>
            </div>
          </div>
        </div>
      </div>
    `;
  }

  function renderWelcomeMessage() {
    return `
      <div class="welcome-message">
        <h2>Welcome to Corex!</h2>
        <p>Ask me anything and I'll help you with accurate, document-backed answers.</p>
      </div>
    `;
  }

  function escapeHtml(text) {
    const div = document.createElement('div');
    div.textContent = text;
    return div.innerHTML;
  }

  function displayAllMessages() {
    if (conversationHistory.length === 0) {
      chatMessages.innerHTML = renderWelcomeMessage();
      return;
    }

    let html = '';
    conversationHistory.forEach(msg => {
      if (msg.role === 'user') {
        html += renderUserMessage(msg.content);
      } else {
        html += renderAssistantMessage(msg.content, msg.source);
      }
    });
    chatMessages.innerHTML = html;
    scrollToBottom();
  }

  function scrollToBottom() {
    chatMessages.scrollTop = chatMessages.scrollHeight;
  }

function formatAnswer(text) {
  if (typeof text !== "string") {
    text = String(text ?? "No response received.");
  }
  return text
    .split('\n')
    .filter(line => line.trim())
    .map(line => `<p>${line}</p>`)
    .join('');
}

  // 📷 OCR Functions
  async function handleFileUpload(event) {
    const file = event.target.files[0];
    if (!file) return;

    // Validate file type
    if (!file.type.startsWith('image/')) {
      alert('Please select an image file.');
      return;
    }

    // Show modal
    ocrModal.style.display = 'block';
    ocrPreview.innerHTML = `<img src="${URL.createObjectURL(file)}" alt="Preview" style="max-width: 100%; max-height: 300px;">`;
    ocrResults.innerHTML = '<div class="loading">Processing image...</div>';

    try {
      // Extract text using OCR
      const formData = new FormData();
      formData.append('file', file);

      const response = await fetch('/ocr/extract-text/', {
        method: 'POST',
        body: formData
      });

      if (!response.ok) throw new Error(`Server returned ${response.status}`);
      const data = await response.json();

      if (data.success) {
        currentOcrText = data.extracted_text;
        currentOcrImage = file;
        ocrResults.innerHTML = `
          <div class="ocr-success">
            <h4>Extracted Text:</h4>
            <div class="extracted-text">${escapeHtml(data.extracted_text)}</div>
          </div>
        `;
      } else {
        ocrResults.innerHTML = `
          <div class="ocr-error">
            <h4>Error:</h4>
            <p>${data.error || 'Failed to extract text'}</p>
          </div>
        `;
      }
    } catch (error) {
      ocrResults.innerHTML = `
        <div class="ocr-error">
          <h4>Error:</h4>
          <p>Failed to process image: ${error.message}</p>
        </div>
      `;
    }
  }

  async function handleOcrQuery(query) {
    if (!currentOcrText) {
      alert('No OCR text available. Please upload an image first.');
      return;
    }

    // Add user message
    addUserMessage(`[Image Query] ${query}`);
    displayAllMessages();

    // Show loading
    const loadingMessage = renderLoadingMessage();
    chatMessages.innerHTML += loadingMessage;
    scrollToBottom();

    try {
      const response = await fetch('/ocr/query/', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({
          query: query,
          conversation_history: getHistory(),
          extracted_text: currentOcrText
        })
      });

      if (!response.ok) throw new Error(`Server returned ${response.status}`);
      const data = await response.json();

      // Remove loading and add response
      chatMessages.innerHTML = chatMessages.innerHTML.replace(loadingMessage, '');
      addAssistantMessage(data.response, data.source);
      displayAllMessages();
      ocrModal.style.display = 'none';

    } catch (error) {
      chatMessages.innerHTML = chatMessages.innerHTML.replace(loadingMessage, '');
      addAssistantMessage(`Failed to get response: ${error.message}`, 'Error');
      displayAllMessages();
    }
  }

  // 🔍 Query handler
  async function handleQuery() {
    const query = queryInput.value.trim();
    if (!query) return;

    // Add user message to conversation
    addUserMessage(query);
    displayAllMessages();

    // Clear input
    queryInput.value = '';

    // Show loading message
    const loadingMessage = renderLoadingMessage();
    chatMessages.innerHTML += loadingMessage;
    scrollToBottom();

    try {
      // Send conversation history to backend
      const response = await fetch('/query/', {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({
          query: query,
          conversation_history: getHistory()
        })
      });

      if (!response.ok) throw new Error(`Server returned ${response.status}`);
      const data = await response.json();

      // Remove loading message and add assistant response
      chatMessages.innerHTML = chatMessages.innerHTML.replace(loadingMessage, '');
      addAssistantMessage(data.response, data.source);
      displayAllMessages();

    } catch (err) {
      // Remove loading message and add error message
      chatMessages.innerHTML = chatMessages.innerHTML.replace(loadingMessage, '');
      addAssistantMessage(`Failed to get response: ${err.message}`, 'Error');
      displayAllMessages();
    }
}

  // 🔗 Event listeners
  askButton.addEventListener('click', handleQuery);
  queryInput.addEventListener('keypress', e => {
    if (e.key === 'Enter') handleQuery();
  });

  // OCR functionality
  attachFileBtn.addEventListener('click', () => {
    fileInput.click();
  });

  fileInput.addEventListener('change', handleFileUpload);

  closeOcrModal.addEventListener('click', () => {
    ocrModal.style.display = 'none';
  });

  useOcrText.addEventListener('click', () => {
    queryInput.value = currentOcrText;
    ocrModal.style.display = 'none';
    queryInput.focus();
  });

  queryWithOcr.addEventListener('click', () => {
    const query = prompt('Enter your question about the image:');
    if (query) {
      handleOcrQuery(query);
    }
  });

  cancelOcr.addEventListener('click', () => {
    ocrModal.style.display = 'none';
  });

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

  // Auto-resize input
  queryInput.addEventListener('input', () => {
    queryInput.style.height = 'auto';
    queryInput.style.height = queryInput.scrollHeight + 'px';
  });

  // Dropdown menu functionality
  const optionsBtn = document.getElementById('optionsBtn');
  const optionsMenu = document.getElementById('optionsMenu');
  const downloadTxtBtn = document.getElementById('downloadTxt');
  const downloadPdfBtn = document.getElementById('downloadPdf');
  const clearChatBtn = document.getElementById('clearChat');

  // Toggle dropdown menu
  optionsBtn.addEventListener('click', (e) => {
    e.stopPropagation();
    optionsMenu.classList.toggle('show');
  });

  // Close dropdown when clicking outside
  document.addEventListener('click', (e) => {
    if (!optionsBtn.contains(e.target) && !optionsMenu.contains(e.target)) {
      optionsMenu.classList.remove('show');
    }
  });

  // Download as TXT
  downloadTxtBtn.addEventListener('click', () => {
    downloadChatAsTxt();
    optionsMenu.classList.remove('show');
  });

  // Download as PDF
  downloadPdfBtn.addEventListener('click', () => {
    downloadChatAsPdf();
    optionsMenu.classList.remove('show');
  });

  // Clear chat
  clearChatBtn.addEventListener('click', () => {
    clearConversation();
    displayAllMessages();
    optionsMenu.classList.remove('show');
  });

  // Download functions
  function downloadChatAsTxt() {
    if (conversationHistory.length === 0) {
      alert('No conversation to download');
      return;
    }

    let content = 'Corex Chat History\n';
    content += '='.repeat(50) + '\n\n';
    
    conversationHistory.forEach((msg, index) => {
      const timestamp = new Date(msg.timestamp).toLocaleString();
      const role = msg.role === 'user' ? 'You' : 'Corex';
      const source = msg.source ? ` (${msg.source})` : '';
      
      content += `[${timestamp}] ${role}${source}:\n`;
      content += msg.content + '\n\n';
    });

    const blob = new Blob([content], { type: 'text/plain' });
    const url = URL.createObjectURL(blob);
    const a = document.createElement('a');
    a.href = url;
    a.download = `corex-chat-${new Date().toISOString().split('T')[0]}.txt`;
    document.body.appendChild(a);
    a.click();
    document.body.removeChild(a);
    URL.revokeObjectURL(url);
  }

  function downloadChatAsPdf() {
    if (conversationHistory.length === 0) {
      alert('No conversation to download');
      return;
    }

    try {
      const { jsPDF } = window.jspdf;
      const doc = new jsPDF();
      
      // Set up the document
      let yPosition = 20;
      const pageHeight = doc.internal.pageSize.height;
      const pageWidth = doc.internal.pageSize.width;
      const margin = 20;
      const maxWidth = pageWidth - (margin * 2);
      
      // Helper function to add text with word wrapping
      function addTextWithWrap(text, x, y, maxWidth, fontSize = 10) {
        doc.setFontSize(fontSize);
        const lines = doc.splitTextToSize(text, maxWidth);
        doc.text(lines, x, y);
        return y + (lines.length * (fontSize * 0.4));
      }
      
      // Helper function to check if we need a new page
      function checkNewPage(requiredSpace) {
        if (yPosition + requiredSpace > pageHeight - 20) {
          doc.addPage();
          yPosition = 20;
          return true;
        }
        return false;
      }
      
      // Title
      doc.setFontSize(16);
      doc.setFont(undefined, 'bold');
      doc.text('Corex Chat History', pageWidth / 2, yPosition, { align: 'center' });
      yPosition += 10;
      
      // Date
      doc.setFontSize(10);
      doc.setFont(undefined, 'normal');
      doc.text(`Generated on: ${new Date().toLocaleString()}`, pageWidth / 2, yPosition, { align: 'center' });
      yPosition += 15;
      
      // Add a line
      doc.line(margin, yPosition, pageWidth - margin, yPosition);
      yPosition += 10;
      
      // Process each message
      conversationHistory.forEach((msg, index) => {
        const timestamp = new Date(msg.timestamp).toLocaleString();
        const role = msg.role === 'user' ? 'You' : 'Corex';
        const source = msg.source ? ` (${msg.source})` : '';
        
        // Check if we need a new page for this message
        const messageText = `[${timestamp}] ${role}${source}:\n${msg.content}`;
        const estimatedHeight = (messageText.split('\n').length * 4) + 10;
        
        if (checkNewPage(estimatedHeight)) {
          // Add a continuation marker
          doc.setFontSize(8);
          doc.text('...continued from previous page...', margin, yPosition);
          yPosition += 5;
        }
        
        // Message header
        doc.setFontSize(10);
        doc.setFont(undefined, 'bold');
        yPosition = addTextWithWrap(`[${timestamp}] ${role}${source}:`, margin, yPosition, maxWidth, 10);
        
        // Message content
        doc.setFont(undefined, 'normal');
        yPosition = addTextWithWrap(msg.content, margin + 5, yPosition, maxWidth - 5, 9);
        
        // Add some space between messages
        yPosition += 8;
        
        // Add a subtle line between messages (except for the last one)
        if (index < conversationHistory.length - 1) {
          doc.setDrawColor(200, 200, 200);
          doc.line(margin, yPosition, pageWidth - margin, yPosition);
          yPosition += 5;
        }
      });
      
      // Save the PDF
      const fileName = `corex-chat-${new Date().toISOString().split('T')[0]}.pdf`;
      doc.save(fileName);
      
    } catch (error) {
      console.error('Error generating PDF:', error);
      alert('Error generating PDF. Please try downloading as TXT instead.');
    }
  }

  // Scroll to bottom when new messages arrive
  const observer = new MutationObserver(() => {
    scrollToBottom();
  });
  observer.observe(chatMessages, { childList: true, subtree: true });

  // Initialize
  displayAllMessages();
});