app.py

import gradio as gr
import tempfile
import os
import fitz  # PyMuPDF
import uuid
import shutil
from pymilvus import MilvusClient
import json
import sqlite3
from datetime import datetime
import hashlib
import bcrypt
import re
from typing import List, Dict, Tuple, Optional
import threading
import requests
import base64
from PIL import Image
import io

from middleware import Middleware
from rag import Rag
from pathlib import Path
import subprocess

# importing necessary functions from dotenv library
from dotenv import load_dotenv, dotenv_values 
import dotenv
import platform
import time
from pptxtopdf import convert

# Import libraries for DOC and Excel export
try:
    from docx import Document
    from docx.shared import Inches, Pt
    from docx.enum.text import WD_ALIGN_PARAGRAPH
    from docx.enum.style import WD_STYLE_TYPE
    from docx.oxml.shared import OxmlElement, qn
    from docx.oxml.ns import nsdecls
    from docx.oxml import parse_xml
    DOCX_AVAILABLE = True
except ImportError:
    DOCX_AVAILABLE = False
    print("Warning: python-docx not available. DOC export will be disabled.")

try:
    import openpyxl
    from openpyxl import Workbook
    from openpyxl.styles import Font, PatternFill, Alignment, Border, Side
    from openpyxl.chart import BarChart, LineChart, PieChart, Reference
    from openpyxl.utils.dataframe import dataframe_to_rows
    import pandas as pd
    EXCEL_AVAILABLE = True
except ImportError:
    EXCEL_AVAILABLE = False
    print("Warning: openpyxl/pandas not available. Excel export will be disabled.")

# loading variables from .env file
dotenv_file = dotenv.find_dotenv()
dotenv.load_dotenv(dotenv_file)

#kickstart docker and ollama servers

rag = Rag()

# Database for user management and chat history
class DatabaseManager:
    def __init__(self, db_path="app_database.db"):
        self.db_path = db_path
        self.init_database()
    
    def init_database(self):
        """Initialize database tables"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # Users table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                password_hash TEXT NOT NULL,
                team TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        ''')
        

        
        # Document collections table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS document_collections (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                collection_name TEXT UNIQUE NOT NULL,
                team TEXT NOT NULL,
                uploaded_by INTEGER,
                upload_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                file_count INTEGER DEFAULT 0,
                FOREIGN KEY (uploaded_by) REFERENCES users (id)
            )
        ''')
        
        conn.commit()
        conn.close()
    
    def create_user(self, username: str, password: str, team: str) -> bool:
        """Create a new user"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Hash password
            password_hash = bcrypt.hashpw(password.encode('utf-8'), bcrypt.gensalt())
            
            cursor.execute(
                'INSERT INTO users (username, password_hash, team) VALUES (?, ?, ?)',
                (username, password_hash.decode('utf-8'), team)
            )
            conn.commit()
            conn.close()
            return True
        except sqlite3.IntegrityError:
            return False
    
    def authenticate_user(self, username: str, password: str) -> Optional[Dict]:
        """Authenticate user and return user info"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('SELECT id, username, password_hash, team FROM users WHERE username = ?', (username,))
            user = cursor.fetchone()
            conn.close()
            
            if user and bcrypt.checkpw(password.encode('utf-8'), user[2].encode('utf-8')):
                return {
                    'id': user[0],
                    'username': user[1],
                    'team': user[3]
                }
            return None
        except Exception as e:
            print(f"Authentication error: {e}")
            return None
    

    
    def save_document_collection(self, collection_name: str, team: str, user_id: int, file_count: int):
        """Save document collection info"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute(
                'INSERT OR REPLACE INTO document_collections (collection_name, team, uploaded_by, file_count) VALUES (?, ?, ?, ?)',
                (collection_name, team, user_id, file_count)
            )
            conn.commit()
            conn.close()
        except Exception as e:
            print(f"Error saving document collection: {e}")
    
    def get_team_collections(self, team: str) -> List[str]:
        """Get all collections for a team"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('SELECT collection_name FROM document_collections WHERE team = ?', (team,))
            collections = [row[0] for row in cursor.fetchall()]
            conn.close()
            return collections
        except Exception as e:
            print(f"Error getting team collections: {e}")
            return []
    


# User session management
class SessionManager:
    def __init__(self):
        self.active_sessions = {}
        self.session_lock = threading.Lock()
    
    def create_session(self, user_info: Dict) -> str:
        """Create a new session for user"""
        session_id = str(uuid.uuid4())
        with self.session_lock:
            self.active_sessions[session_id] = {
                'user_info': user_info,
                'created_at': datetime.now(),
                'last_activity': datetime.now()
            }
        return session_id
    
    def get_session(self, session_id: str) -> Optional[Dict]:
        """Get session info"""
        with self.session_lock:
            if session_id in self.active_sessions:
                self.active_sessions[session_id]['last_activity'] = datetime.now()
                return self.active_sessions[session_id]
        return None
    
    def remove_session(self, session_id: str):
        """Remove session"""
        with self.session_lock:
            if session_id in self.active_sessions:
                del self.active_sessions[session_id]

# Initialize managers
db_manager = DatabaseManager()
session_manager = SessionManager()

# Create default users if they don't exist
def create_default_users():
    """Create default team users"""
    teams = ["Team_A", "Team_B"]
    for team in teams:
        username = f"admin_{team.lower()}"
        password = f"admin123_{team.lower()}"
        if not db_manager.authenticate_user(username, password):
            db_manager.create_user(username, password, team)
            print(f"Created default user: {username} for {team}")

create_default_users()


def start_services():
    # --- Docker Desktop (Windows Only) ---
    if platform.system() == "Windows":
        def is_docker_desktop_running():
            try:
                # Check if "Docker Desktop.exe" is in the task list.
                result = subprocess.run(
                    ["tasklist", "/FI", "IMAGENAME eq Docker Desktop.exe"],
                    stdout=subprocess.PIPE, stderr=subprocess.PIPE
                )
                return "Docker Desktop.exe" in result.stdout.decode()
            except Exception as e:
                print("Error checking Docker Desktop:", e)
                return False

        def start_docker_desktop():
            # Adjust this path if your Docker Desktop executable is located elsewhere.
            docker_desktop_path = r"C:\Program Files\Docker\Docker\Docker Desktop.exe"
            if not os.path.exists(docker_desktop_path):
                print("Docker Desktop executable not found. Please verify the installation path.")
                return
            try:
                subprocess.Popen([docker_desktop_path], shell=True)
                print("Docker Desktop is starting...")
            except Exception as e:
                print("Error starting Docker Desktop:", e)

        if is_docker_desktop_running():
            print("Docker Desktop is already running.")
        else:
            print("Docker Desktop is not running. Starting it now...")
            start_docker_desktop()
            # Wait for Docker Desktop to initialize (adjust delay as needed)
            time.sleep(15)

    # --- Ollama Server Management ---
    def is_ollama_running():
        if platform.system() == "Windows":
            try:
                # Check for "ollama.exe" in the task list (adjust if the executable name differs)
                result = subprocess.run(
                    ['tasklist', '/FI', 'IMAGENAME eq ollama.exe'],
                    stdout=subprocess.PIPE, stderr=subprocess.PIPE
                )
                return "ollama.exe" in result.stdout.decode().lower()
            except Exception as e:
                print("Error checking Ollama on Windows:", e)
                return False
        else:
            try:
                result = subprocess.run(
                    ['pgrep', '-f', 'ollama'],
                    stdout=subprocess.PIPE, stderr=subprocess.PIPE
                )
                return result.returncode == 0
            except Exception as e:
                print("Error checking Ollama:", e)
                return False

    def start_ollama():
        if platform.system() == "Windows":
            try:
                subprocess.Popen(['ollama', 'serve'], shell=True)
                print("Ollama server started on Windows.")
            except Exception as e:
                print("Failed to start Ollama server on Windows:", e)
        else:
            try:
                subprocess.Popen(['ollama', 'serve'])
                print("Ollama server started.")
            except Exception as e:
                print("Failed to start Ollama server:", e)

    # Ollama is no longer used; replaced by Gemini API calls.
    # Skip Ollama server checks and startup.

    # --- Docker Containers Management ---
    def get_docker_containers():
        try:
            result = subprocess.run(
                ['docker', 'ps', '-aq'],
                stdout=subprocess.PIPE, stderr=subprocess.PIPE
            )
            if result.returncode != 0:
                print("Error retrieving Docker containers:", result.stderr.decode())
                return []
            return result.stdout.decode().splitlines()
        except Exception as e:
            print("Error retrieving Docker containers:", e)
            return []

    def get_running_docker_containers():
        try:
            result = subprocess.run(
                ['docker', 'ps', '-q'],
                stdout=subprocess.PIPE, stderr=subprocess.PIPE
            )
            if result.returncode != 0:
                print("Error retrieving running Docker containers:", result.stderr.decode())
                return []
            return result.stdout.decode().splitlines()
        except Exception as e:
            print("Error retrieving running Docker containers:", e)
            return []

    def start_docker_container(container_id):
        try:
            result = subprocess.run(
                ['docker', 'start', container_id],
                stdout=subprocess.PIPE, stderr=subprocess.PIPE
            )
            if result.returncode == 0:
                print(f"Started Docker container {container_id}.")
            else:
                print(f"Failed to start Docker container {container_id}: {result.stderr.decode()}")
        except Exception as e:
            print(f"Error starting Docker container {container_id}: {e}")

    all_containers = set(get_docker_containers())
    running_containers = set(get_running_docker_containers())
    stopped_containers = all_containers - running_containers

    if stopped_containers:
        print(f"Found {len(stopped_containers)} stopped Docker container(s). Starting them...")
        for container_id in stopped_containers:
            start_docker_container(container_id)
    else:
        print("All Docker containers are already running.")

    
# Skip Docker services when running on Hugging Face Spaces
if not os.getenv("SPACE_ID"):
    start_services()
else:
    print("Running on Hugging Face Spaces - skipping Docker services")

def generate_uuid(state):
    # Check if UUID already exists in session state
    if state["user_uuid"] is None:
        # Generate a new UUID if not already set
        state["user_uuid"] = str(uuid.uuid4())

    return state["user_uuid"]


class PDFSearchApp:
    def __init__(self):
        self.indexed_docs = {}
        self.current_pdf = None
        self.db_manager = db_manager
        self.session_manager = session_manager
        
    def upload_and_convert(self, state, files, max_pages, session_id=None, folder_name=None):
        """Upload and convert files with team-based organization"""

        if files is None:
            return "No file uploaded"
        
        try:
            # Get user info from session if available
            user_info = None
            team = "default"
            if session_id:
                session = self.session_manager.get_session(session_id)
                if session:
                    user_info = session['user_info']
                    team = user_info['team']
            
            total_pages = 0
            uploaded_files = []
            
            # Create team-specific folder if folder_name is provided
            if folder_name:
                folder_name = folder_name.replace(" ", "_").replace("-", "_")
                collection_name = f"{team}_{folder_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            else:
                collection_name = f"{team}_documents_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            
            for file in files[:]:
                    # Extract the last part of the path (file name)
                filename = os.path.basename(file.name)
                name, ext = os.path.splitext(filename)
                pdf_path = file.name
                
                # Convert PPT to PDF if needed
                if ext.lower() in [".ppt", ".pptx"]:
                    output_file = os.path.splitext(file.name)[0] + '.pdf'
                    output_directory = os.path.dirname(file.name)
                    outfile = os.path.join(output_directory, output_file)
                    convert(file.name, outfile)
                    pdf_path = outfile
                    name = os.path.basename(outfile)
                    name, ext = os.path.splitext(name)
                
                # Create unique document ID
                doc_id = f"{collection_name}_{name.replace(' ', '_').replace('-', '_')}"
                
                print(f"Uploading file: {doc_id}")
                middleware = Middleware(collection_name, create_collection=True)
                
                pages = middleware.index(pdf_path, id=doc_id, max_pages=max_pages)
                total_pages += len(pages) if pages else 0
                uploaded_files.append(doc_id)
                
                self.indexed_docs[doc_id] = True
            
            # Save collection info to database
            if user_info:
                self.db_manager.save_document_collection(
                    collection_name, 
                    team, 
                    user_info['id'], 
                    len(uploaded_files)
                )
            
            return f"Uploaded {len(uploaded_files)} files with {total_pages} total pages to collection: {collection_name}"
            
        except Exception as e:
            return f"Error processing files: {str(e)}"
    

    def display_file_list(text):
        try:
        # Retrieve all entries in the specified directory
            directory_path = "pages"
            current_working_directory = os.getcwd()
            directory_path = os.path.join(current_working_directory, directory_path)
            entries = os.listdir(directory_path)
            # Filter out entries that are directories
            directories = [entry for entry in entries if os.path.isdir(os.path.join(directory_path, entry))]
            return directories
        except FileNotFoundError:
            return f"The directory {directory_path} does not exist."
        except PermissionError:
            return f"Permission denied to access {directory_path}."
        except Exception as e:
            return str(e)

    
    def search_documents(self, state, query, num_results, session_id=None):
        print(f"Searching for query: {query}")
       
        if not query:
            print("Please enter a search query")
            return "Please enter a search query", "--", "Please enter a search query", [], None
            
        try:
            # Get user info from session if available
            user_info = None
            if session_id:
                session = self.session_manager.get_session(session_id)
                if session:
                    user_info = session['user_info']
            
            middleware = Middleware("test", create_collection=False)
            
            # Enhanced multi-page retrieval with vision-guided chunking approach
            # Get more results than requested to allow for intelligent filtering
            # Request 3x the number of results for better selection
            search_results = middleware.search([query], topk=max(num_results * 3, 20))[0]
            
            # Debug: Log the number of results retrieved
            print(f"🔍 Retrieved {len(search_results)} total results from search")
            if len(search_results) > 0:
                print(f"🔍 Top result score: {search_results[0][0]:.3f}")
                print(f"🔍 Bottom result score: {search_results[-1][0]:.3f}")
            
            if not search_results:
                return "No search results found", "--", "No search results found for your query", [], None
            
            # Implement intelligent multi-page selection based on research
            selected_results = self._select_relevant_pages(search_results, query, num_results)
            
            # Process selected results
            cited_pages = []
            img_paths = []
            all_paths = []
            page_scores = []
            
            print(f"📄 Processing {len(selected_results)} selected results...")
            
            for i, (score, page_num, coll_num) in enumerate(selected_results):
                # Convert 0-based page number to 1-based for file naming
                display_page_num = page_num + 1
                img_path = f"pages/{coll_num}/page_{display_page_num}.png"
                path = f"pages/{coll_num}/page_{display_page_num}"

                if os.path.exists(img_path):
                    img_paths.append(img_path)
                    all_paths.append(path)
                    page_scores.append(score)
                    cited_pages.append(f"Page {display_page_num} from {coll_num}")
                    print(f"✅ Retrieved page {i+1}: {img_path} (Score: {score:.3f})")
                else:
                    print(f"❌ Image file not found: {img_path}")
            
            print(f"📊 Final count: {len(img_paths)} valid pages out of {len(selected_results)} selected")
            
            if not img_paths:
                return "No valid image files found", "--", "Error: No valid image files found for the search results", [], None
            
            # Generate RAG response with multiple pages using enhanced approach
            rag_response, csv_filepath, doc_filepath, excel_filepath = self._generate_multi_page_response(query, img_paths, cited_pages, page_scores)
            

            
            # Prepare downloads
            csv_download = self._prepare_csv_download(csv_filepath)
            doc_download = self._prepare_doc_download(doc_filepath)
            excel_download = self._prepare_excel_download(excel_filepath)
            
            # Return multiple images if available, otherwise single image
            if len(img_paths) > 1:
                # Format for Gallery component: list of (image_path, caption) tuples
                # Extract page numbers from cited_pages for accurate captions
                gallery_images = []
                for i, img_path in enumerate(img_paths):
                    # Extract page number from cited_pages
                    page_info = cited_pages[i].split(" from ")[0]  # "Page X"
                    page_num = page_info.split("Page ")[1]  # "X"
                    gallery_images.append((img_path, f"Page {page_num}"))
                return ", ".join(all_paths), gallery_images, rag_response, cited_pages, csv_download, doc_download, excel_download
            else:
                # Single image format
                page_info = cited_pages[0].split(" from ")[0]  # "Page X"
                page_num = page_info.split("Page ")[1]  # "X"
                return all_paths[0], [(img_paths[0], f"Page {page_num}")], rag_response, cited_pages, csv_download, doc_download, excel_download
            
        except Exception as e:
            error_msg = f"Error during search: {str(e)}"
            return error_msg, "--", error_msg, [], None, None, None, None
    
    def _select_relevant_pages(self, search_results, query, num_results):
        """
        Intelligent page selection using vision-guided chunking principles
        Based on research from M3DocRAG and multi-modal retrieval models
        """
        if len(search_results) <= num_results:
            return search_results
        
        # Detect if query needs multiple pages
        multi_page_keywords = [
            'compare', 'difference', 'similarities', 'both', 'multiple', 'various',
            'different', 'types', 'kinds', 'categories', 'procedures', 'methods',
            'approaches', 'techniques', 'safety', 'protocols', 'guidelines',
            'overview', 'summary', 'comprehensive', 'complete', 'all', 'everything'
        ]
        
        query_lower = query.lower()
        needs_multiple_pages = any(keyword in query_lower for keyword in multi_page_keywords)
        
        # Sort by relevance score
        sorted_results = sorted(search_results, key=lambda x: x[0], reverse=True)
        
        # CRITICAL FIX: Ensure we return exactly the number of pages requested
        # This addresses the ColPali retrieval configuration issue mentioned in research
        
        # Strategy 1: Include highest scoring result from each collection (diversity)
        selected = []
        seen_collections = set()
        
        # First pass: get one page from each collection for diversity
        for score, page_num, coll_num in sorted_results:
            if coll_num not in seen_collections and len(selected) < min(num_results // 2, len(search_results)):
                selected.append((score, page_num, coll_num))
                seen_collections.add(coll_num)
        
        # Strategy 2: Fill remaining slots with highest scoring results
        for score, page_num, coll_num in sorted_results:
            if (score, page_num, coll_num) not in selected and len(selected) < num_results:
                selected.append((score, page_num, coll_num))
        
        # Strategy 3: If we still don't have enough, add more from any collection
        if len(selected) < num_results:
            for score, page_num, coll_num in sorted_results:
                if (score, page_num, coll_num) not in selected and len(selected) < num_results:
                    selected.append((score, page_num, coll_num))
        
        # Strategy 4: If we have too many, trim to exact number requested
        if len(selected) > num_results:
            selected = selected[:num_results]
        
        # Strategy 5: If we have too few, add more from the sorted results
        if len(selected) < num_results and len(sorted_results) >= num_results:
            for score, page_num, coll_num in sorted_results:
                if (score, page_num, coll_num) not in selected and len(selected) < num_results:
                    selected.append((score, page_num, coll_num))
        
        # Sort selected results by score for consistency
        selected.sort(key=lambda x: x[0], reverse=True)
        
        print(f"Requested {num_results} pages, selected {len(selected)} pages from {len(seen_collections)} collections")
        
        # Final verification: ensure we return exactly the requested number
        if len(selected) != num_results:
            print(f"⚠️ Warning: Requested {num_results} pages but selected {len(selected)} pages")
            if len(selected) < num_results and len(sorted_results) >= num_results:
                # Add more pages to reach the target
                for score, page_num, coll_num in sorted_results:
                    if (score, page_num, coll_num) not in selected and len(selected) < num_results:
                        selected.append((score, page_num, coll_num))
                print(f"Added more pages to reach target: {len(selected)} pages")
        
        return selected
    
    def _optimize_consecutive_pages(self, selected, all_results, target_count=None):
        """
        Optimize selection to include consecutive pages when beneficial
        """
        # Group by collection
        collection_pages = {}
        for score, page_num, coll_num in selected:
            if coll_num not in collection_pages:
                collection_pages[coll_num] = []
            collection_pages[coll_num].append((score, page_num, coll_num))
        
        optimized = []
        for coll_num, pages in collection_pages.items():
            if len(pages) > 1:
                # Check if pages are consecutive
                page_nums = [p[1] for p in pages]
                page_nums.sort()
                
                # If pages are consecutive, add any missing pages in between
                if max(page_nums) - min(page_nums) == len(page_nums) - 1:
                    # Find all pages in this range from all_results
                    for score, page_num, coll in all_results:
                        if (coll == coll_num and 
                            min(page_nums) <= page_num <= max(page_nums) and
                            (score, page_num, coll) not in optimized):
                            optimized.append((score, page_num, coll))
                else:
                    optimized.extend(pages)
            else:
                optimized.extend(pages)
        
        # Ensure we maintain the target count if specified
        if target_count and len(optimized) != target_count:
            if len(optimized) > target_count:
                # Trim to target count, keeping highest scoring
                optimized.sort(key=lambda x: x[0], reverse=True)
                optimized = optimized[:target_count]
            elif len(optimized) < target_count:
                # Add more pages to reach target
                for score, page_num, coll in all_results:
                    if (score, page_num, coll) not in optimized and len(optimized) < target_count:
                        optimized.append((score, page_num, coll))
        
        return optimized
    
    def _generate_comprehensive_analysis(self, query, cited_pages, page_scores):
        """
        Generate comprehensive analysis section based on research strategies
        Implements hierarchical retrieval insights and cross-reference analysis
        """
        try:
            # Analyze query complexity and information needs
            query_lower = query.lower()
            
            # Determine query type for targeted analysis
            query_types = []
            if any(word in query_lower for word in ['compare', 'difference', 'similarities', 'versus']):
                query_types.append("Comparative Analysis")
            if any(word in query_lower for word in ['procedure', 'method', 'how to', 'steps']):
                query_types.append("Procedural Information")
            if any(word in query_lower for word in ['safety', 'warning', 'danger', 'risk']):
                query_types.append("Safety Information")
            if any(word in query_lower for word in ['specification', 'technical', 'measurement', 'data']):
                query_types.append("Technical Specifications")
            if any(word in query_lower for word in ['overview', 'summary', 'comprehensive', 'complete']):
                query_types.append("Comprehensive Overview")
            if any(word in query_lower for word in ['table', 'csv', 'spreadsheet', 'data', 'list', 'chart']):
                query_types.append("Tabular Data Request")
            
            # Calculate information quality metrics
            avg_score = sum(page_scores) / len(page_scores) if page_scores else 0
            score_variance = sum((score - avg_score) ** 2 for score in page_scores) / len(page_scores) if page_scores else 0
            
            # Generate analysis insights
            analysis = f"""
🔬 **Comprehensive Analysis & Insights**:

📝 **Query Analysis**:
• Query Type: {', '.join(query_types) if query_types else 'General Information'}
• Information Complexity: {'High' if len(cited_pages) > 3 else 'Medium' if len(cited_pages) > 1 else 'Low'}
• Cross-Reference Depth: {'Excellent' if len(set([p.split(' from ')[1].split(' (')[0] for p in cited_pages])) > 2 else 'Good' if len(set([p.split(' from ')[1].split(' (')[0] for p in cited_pages])) > 1 else 'Limited'}

📊 **Information Quality Assessment**:
• Average Relevance: {avg_score:.3f} ({'Excellent' if avg_score > 0.9 else 'Very Good' if avg_score > 0.8 else 'Good' if avg_score > 0.7 else 'Moderate' if avg_score > 0.6 else 'Basic'})
• Information Consistency: {'High' if score_variance < 0.1 else 'Moderate' if score_variance < 0.2 else 'Variable'}
• Source Reliability: {'High' if avg_score > 0.8 and len(cited_pages) > 2 else 'Moderate' if avg_score > 0.6 else 'Requires Verification'}

🎯 **Information Coverage Analysis**:
• Primary Information: {'Comprehensive' if any('primary' in p.lower() or 'main' in p.lower() for p in cited_pages) else 'Standard'}
• Supporting Details: {'Extensive' if len(cited_pages) > 3 else 'Adequate' if len(cited_pages) > 1 else 'Basic'}
• Technical Depth: {'High' if any('technical' in p.lower() or 'specification' in p.lower() for p in cited_pages) else 'Standard'}

💡 **Strategic Insights**:
• Information Gaps: {'Minimal' if avg_score > 0.8 and len(cited_pages) > 3 else 'Moderate' if avg_score > 0.6 else 'Significant - consider additional sources'}
• Cross-Validation: {'Strong' if len(set([p.split(' from ')[1].split(' (')[0] for p in cited_pages])) > 1 else 'Limited to single source'}
• Practical Applicability: {'High' if any('procedure' in p.lower() or 'method' in p.lower() for p in cited_pages) else 'Moderate'}

🔍 **Recommendations for Further Research**:
• {'Consider additional technical specifications' if not any('technical' in p.lower() for p in cited_pages) else 'Technical coverage adequate'}
• {'Seek safety guidelines and warnings' if not any('safety' in p.lower() for p in cited_pages) else 'Safety information included'}
• {'Look for comparative analysis' if not any('compare' in p.lower() for p in cited_pages) else 'Comparative analysis available'}
"""
            
            return analysis
            
        except Exception as e:
            print(f"Error generating comprehensive analysis: {e}")
            return "🔬 **Analysis**: Comprehensive analysis of retrieved information completed."
    

    
    def _detect_table_request(self, query):
        """
        Detect if the user is requesting tabular data
        """
        query_lower = query.lower()
        table_keywords = [
            'table', 'csv', 'spreadsheet', 'data table', 'list', 'chart',
            'tabular', 'matrix', 'grid', 'dataset', 'data set',
            'show me a table', 'create a table', 'generate table',
            'in table format', 'as a table', 'tabular format'
        ]
        
        return any(keyword in query_lower for keyword in table_keywords)
    
    def _detect_report_request(self, query):
        """
        Detect if the user is requesting a comprehensive report
        """
        query_lower = query.lower()
        report_keywords = [
            'report', 'comprehensive report', 'detailed report', 'full report',
            'complete report', 'comprehensive analysis', 'detailed analysis',
            'full analysis', 'complete analysis', 'comprehensive overview',
            'detailed overview', 'full overview', 'complete overview',
            'comprehensive summary', 'detailed summary', 'full summary',
            'complete summary', 'comprehensive document', 'detailed document',
            'full document', 'complete document', 'comprehensive review',
            'detailed review', 'full review', 'complete review',
            'export report', 'generate report', 'create report',
            'doc format', 'word document', 'word doc', 'document format'
        ]
        
        return any(keyword in query_lower for keyword in report_keywords)
    
    def _detect_chart_request(self, query):
        """
        Detect if the user is requesting charts, graphs, or visualizations
        """
        query_lower = query.lower()
        chart_keywords = [
            'chart', 'graph', 'bar chart', 'line chart', 'pie chart',
            'bar graph', 'line graph', 'pie graph', 'histogram',
            'scatter plot', 'scatter chart', 'area chart', 'column chart',
            'visualization', 'visualize', 'plot', 'figure', 'diagram',
            'excel chart', 'excel graph', 'spreadsheet chart',
            'create chart', 'generate chart', 'make chart',
            'create graph', 'generate graph', 'make graph',
            'chart data', 'graph data', 'plot data', 'visualize data',
            'bar graph', 'line graph', 'pie graph', 'histogram',
            'scatter plot', 'area chart', 'column chart'
        ]
        
        return any(keyword in query_lower for keyword in chart_keywords)
    
    def _extract_custom_headers(self, query):
        """
        Extract custom headers from user query for both tables and charts
        Examples: 
        - "create table with columns: Name, Age, Department"
        - "create chart with headers: Threat Type, Frequency, Risk Level"
        - "excel export with columns: Category, Value, Description"
        """
        try:
            # Look for header specifications in the query
            header_patterns = [
                r'columns?:\s*([^,]+(?:,\s*[^,]+)*)',  # "columns: A, B, C"
                r'headers?:\s*([^,]+(?:,\s*[^,]+)*)',  # "headers: A, B, C"
                r'\bwith\s+columns?\s*([^,]+(?:,\s*[^,]+)*)',  # "with columns A, B, C"
                r'\bwith\s+headers?\s*([^,]+(?:,\s*[^,]+)*)',  # "with headers A, B, C"
                r'headers?\s*=\s*([^,]+(?:,\s*[^,]+)*)',  # "headers = A, B, C"
                r'format:\s*([^,]+(?:,\s*[^,]+)*)',  # "format: A, B, C"
                r'chart\s+headers?:\s*([^,]+(?:,\s*[^,]+)*)',  # "chart headers: A, B, C"
                r'excel\s+headers?:\s*([^,]+(?:,\s*[^,]+)*)',  # "excel headers: A, B, C"
                r'chart\s+with\s+headers?:\s*([^,]+(?:,\s*[^,]+)*)',  # "chart with headers: A, B, C"
                r'excel\s+with\s+headers?:\s*([^,]+(?:,\s*[^,]+)*)',  # "excel with headers: A, B, C"
            ]
            
            for pattern in header_patterns:
                match = re.search(pattern, query, re.IGNORECASE)
                if match:
                    headers_str = match.group(1)
                    # Split by comma and clean up
                    headers = [h.strip() for h in headers_str.split(',')]
                    # Remove empty headers
                    headers = [h for h in headers if h]
                    if headers:
                        print(f"📋 Custom headers detected: {headers}")
                        return headers
            
            return None
            
        except Exception as e:
            print(f"Error extracting custom headers: {e}")
            return None
    
    def _generate_csv_table_response(self, query, rag_response, cited_pages, page_scores):
        """
        Generate a CSV table response when user requests tabular data
        """
        try:
            # Extract custom headers from query if specified
            custom_headers = self._extract_custom_headers(query)
            
            # Extract structured data from the RAG response
            csv_data = self._extract_structured_data(rag_response, cited_pages, page_scores, custom_headers)
            
            if csv_data:
                # Format as CSV
                csv_content = self._format_as_csv(csv_data)
                
                # Generate a unique filename for the CSV
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                safe_query = "".join(c for c in query[:30] if c.isalnum() or c in (' ', '-', '_')).rstrip()
                safe_query = safe_query.replace(' ', '_')
                filename = f"table_{safe_query}_{timestamp}.csv"
                filepath = os.path.join("temp", filename)
                
                # Ensure temp directory exists
                os.makedirs("temp", exist_ok=True)
                
                # Save CSV file
                with open(filepath, 'w', encoding='utf-8') as f:
                    f.write(csv_content)
                
                # Create enhanced response with CSV and download link
                header_info = ""
                if custom_headers:
                    header_info = f"""
📋 **Custom Headers Applied**:
• Headers: {', '.join(custom_headers)}
• Data automatically mapped to your specified columns
"""
                
                table_response = f"""
{rag_response}

📊 **CSV Table Generated Successfully**:

```csv
{csv_content}
```

{header_info}

💾 **Download Options**:
• **Direct Download**: Click the download button below
• **Manual Copy**: Copy the CSV content above and save as .csv file

📋 **Table Information**:
• Rows: {len(csv_data) if csv_data else 0}
• Columns: {len(csv_data[0]) if csv_data and len(csv_data) > 0 else 0}
• Data Source: {len(cited_pages)} document pages
• Filename: {filename}
"""
                return table_response, filepath
            else:
                # Fallback if no structured data found
                header_suggestion = ""
                if custom_headers:
                    header_suggestion = f"""
📋 **Custom Headers Detected**: {', '.join(custom_headers)}
The system found your specified headers but couldn't extract matching data from the response.
"""
                
                fallback_response = f"""
{rag_response}

📊 **Table Request Detected**:
The system detected you requested tabular data, but the current response doesn't contain structured information suitable for a CSV table.

{header_suggestion}

💡 **Suggestions**:
• Try asking for specific data types (e.g., "list of safety procedures", "compare different methods")
• Request numerical data or comparisons
• Ask for categorized information
• Specify custom headers: "create table with columns: Name, Age, Department"
"""
                return fallback_response, None
                
        except Exception as e:
            print(f"Error generating CSV table response: {e}")
            return rag_response, None
    
    def _extract_structured_data(self, rag_response, cited_pages, page_scores, custom_headers=None):
        """
        Extract ANY structured data from RAG response - no predefined templates
        """
        try:
            lines = rag_response.split('\n')
            structured_data = []
            
            # If user specified custom headers, try to extract data that fits
            if custom_headers:
                headers = custom_headers
                structured_data = [headers]
                
                # Extract any data that could fit the headers
                data_rows = []
                
                # Look for any structured content in the response
                for line in lines:
                    line = line.strip()
                    if line and not line.startswith('#'):  # Skip markdown headers
                        # Try to extract meaningful data from each line
                        data_row = self._extract_data_from_line(line, headers)
                        if data_row:
                            data_rows.append(data_row)
                
                # If we found data, use it; otherwise create placeholder rows
                if data_rows:
                    structured_data.extend(data_rows)
                else:
                    # Create placeholder rows based on available content
                    for i, citation in enumerate(cited_pages):
                        row = self._create_placeholder_row(citation, headers, i)
                        structured_data.append(row)
                
                return structured_data
            
            # No custom headers - let's be smart about what we find
            else:
                # Look for any obvious table-like structures first
                table_data = self._find_table_structures(lines)
                if table_data:
                    return table_data
                
                # Look for any structured lists or data
                list_data = self._find_list_structures(lines)
                if list_data:
                    return list_data
                
                # Look for any key-value patterns
                kv_data = self._find_key_value_structures(lines)
                if kv_data:
                    return kv_data
                
                # Last resort: create a simple summary
                return self._create_summary_table(cited_pages)
            
        except Exception as e:
            print(f"Error extracting structured data: {e}")
            return None
    
    def _extract_data_from_line(self, line, headers):
        """Extract data from a line that could fit the specified headers"""
        try:
            # Remove common prefixes
            line = re.sub(r'^[\d•\-\.\s]+', '', line)
            
            # If we have multiple headers, try to split the line
            if len(headers) > 1:
                # Look for natural splits (commas, semicolons, etc.)
                if ',' in line:
                    parts = [p.strip() for p in line.split(',')]
                elif ';' in line:
                    parts = [p.strip() for p in line.split(';')]
                elif ' - ' in line:
                    parts = [p.strip() for p in line.split(' - ')]
                elif ':' in line:
                    parts = [p.strip() for p in line.split(':', 1)]
                else:
                    # Just put the whole line in the first column
                    parts = [line] + [''] * (len(headers) - 1)
                
                # Pad or truncate to match header count
                while len(parts) < len(headers):
                    parts.append('')
                return parts[:len(headers)]
            else:
                return [line]
                
        except Exception as e:
            print(f"Error extracting data from line: {e}")
            return None
    
    def _create_placeholder_row(self, citation, headers, index):
        """Create a placeholder row based on available data"""
        try:
            row = []
            for header in headers:
                header_lower = header.lower()
                
                if 'page' in header_lower or 'number' in header_lower:
                    page_num = citation.split('Page ')[1].split(' from')[0] if 'Page ' in citation else str(index + 1)
                    row.append(page_num)
                elif 'collection' in header_lower or 'source' in header_lower or 'document' in header_lower:
                    collection = citation.split(' from ')[1] if ' from ' in citation else 'Unknown'
                    row.append(collection)
                elif 'content' in header_lower or 'description' in header_lower or 'summary' in header_lower:
                    row.append(f"Content from {citation}")
                else:
                    # For unknown headers, try to extract something relevant
                    if 'page' in citation:
                        row.append(citation)
                    else:
                        row.append('')
            
            return row
            
        except Exception as e:
            print(f"Error creating placeholder row: {e}")
            return [''] * len(headers)
    
    def _find_table_structures(self, lines):
        """Find any table-like structures in the text"""
        try:
            table_lines = []
            for line in lines:
                line = line.strip()
                # Look for lines with multiple columns (separated by |, tabs, or multiple spaces)
                if '|' in line or '\t' in line or re.search(r'\s{3,}', line):
                    table_lines.append(line)
            
            if table_lines:
                # Try to determine headers from the first line
                first_line = table_lines[0]
                if '|' in first_line:
                    headers = [h.strip() for h in first_line.split('|')]
                else:
                    headers = re.split(r'\s{3,}', first_line)
                
                structured_data = [headers]
                
                # Process remaining lines
                for line in table_lines[1:]:
                    if '|' in line:
                        columns = [col.strip() for col in line.split('|')]
                    else:
                        columns = re.split(r'\s{3,}', line)
                    
                    if len(columns) >= 2:
                        structured_data.append(columns)
                
                return structured_data
            
            return None
            
        except Exception as e:
            print(f"Error finding table structures: {e}")
            return None
    
    def _find_list_structures(self, lines):
        """Find any list-like structures in the text"""
        try:
            items = []
            for line in lines:
                line = line.strip()
                # Remove common list markers
                if re.match(r'^[\d•\-\.]+', line):
                    item = re.sub(r'^[\d•\-\.\s]+', '', line)
                    if item:
                        items.append(item)
            
            if items:
                # Create a simple list structure
                structured_data = [['Item', 'Description']]
                for i, item in enumerate(items, 1):
                    structured_data.append([str(i), item])
                
                return structured_data
            
            return None
            
        except Exception as e:
            print(f"Error finding list structures: {e}")
            return None
    
    def _find_key_value_structures(self, lines):
        """Find any key-value structures in the text"""
        try:
            kv_pairs = []
            for line in lines:
                line = line.strip()
                # Look for key: value patterns
                if re.match(r'^[A-Za-z\s]+:\s+', line):
                    kv_pairs.append(line)
            
            if kv_pairs:
                structured_data = [['Property', 'Value']]
                for pair in kv_pairs:
                    if ':' in pair:
                        key, value = pair.split(':', 1)
                        structured_data.append([key.strip(), value.strip()])
                
                return structured_data
            
            return None
            
        except Exception as e:
            print(f"Error finding key-value structures: {e}")
            return None
    
    def _create_summary_table(self, cited_pages):
        """Create a simple summary table as last resort"""
        try:
            structured_data = [['Page', 'Collection', 'Content']]
            for i, citation in enumerate(cited_pages):
                collection = citation.split(' from ')[1] if ' from ' in citation else 'Unknown'
                page_num = citation.split('Page ')[1].split(' from')[0] if 'Page ' in citation else str(i+1)
                structured_data.append([page_num, collection, f"Content from {citation}"])
            
            return structured_data
            
        except Exception as e:
            print(f"Error creating summary table: {e}")
            return None
            
        except Exception as e:
            print(f"Error extracting structured data: {e}")
            return None
    
    def _format_as_csv(self, data):
        """
        Format structured data as CSV
        """
        try:
            csv_lines = []
            for row in data:
                # Escape commas and quotes in CSV
                escaped_row = []
                for cell in row:
                    cell_str = str(cell)
                    if ',' in cell_str or '"' in cell_str or '\n' in cell_str:
                        # Escape quotes and wrap in quotes
                        cell_str = f'"{cell_str.replace('"', '""')}"'
                    escaped_row.append(cell_str)
                csv_lines.append(','.join(escaped_row))
            
            return '\n'.join(csv_lines)
            
        except Exception as e:
            print(f"Error formatting CSV: {e}")
            return "Error,Generating,CSV,Format"
    
    def _prepare_csv_download(self, csv_filepath):
        """
        Prepare CSV file for download in Gradio
        """
        if csv_filepath and os.path.exists(csv_filepath):
            return csv_filepath
        else:
            return None
    
    def _generate_comprehensive_doc_report(self, query, rag_response, cited_pages, page_scores, user_info=None):
        """
        Generate a comprehensive DOC report with proper formatting and structure
        """
        if not DOCX_AVAILABLE:
            return None, "DOC export not available - python-docx library not installed"
        
        try:
            print("📄 [REPORT] Generating comprehensive DOC report...")
            
            # Create a new Document
            doc = Document()
            
            # Set up document styles
            self._setup_document_styles(doc)
            
            # Add title page
            self._add_title_page(doc, query, user_info)
            
            # Add executive summary
            self._add_executive_summary(doc, query, rag_response)
            
            # Add detailed analysis
            self._add_detailed_analysis(doc, rag_response, cited_pages, page_scores)
            
            # Add methodology
            self._add_methodology_section(doc, cited_pages, page_scores)
            
            # Add findings and conclusions
            self._add_findings_conclusions(doc, rag_response, cited_pages)
            
            # Add appendices
            self._add_appendices(doc, cited_pages, page_scores)
            
            # Generate unique filename
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            safe_query = "".join(c for c in query[:30] if c.isalnum() or c in (' ', '-', '_')).rstrip()
            safe_query = safe_query.replace(' ', '_')
            filename = f"comprehensive_report_{safe_query}_{timestamp}.docx"
            filepath = os.path.join("temp", filename)
            
            # Ensure temp directory exists
            os.makedirs("temp", exist_ok=True)
            
            # Save the document
            doc.save(filepath)
            
            print(f"✅ [REPORT] Comprehensive DOC report generated: {filepath}")
            return filepath, None
            
        except Exception as e:
            error_msg = f"Error generating DOC report: {str(e)}"
            print(f"❌ [REPORT] {error_msg}")
            return None, error_msg
    
    def _setup_document_styles(self, doc):
        """Set up professional document styles"""
        try:
            # Import RGBColor for proper color handling
            from docx.shared import RGBColor
            
            # Title style
            title_style = doc.styles.add_style('CustomTitle', WD_STYLE_TYPE.PARAGRAPH)
            title_font = title_style.font
            title_font.name = 'Calibri'
            title_font.size = Pt(24)
            title_font.bold = True
            title_font.color.rgb = RGBColor(47, 84, 150)  # #2F5496
            
            # Heading 1 style
            h1_style = doc.styles.add_style('CustomHeading1', WD_STYLE_TYPE.PARAGRAPH)
            h1_font = h1_style.font
            h1_font.name = 'Calibri'
            h1_font.size = Pt(16)
            h1_font.bold = True
            h1_font.color.rgb = RGBColor(47, 84, 150)  # #2F5496
            
            # Heading 2 style
            h2_style = doc.styles.add_style('CustomHeading2', WD_STYLE_TYPE.PARAGRAPH)
            h2_font = h2_style.font
            h2_font.name = 'Calibri'
            h2_font.size = Pt(14)
            h2_font.bold = True
            h2_font.color.rgb = RGBColor(47, 84, 150)  # #2F5496
            
            # Body text style
            body_style = doc.styles.add_style('CustomBody', WD_STYLE_TYPE.PARAGRAPH)
            body_font = body_style.font
            body_font.name = 'Calibri'
            body_font.size = Pt(11)
            
        except Exception as e:
            print(f"Warning: Could not set up custom styles: {e}")
    
    def _add_title_page(self, doc, query, user_info):
        """Add professional title page for security analysis report"""
        try:
            # Import RGBColor for proper color handling
            from docx.shared import RGBColor
            
            # Title
            title = doc.add_paragraph()
            title.alignment = WD_ALIGN_PARAGRAPH.CENTER
            title_run = title.add_run("SECURITY THREAT ANALYSIS REPORT")
            title_run.font.name = 'Calibri'
            title_run.font.size = Pt(24)
            title_run.font.bold = True
            title_run.font.color.rgb = RGBColor(47, 84, 150)  # #2F5496
            
            # Subtitle
            subtitle = doc.add_paragraph()
            subtitle.alignment = WD_ALIGN_PARAGRAPH.CENTER
            subtitle_run = subtitle.add_run(f"Threat Intelligence Query: {query}")
            subtitle_run.font.name = 'Calibri'
            subtitle_run.font.size = Pt(14)
            subtitle_run.font.italic = True
            
            # Add spacing
            doc.add_paragraph()
            doc.add_paragraph()
            
            # Report classification
            classification = doc.add_paragraph()
            classification.alignment = WD_ALIGN_PARAGRAPH.CENTER
            classification_run = classification.add_run("SECURITY ANALYSIS & THREAT INTELLIGENCE")
            classification_run.font.name = 'Calibri'
            classification_run.font.size = Pt(12)
            classification_run.font.bold = True
            classification_run.font.color.rgb = RGBColor(220, 53, 69)  # #dc3545
            
            # Report details
            details = doc.add_paragraph()
            details.alignment = WD_ALIGN_PARAGRAPH.CENTER
            details_run = details.add_run(f"Generated on: {datetime.now().strftime('%B %d, %Y at %I:%M %p')}")
            details_run.font.name = 'Calibri'
            details_run.font.size = Pt(11)
            
            if user_info:
                user_details = doc.add_paragraph()
                user_details.alignment = WD_ALIGN_PARAGRAPH.CENTER
                user_run = user_details.add_run(f"Generated by: {user_info['username']} ({user_info['team']})")
                user_run.font.name = 'Calibri'
                user_run.font.size = Pt(11)
            
            # Add page break
            doc.add_page_break()
            
        except Exception as e:
            print(f"Warning: Could not add title page: {e}")
    
    def _add_executive_summary(self, doc, query, rag_response):
        """Add executive summary section aligned with security analysis framework"""
        try:
            # Import RGBColor for proper color handling
            from docx.shared import RGBColor
            
            # Section heading
            heading = doc.add_paragraph()
            heading_run = heading.add_run("EXECUTIVE SUMMARY")
            heading_run.font.name = 'Calibri'
            heading_run.font.size = Pt(16)
            heading_run.font.bold = True
            heading_run.font.color.rgb = RGBColor(47, 84, 150)  # #2F5496
            
            # Report purpose
            purpose = doc.add_paragraph()
            purpose_run = purpose.add_run("This security analysis report provides comprehensive threat assessment and operational insights based on the query: ")
            purpose_run.font.name = 'Calibri'
            purpose_run.font.size = Pt(11)
            
            # Query in bold
            query_text = doc.add_paragraph()
            query_run = query_text.add_run(f'"{query}"')
            query_run.font.name = 'Calibri'
            query_run.font.size = Pt(11)
            query_run.font.bold = True
            
            # Analysis framework overview
            framework_heading = doc.add_paragraph()
            framework_run = framework_heading.add_run("Analysis Framework:")
            framework_run.font.name = 'Calibri'
            framework_run.font.size = Pt(12)
            framework_run.font.bold = True
            
            # Framework components
            framework_components = [
                "• Fact-Finding & Contextualization: Background information and context development",
                "• Case Study Identification: Incident prevalence and TTP extraction",
                "• Analytical Assessment: Intent, motivation, and threat landscape evaluation",
                "• Operational Relevance: Ground-level actionable insights and recommendations"
            ]
            
            for component in framework_components:
                comp_para = doc.add_paragraph()
                comp_run = comp_para.add_run(component)
                comp_run.font.name = 'Calibri'
                comp_run.font.size = Pt(11)
            
            # Key findings
            findings_heading = doc.add_paragraph()
            findings_run = findings_heading.add_run("Key Findings:")
            findings_run.font.name = 'Calibri'
            findings_run.font.size = Pt(12)
            findings_run.font.bold = True
            
            # Extract key points from RAG response
            key_points = self._extract_key_points(rag_response)
            for point in key_points[:5]:  # Top 5 key points
                point_para = doc.add_paragraph()
                point_run = point_para.add_run(f"• {point}")
                point_run.font.name = 'Calibri'
                point_run.font.size = Pt(11)
            
            doc.add_paragraph()
            
        except Exception as e:
            print(f"Warning: Could not add executive summary: {e}")
    
    def _add_detailed_analysis(self, doc, rag_response, cited_pages, page_scores):
        """Add detailed analysis section aligned with security analysis framework"""
        try:
            # Import RGBColor for proper color handling
            from docx.shared import RGBColor
            
            # Section heading
            heading = doc.add_paragraph()
            heading_run = heading.add_run("DETAILED ANALYSIS")
            heading_run.font.name = 'Calibri'
            heading_run.font.size = Pt(16)
            heading_run.font.bold = True
            heading_run.font.color.rgb = RGBColor(47, 84, 150)  # #2F5496
            
            # 1. Fact-Finding & Contextualization
            fact_finding_heading = doc.add_paragraph()
            fact_finding_run = fact_finding_heading.add_run("1. FACT-FINDING & CONTEXTUALIZATION")
            fact_finding_run.font.name = 'Calibri'
            fact_finding_run.font.size = Pt(14)
            fact_finding_run.font.bold = True
            fact_finding_run.font.color.rgb = RGBColor(40, 167, 69)  # #28a745
            
            fact_finding_para = doc.add_paragraph()
            fact_finding_para_run = fact_finding_para.add_run("This section provides background information for readers to understand the origin, development, and context of the subject topic.")
            fact_finding_para_run.font.name = 'Calibri'
            fact_finding_para_run.font.size = Pt(11)
            
            # Extract contextual information
            context_info = self._extract_contextual_info(rag_response)
            for info in context_info:
                info_para = doc.add_paragraph()
                info_run = info_para.add_run(f"• {info}")
                info_run.font.name = 'Calibri'
                info_run.font.size = Pt(11)
            
            doc.add_paragraph()
            
            # 2. Case Study Identification
            case_study_heading = doc.add_paragraph()
            case_study_run = case_study_heading.add_run("2. CASE STUDY IDENTIFICATION")
            case_study_run.font.name = 'Calibri'
            case_study_run.font.size = Pt(14)
            case_study_run.font.bold = True
            case_study_run.font.color.rgb = RGBColor(255, 193, 7)  # #ffc107
            
            case_study_para = doc.add_paragraph()
            case_study_para_run = case_study_para.add_run("This section provides context and prevalence assessment, highlighting past incidents to establish patterns and extract relevant TTPs for analysis.")
            case_study_para_run.font.name = 'Calibri'
            case_study_para_run.font.size = Pt(11)
            
            # Extract case study information
            case_studies = self._extract_case_studies(rag_response)
            for case in case_studies:
                case_para = doc.add_paragraph()
                case_run = case_para.add_run(f"• {case}")
                case_run.font.name = 'Calibri'
                case_run.font.size = Pt(11)
            
            doc.add_paragraph()
            
            # 3. Analytical Assessment
            analytical_heading = doc.add_paragraph()
            analytical_run = analytical_heading.add_run("3. ANALYTICAL ASSESSMENT")
            analytical_run.font.name = 'Calibri'
            analytical_run.font.size = Pt(14)
            analytical_run.font.bold = True
            analytical_run.font.color.rgb = RGBColor(220, 53, 69)  # #dc3545
            
            analytical_para = doc.add_paragraph()
            analytical_para_run = analytical_para.add_run("This section evaluates gathered information to assess intent, motivation, TTPs, emerging trends, and relevance to threat landscapes.")
            analytical_para_run.font.name = 'Calibri'
            analytical_para_run.font.size = Pt(11)
            
            # Extract analytical insights
            analytical_insights = self._extract_analytical_insights(rag_response)
            for insight in analytical_insights:
                insight_para = doc.add_paragraph()
                insight_run = insight_para.add_run(f"• {insight}")
                insight_run.font.name = 'Calibri'
                insight_run.font.size = Pt(11)
            
            doc.add_paragraph()
            
            # 4. Operational Relevance
            operational_heading = doc.add_paragraph()
            operational_run = operational_heading.add_run("4. OPERATIONAL RELEVANCE")
            operational_run.font.name = 'Calibri'
            operational_run.font.size = Pt(14)
            operational_run.font.bold = True
            operational_run.font.color.rgb = RGBColor(111, 66, 193)  # #6f42c1
            
            operational_para = doc.add_paragraph()
            operational_para_run = operational_para.add_run("This section translates research insights into actionable knowledge for ground-level personnel, highlighting operational risks and procedural recommendations.")
            operational_para_run.font.name = 'Calibri'
            operational_para_run.font.size = Pt(11)
            
            # Extract operational insights
            operational_insights = self._extract_operational_insights(rag_response)
            for insight in operational_insights:
                insight_para = doc.add_paragraph()
                insight_run = insight_para.add_run(f"• {insight}")
                insight_run.font.name = 'Calibri'
                insight_run.font.size = Pt(11)
            
            doc.add_paragraph()
            
            # Main RAG response as comprehensive analysis
            main_analysis_heading = doc.add_paragraph()
            main_analysis_run = main_analysis_heading.add_run("COMPREHENSIVE ANALYSIS")
            main_analysis_run.font.name = 'Calibri'
            main_analysis_run.font.size = Pt(12)
            main_analysis_run.font.bold = True
            
            response_para = doc.add_paragraph()
            response_run = response_para.add_run(rag_response)
            response_run.font.name = 'Calibri'
            response_run.font.size = Pt(11)
            
            doc.add_paragraph()
            
        except Exception as e:
            print(f"Warning: Could not add detailed analysis: {e}")
    
    def _add_methodology_section(self, doc, cited_pages, page_scores):
        """Add methodology section aligned with security analysis framework"""
        try:
            # Import RGBColor for proper color handling
            from docx.shared import RGBColor
            
            # Section heading
            heading = doc.add_paragraph()
            heading_run = heading.add_run("METHODOLOGY")
            heading_run.font.name = 'Calibri'
            heading_run.font.size = Pt(16)
            heading_run.font.bold = True
            heading_run.font.color.rgb = RGBColor(47, 84, 150)  # #2F5496
            
            # Methodology content
            method_para = doc.add_paragraph()
            method_run = method_para.add_run("This security analysis was conducted using advanced AI-powered threat intelligence and document analysis techniques:")
            method_run.font.name = 'Calibri'
            method_run.font.size = Pt(11)
            
            # Analysis Framework
            framework_heading = doc.add_paragraph()
            framework_run = framework_heading.add_run("Security Analysis Framework:")
            framework_run.font.name = 'Calibri'
            framework_run.font.size = Pt(12)
            framework_run.font.bold = True
            
            framework_components = [
                "• Fact-Finding & Contextualization: Background research and context development",
                "• Case Study Identification: Incident analysis and TTP extraction",
                "• Analytical Assessment: Threat landscape evaluation and risk assessment",
                "• Operational Relevance: Ground-level actionable intelligence generation"
            ]
            
            for component in framework_components:
                comp_para = doc.add_paragraph()
                comp_run = comp_para.add_run(component)
                comp_run.font.name = 'Calibri'
                comp_run.font.size = Pt(11)
            
            # Document sources
            sources_heading = doc.add_paragraph()
            sources_run = sources_heading.add_run("Intelligence Sources:")
            sources_run.font.name = 'Calibri'
            sources_run.font.size = Pt(12)
            sources_run.font.bold = True
            
            # List sources
            for i, citation in enumerate(cited_pages):
                source_para = doc.add_paragraph()
                source_run = source_para.add_run(f"{i+1}. {citation}")
                source_run.font.name = 'Calibri'
                source_run.font.size = Pt(11)
            
            # Analysis approach
            approach_heading = doc.add_paragraph()
            approach_run = approach_heading.add_run("Technical Analysis Approach:")
            approach_run.font.name = 'Calibri'
            approach_run.font.size = Pt(12)
            approach_run.font.bold = True
            
            approach_para = doc.add_paragraph()
            approach_run = approach_para.add_run("• Multi-modal document analysis using AI vision models for threat pattern recognition")
            approach_run.font.name = 'Calibri'
            approach_run.font.size = Pt(11)
            
            approach2_para = doc.add_paragraph()
            approach2_run = approach2_para.add_run("• Intelligent content retrieval and relevance scoring for threat intelligence prioritization")
            approach2_run.font.name = 'Calibri'
            approach2_run.font.size = Pt(11)
            
            approach3_para = doc.add_paragraph()
            approach3_run = approach3_para.add_run("• Comprehensive threat synthesis and actionable intelligence generation")
            approach3_run.font.name = 'Calibri'
            approach3_run.font.size = Pt(11)
            
            approach4_para = doc.add_paragraph()
            approach4_run = approach4_para.add_run("• Evidence-based risk assessment and operational recommendation development")
            approach4_run.font.name = 'Calibri'
            approach4_run.font.size = Pt(11)
            
            doc.add_paragraph()
            
        except Exception as e:
            print(f"Warning: Could not add methodology section: {e}")
    
    def _add_findings_conclusions(self, doc, rag_response, cited_pages):
        """Add findings and conclusions section aligned with security analysis framework"""
        try:
            # Import RGBColor for proper color handling
            from docx.shared import RGBColor
            
            # Section heading
            heading = doc.add_paragraph()
            heading_run = heading.add_run("FINDINGS AND CONCLUSIONS")
            heading_run.font.name = 'Calibri'
            heading_run.font.size = Pt(16)
            heading_run.font.bold = True
            heading_run.font.color.rgb = RGBColor(47, 84, 150)  # #2F5496
            
            # Threat Assessment Summary
            threat_heading = doc.add_paragraph()
            threat_run = threat_heading.add_run("Threat Assessment Summary:")
            threat_run.font.name = 'Calibri'
            threat_run.font.size = Pt(12)
            threat_run.font.bold = True
            
            # Extract threat-related findings
            threat_findings = self._extract_threat_findings(rag_response)
            for finding in threat_findings:
                finding_para = doc.add_paragraph()
                finding_run = finding_para.add_run(f"• {finding}")
                finding_run.font.name = 'Calibri'
                finding_run.font.size = Pt(11)
            
            # TTP Analysis
            ttp_heading = doc.add_paragraph()
            ttp_run = ttp_heading.add_run("Tactics, Techniques, and Procedures (TTPs):")
            ttp_run.font.name = 'Calibri'
            ttp_run.font.size = Pt(12)
            ttp_run.font.bold = True
            
            # Extract TTP information
            ttps = self._extract_ttps(rag_response)
            for ttp in ttps:
                ttp_para = doc.add_paragraph()
                ttp_run = ttp_para.add_run(f"• {ttp}")
                ttp_run.font.name = 'Calibri'
                ttp_run.font.size = Pt(11)
            
            # Operational Recommendations
            recommendations_heading = doc.add_paragraph()
            recommendations_run = recommendations_heading.add_run("Operational Recommendations:")
            recommendations_run.font.name = 'Calibri'
            recommendations_run.font.size = Pt(12)
            recommendations_run.font.bold = True
            
            # Extract operational recommendations
            recommendations = self._extract_operational_recommendations(rag_response)
            for rec in recommendations:
                rec_para = doc.add_paragraph()
                rec_run = rec_para.add_run(f"• {rec}")
                rec_run.font.name = 'Calibri'
                rec_run.font.size = Pt(11)
            
            # Risk Assessment
            risk_heading = doc.add_paragraph()
            risk_run = risk_heading.add_run("Risk Assessment:")
            risk_run.font.name = 'Calibri'
            risk_run.font.size = Pt(12)
            risk_run.font.bold = True
            
            # Extract risk information
            risks = self._extract_risk_assessment(rag_response)
            for risk in risks:
                risk_para = doc.add_paragraph()
                risk_run = risk_para.add_run(f"• {risk}")
                risk_run.font.name = 'Calibri'
                risk_run.font.size = Pt(11)
            
            # Conclusions
            conclusions_heading = doc.add_paragraph()
            conclusions_run = conclusions_heading.add_run("Conclusions:")
            conclusions_run.font.name = 'Calibri'
            conclusions_run.font.size = Pt(12)
            conclusions_run.font.bold = True
            
            conclusions_para = doc.add_paragraph()
            conclusions_run = conclusions_para.add_run("This security analysis provides actionable intelligence for threat mitigation and operational preparedness. The findings support evidence-based decision making for security operations and risk management.")
            conclusions_run.font.name = 'Calibri'
            conclusions_run.font.size = Pt(11)
            
            doc.add_paragraph()
            
        except Exception as e:
            print(f"Warning: Could not add findings and conclusions: {e}")
    
    def _add_appendices(self, doc, cited_pages, page_scores):
        """Add appendices section"""
        try:
            # Import RGBColor for proper color handling
            from docx.shared import RGBColor
            
            # Section heading
            heading = doc.add_paragraph()
            heading_run = heading.add_run("APPENDICES")
            heading_run.font.name = 'Calibri'
            heading_run.font.size = Pt(16)
            heading_run.font.bold = True
            heading_run.font.color.rgb = RGBColor(47, 84, 150)  # #2F5496
            
            # Appendix A: Document Sources
            appendix_a = doc.add_paragraph()
            appendix_a_run = appendix_a.add_run("Appendix A: Document Sources and Relevance Scores")
            appendix_a_run.font.name = 'Calibri'
            appendix_a_run.font.size = Pt(12)
            appendix_a_run.font.bold = True
            
            for i, (citation, score) in enumerate(zip(cited_pages, page_scores)):
                source_para = doc.add_paragraph()
                source_run = source_para.add_run(f"{i+1}. {citation} (Relevance Score: {score:.3f})")
                source_run.font.name = 'Calibri'
                source_run.font.size = Pt(11)
            
            doc.add_paragraph()
            
        except Exception as e:
            print(f"Warning: Could not add appendices: {e}")
    
    def _extract_key_points(self, rag_response):
        """Extract key points from RAG response"""
        try:
            # Split response into sentences
            sentences = re.split(r'[.!?]+', rag_response)
            key_points = []
            
            # Look for sentences with key indicators
            key_indicators = ['important', 'key', 'critical', 'essential', 'significant', 'major', 'primary', 'main']
            
            for sentence in sentences:
                sentence = sentence.strip()
                if len(sentence) > 20 and any(indicator in sentence.lower() for indicator in key_indicators):
                    key_points.append(sentence)
            
            # If not enough key points found, use first few sentences
            if len(key_points) < 3:
                key_points = [s.strip() for s in sentences[:5] if len(s.strip()) > 20]
            
            return key_points[:5]  # Return top 5
            
        except Exception as e:
            print(f"Warning: Could not extract key points: {e}")
            return ["Analysis completed successfully", "Comprehensive review performed", "Key insights identified"]
    
    def _extract_contextual_info(self, rag_response):
        """Extract contextual information for fact-finding section"""
        try:
            sentences = re.split(r'[.!?]+', rag_response)
            contextual_info = []
            
            # Look for contextual indicators
            context_indicators = [
                'background', 'history', 'origin', 'development', 'context', 'definition',
                'introduction', 'overview', 'description', 'characteristics', 'features',
                'components', 'types', 'categories', 'classification', 'structure'
            ]
            
            for sentence in sentences:
                sentence = sentence.strip()
                if len(sentence) > 15 and any(indicator in sentence.lower() for indicator in context_indicators):
                    contextual_info.append(sentence)
            
            # If not enough contextual info, use general descriptive sentences
            if len(contextual_info) < 3:
                contextual_info = [s.strip() for s in sentences[:3] if len(s.strip()) > 15]
            
            return contextual_info[:5]  # Return top 5
            
        except Exception as e:
            print(f"Warning: Could not extract contextual info: {e}")
            return ["Background information extracted from analysis", "Contextual details identified", "Historical context established"]
    
    def _extract_case_studies(self, rag_response):
        """Extract case study information for incident identification"""
        try:
            sentences = re.split(r'[.!?]+', rag_response)
            case_studies = []
            
            # Look for case study indicators
            case_indicators = [
                'incident', 'case', 'example', 'instance', 'occurrence', 'event',
                'attack', 'threat', 'vulnerability', 'exploit', 'breach', 'compromise',
                'pattern', 'trend', 'frequency', 'prevalence', 'statistics', 'data'
            ]
            
            for sentence in sentences:
                sentence = sentence.strip()
                if len(sentence) > 15 and any(indicator in sentence.lower() for indicator in case_indicators):
                    case_studies.append(sentence)
            
            # If not enough case studies, use sentences with numbers or dates
            if len(case_studies) < 3:
                for sentence in sentences:
                    sentence = sentence.strip()
                    if len(sentence) > 15 and (re.search(r'\d+', sentence) or any(word in sentence.lower() for word in ['first', 'second', 'third', 'recent', 'previous'])):
                        case_studies.append(sentence)
            
            return case_studies[:5]  # Return top 5
            
        except Exception as e:
            print(f"Warning: Could not extract case studies: {e}")
            return ["Incident patterns identified", "Case study information extracted", "Prevalence data analyzed"]
    
    def _extract_analytical_insights(self, rag_response):
        """Extract analytical insights for threat assessment"""
        try:
            sentences = re.split(r'[.!?]+', rag_response)
            analytical_insights = []
            
            # Look for analytical indicators
            analytical_indicators = [
                'intent', 'motivation', 'purpose', 'objective', 'goal', 'target',
                'technique', 'procedure', 'method', 'approach', 'strategy', 'tactic',
                'trend', 'emerging', 'evolution', 'development', 'change', 'shift',
                'threat', 'risk', 'vulnerability', 'impact', 'consequence', 'effect'
            ]
            
            for sentence in sentences:
                sentence = sentence.strip()
                if len(sentence) > 15 and any(indicator in sentence.lower() for indicator in analytical_indicators):
                    analytical_insights.append(sentence)
            
            # If not enough insights, use sentences with analytical language
            if len(analytical_insights) < 3:
                for sentence in sentences:
                    sentence = sentence.strip()
                    if len(sentence) > 15 and any(word in sentence.lower() for word in ['because', 'therefore', 'however', 'although', 'while', 'despite']):
                        analytical_insights.append(sentence)
            
            return analytical_insights[:5]  # Return top 5
            
        except Exception as e:
            print(f"Warning: Could not extract analytical insights: {e}")
            return ["Analytical assessment completed", "Threat landscape evaluated", "Risk factors identified"]
    
    def _extract_operational_insights(self, rag_response):
        """Extract operational insights for ground-level recommendations"""
        try:
            sentences = re.split(r'[.!?]+', rag_response)
            operational_insights = []
            
            # Look for operational indicators
            operational_indicators = [
                'recommendation', 'action', 'procedure', 'protocol', 'guideline',
                'training', 'awareness', 'vigilance', 'monitoring', 'detection',
                'prevention', 'mitigation', 'response', 'recovery', 'preparation',
                'equipment', 'tool', 'technology', 'system', 'process', 'workflow'
            ]
            
            for sentence in sentences:
                sentence = sentence.strip()
                if len(sentence) > 15 and any(indicator in sentence.lower() for indicator in operational_indicators):
                    operational_insights.append(sentence)
            
            # If not enough operational insights, use sentences with actionable language
            if len(operational_insights) < 3:
                for sentence in sentences:
                    sentence = sentence.strip()
                    if len(sentence) > 15 and any(word in sentence.lower() for word in ['should', 'must', 'need', 'require', 'implement', 'establish', 'develop']):
                        operational_insights.append(sentence)
            
            return operational_insights[:5]  # Return top 5
            
        except Exception as e:
            print(f"Warning: Could not extract operational insights: {e}")
            return ["Operational recommendations identified", "Ground-level procedures suggested", "Training requirements outlined"]
    
    def _extract_findings(self, rag_response):
        """Extract findings from RAG response"""
        try:
            # Split response into sentences
            sentences = re.split(r'[.!?]+', rag_response)
            findings = []
            
            # Look for sentences that might be findings
            finding_indicators = ['found', 'discovered', 'identified', 'revealed', 'shows', 'indicates', 'demonstrates', 'suggests']
            
            for sentence in sentences:
                sentence = sentence.strip()
                if len(sentence) > 15 and any(indicator in sentence.lower() for indicator in finding_indicators):
                    findings.append(sentence)
            
            # If not enough findings, use meaningful sentences
            if len(findings) < 3:
                findings = [s.strip() for s in sentences[:5] if len(s.strip()) > 15]
            
            return findings[:5]  # Return top 5
            
        except Exception as e:
            print(f"Warning: Could not extract findings: {e}")
            return ["Analysis completed successfully", "Comprehensive review performed", "Key insights identified"]
    
    def _extract_threat_findings(self, rag_response):
        """Extract threat-related findings for security analysis"""
        try:
            sentences = re.split(r'[.!?]+', rag_response)
            threat_findings = []
            
            # Look for threat-related indicators
            threat_indicators = [
                'threat', 'attack', 'vulnerability', 'exploit', 'breach', 'compromise',
                'malware', 'phishing', 'social engineering', 'ransomware', 'ddos',
                'intrusion', 'infiltration', 'espionage', 'sabotage', 'terrorism'
            ]
            
            for sentence in sentences:
                sentence = sentence.strip()
                if len(sentence) > 15 and any(indicator in sentence.lower() for indicator in threat_indicators):
                    threat_findings.append(sentence)
            
            # If not enough threat findings, use general security-related sentences
            if len(threat_findings) < 3:
                for sentence in sentences:
                    sentence = sentence.strip()
                    if len(sentence) > 15 and any(word in sentence.lower() for word in ['security', 'risk', 'danger', 'hazard', 'warning']):
                        threat_findings.append(sentence)
            
            return threat_findings[:5]  # Return top 5
            
        except Exception as e:
            print(f"Warning: Could not extract threat findings: {e}")
            return ["Threat assessment completed", "Security vulnerabilities identified", "Risk factors analyzed"]
    
    def _extract_ttps(self, rag_response):
        """Extract Tactics, Techniques, and Procedures (TTPs)"""
        try:
            sentences = re.split(r'[.!?]+', rag_response)
            ttps = []
            
            # Look for TTP indicators
            ttp_indicators = [
                'technique', 'procedure', 'method', 'approach', 'strategy', 'tactic',
                'process', 'workflow', 'protocol', 'standard', 'practice', 'modus operandi',
                'attack vector', 'exploitation', 'infiltration', 'persistence', 'exfiltration'
            ]
            
            for sentence in sentences:
                sentence = sentence.strip()
                if len(sentence) > 15 and any(indicator in sentence.lower() for indicator in ttp_indicators):
                    ttps.append(sentence)
            
            # If not enough TTPs, use sentences with procedural language
            if len(ttps) < 3:
                for sentence in sentences:
                    sentence = sentence.strip()
                    if len(sentence) > 15 and any(word in sentence.lower() for word in ['step', 'phase', 'stage', 'sequence', 'order']):
                        ttps.append(sentence)
            
            return ttps[:5]  # Return top 5
            
        except Exception as e:
            print(f"Warning: Could not extract TTPs: {e}")
            return ["TTP analysis completed", "Attack methods identified", "Procedural patterns extracted"]
    
    def _extract_operational_recommendations(self, rag_response):
        """Extract operational recommendations for ground-level personnel"""
        try:
            sentences = re.split(r'[.!?]+', rag_response)
            recommendations = []
            
            # Look for recommendation indicators
            recommendation_indicators = [
                'recommend', 'suggest', 'advise', 'propose', 'should', 'must', 'need',
                'implement', 'establish', 'develop', 'create', 'adopt', 'apply',
                'training', 'awareness', 'education', 'preparation', 'readiness'
            ]
            
            for sentence in sentences:
                sentence = sentence.strip()
                if len(sentence) > 15 and any(indicator in sentence.lower() for indicator in recommendation_indicators):
                    recommendations.append(sentence)
            
            # If not enough recommendations, use sentences with actionable language
            if len(recommendations) < 3:
                for sentence in sentences:
                    sentence = sentence.strip()
                    if len(sentence) > 15 and any(word in sentence.lower() for word in ['action', 'measure', 'step', 'procedure', 'protocol']):
                        recommendations.append(sentence)
            
            return recommendations[:5]  # Return top 5
            
        except Exception as e:
            print(f"Warning: Could not extract operational recommendations: {e}")
            return ["Operational procedures recommended", "Training requirements identified", "Security measures suggested"]
    
    def _extract_risk_assessment(self, rag_response):
        """Extract risk assessment information"""
        try:
            sentences = re.split(r'[.!?]+', rag_response)
            risks = []
            
            # Look for risk indicators
            risk_indicators = [
                'risk', 'danger', 'hazard', 'threat', 'vulnerability', 'exposure',
                'probability', 'likelihood', 'impact', 'consequence', 'severity',
                'critical', 'high', 'medium', 'low', 'minimal', 'significant'
            ]
            
            for sentence in sentences:
                sentence = sentence.strip()
                if len(sentence) > 15 and any(indicator in sentence.lower() for indicator in risk_indicators):
                    risks.append(sentence)
            
            # If not enough risks, use sentences with risk-related language
            if len(risks) < 3:
                for sentence in sentences:
                    sentence = sentence.strip()
                    if len(sentence) > 15 and any(word in sentence.lower() for word in ['potential', 'possible', 'likely', 'unlikely', 'certain']):
                        risks.append(sentence)
            
            return risks[:5]  # Return top 5
            
        except Exception as e:
            print(f"Warning: Could not extract risk assessment: {e}")
            return ["Risk assessment completed", "Vulnerability analysis performed", "Threat evaluation conducted"]
    
    def _generate_enhanced_excel_export(self, query, rag_response, cited_pages, page_scores, custom_headers=None):
        """
        Generate enhanced Excel export with proper formatting for charts and graphs
        """
        if not EXCEL_AVAILABLE:
            return None, "Excel export not available - openpyxl/pandas libraries not installed"
        
        try:
            print("📊 [EXCEL] Generating enhanced Excel export...")
            
            # Extract custom headers from query if not provided
            if custom_headers is None:
                custom_headers = self._extract_custom_headers(query)
            
            # Create a new workbook
            wb = Workbook()
            
            # Remove default sheet
            wb.remove(wb.active)
            
            # Create main data sheet
            data_sheet = wb.create_sheet("Data")
            
            # Create summary sheet
            summary_sheet = wb.create_sheet("Summary")
            
            # Create charts sheet
            charts_sheet = wb.create_sheet("Charts")
            
            # Extract structured data
            structured_data = self._extract_structured_data_for_excel(rag_response, cited_pages, page_scores, custom_headers)
            
            # Populate data sheet
            self._populate_data_sheet(data_sheet, structured_data, query)
            
            # Populate summary sheet
            self._populate_summary_sheet(summary_sheet, query, cited_pages, page_scores)
            
            # Create charts if chart request detected
            if self._detect_chart_request(query):
                self._create_excel_charts(charts_sheet, structured_data, query, custom_headers)
            
            # Generate unique filename
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            safe_query = "".join(c for c in query[:30] if c.isalnum() or c in (' ', '-', '_')).rstrip()
            safe_query = safe_query.replace(' ', '_')
            filename = f"enhanced_export_{safe_query}_{timestamp}.xlsx"
            filepath = os.path.join("temp", filename)
            
            # Ensure temp directory exists
            os.makedirs("temp", exist_ok=True)
            
            # Save the workbook
            wb.save(filepath)
            
            print(f"✅ [EXCEL] Enhanced Excel export generated: {filepath}")
            return filepath, None
            
        except Exception as e:
            error_msg = f"Error generating Excel export: {str(e)}"
            print(f"❌ [EXCEL] {error_msg}")
            return None, error_msg
    
    def _extract_structured_data_for_excel(self, rag_response, cited_pages, page_scores, custom_headers=None):
        """Extract structured data specifically for Excel export"""
        try:
            # If custom headers provided, use them
            if custom_headers:
                headers = custom_headers
                print(f"📊 [EXCEL] Using custom headers: {headers}")
            else:
                # Auto-detect headers based on content
                headers = self._auto_detect_excel_headers(rag_response, cited_pages)
                print(f"📊 [EXCEL] Auto-detected headers: {headers}")
            
            # Extract data rows
            data_rows = []
            
            # If custom headers are provided, try to map data to them
            if custom_headers:
                mapped_data = self._map_data_to_custom_headers(rag_response, cited_pages, page_scores, custom_headers)
                if mapped_data:
                    data_rows.extend(mapped_data)
            
            # If no custom data or mapping failed, extract standard data
            if not data_rows:
                # Extract numerical data if present
                numerical_data = self._extract_numerical_data(rag_response)
                if numerical_data:
                    data_rows.extend(numerical_data)
                
                # Extract categorical data
                categorical_data = self._extract_categorical_data(rag_response, cited_pages)
                if categorical_data:
                    data_rows.extend(categorical_data)
                
                # Extract source information
                source_data = self._extract_source_data(cited_pages, page_scores)
                if source_data:
                    data_rows.extend(source_data)
            
            # If still no structured data found, create summary data
            if not data_rows:
                data_rows = self._create_summary_data(rag_response, cited_pages, page_scores)
            
            return {
                'headers': headers,
                'data': data_rows
            }
            
        except Exception as e:
            print(f"Error extracting structured data for Excel: {e}")
            return {
                'headers': ['Category', 'Value', 'Description'],
                'data': [['Analysis', 'Completed', 'Data extracted successfully']]
            }
    
    def _auto_detect_excel_headers(self, rag_response, cited_pages):
        """Auto-detect contextually appropriate headers for Excel export based on query content"""
        try:
            headers = []
            
            # Analyze the content for context clues
            rag_lower = rag_response.lower()
            
            # Security/Analysis context detection
            if any(word in rag_lower for word in ['threat', 'attack', 'vulnerability', 'security', 'risk']):
                if 'threat' in rag_lower or 'attack' in rag_lower:
                    headers.append('Threat Type')
                if 'frequency' in rag_lower or 'count' in rag_lower or 'percentage' in rag_lower:
                    headers.append('Frequency')
                if 'risk' in rag_lower or 'severity' in rag_lower:
                    headers.append('Risk Level')
                if 'impact' in rag_lower or 'damage' in rag_lower:
                    headers.append('Impact')
                if 'mitigation' in rag_lower or 'solution' in rag_lower:
                    headers.append('Mitigation')
            
            # Business/Performance context detection
            elif any(word in rag_lower for word in ['sales', 'revenue', 'performance', 'growth', 'profit']):
                if 'month' in rag_lower or 'quarter' in rag_lower or 'year' in rag_lower:
                    headers.append('Time Period')
                if 'sales' in rag_lower or 'revenue' in rag_lower:
                    headers.append('Sales/Revenue')
                if 'growth' in rag_lower or 'increase' in rag_lower:
                    headers.append('Growth Rate')
                if 'region' in rag_lower or 'location' in rag_lower:
                    headers.append('Region')
            
            # Technical/System context detection
            elif any(word in rag_lower for word in ['system', 'component', 'device', 'technology', 'software']):
                if 'component' in rag_lower or 'device' in rag_lower:
                    headers.append('Component')
                if 'status' in rag_lower or 'condition' in rag_lower:
                    headers.append('Status')
                if 'priority' in rag_lower or 'importance' in rag_lower:
                    headers.append('Priority')
                if 'version' in rag_lower or 'release' in rag_lower:
                    headers.append('Version')
            
            # Data/Statistics context detection
            elif any(word in rag_lower for word in ['data', 'statistics', 'analysis', 'report', 'survey']):
                if 'category' in rag_lower or 'type' in rag_lower:
                    headers.append('Category')
                if 'value' in rag_lower or 'number' in rag_lower or 'count' in rag_lower:
                    headers.append('Value')
                if 'percentage' in rag_lower or 'rate' in rag_lower:
                    headers.append('Percentage')
                if 'trend' in rag_lower or 'change' in rag_lower:
                    headers.append('Trend')
            
            # Generic fallback detection
            else:
                # Check for numerical data
                if re.search(r'\d+', rag_response):
                    headers.append('Value')
                
                # Check for categories or types
                if any(word in rag_lower for word in ['type', 'category', 'class', 'group']):
                    headers.append('Category')
                
                # Check for descriptions
                if len(rag_response) > 100:
                    headers.append('Description')
                
                # Check for sources
                if cited_pages:
                    headers.append('Source')
                
                # Check for scores or ratings
                if any(word in rag_lower for word in ['score', 'rating', 'level', 'grade']):
                    headers.append('Score')
            
            # Ensure we have at least 2-3 headers for chart generation
            if len(headers) < 2:
                if 'Category' not in headers:
                    headers.append('Category')
                if 'Value' not in headers:
                    headers.append('Value')
            
            if len(headers) < 3:
                if 'Description' not in headers:
                    headers.append('Description')
            
            # Limit to 4 headers maximum for chart clarity
            headers = headers[:4]
            
            print(f"📊 [EXCEL] Auto-detected contextually relevant headers: {headers}")
            return headers
            
        except Exception as e:
            print(f"Error auto-detecting headers: {e}")
            return ['Category', 'Value', 'Description']
    
    def _extract_numerical_data(self, rag_response):
        """Extract numerical data from RAG response"""
        try:
            data_rows = []
            
            # Find numbers with context
            number_patterns = [
                r'(\d+(?:\.\d+)?)\s*(percent|%|units|items|components|devices|procedures)',
                r'(\d+(?:\.\d+)?)\s*(voltage|current|resistance|power|frequency)',
                r'(\d+(?:\.\d+)?)\s*(safety|risk|danger|warning)',
                r'(\d+(?:\.\d+)?)\s*(steps|phases|stages|levels)'
            ]
            
            for pattern in number_patterns:
                matches = re.findall(pattern, rag_response, re.IGNORECASE)
                for match in matches:
                    value, category = match
                    data_rows.append([category.title(), value, f"Found in analysis"])
            
            return data_rows
            
        except Exception as e:
            print(f"Error extracting numerical data: {e}")
            return []
    
    def _extract_categorical_data(self, rag_response, cited_pages):
        """Extract categorical data from RAG response"""
        try:
            data_rows = []
            
            # Extract categories mentioned in the response
            categories = []
            
            # Look for common category patterns
            category_patterns = [
                r'(safety|security|warning|danger|risk)',
                r'(procedure|method|technique|approach)',
                r'(component|device|equipment|tool)',
                r'(type|category|class|group)',
                r'(input|output|control|monitoring)'
            ]
            
            for pattern in category_patterns:
                matches = re.findall(pattern, rag_response, re.IGNORECASE)
                categories.extend(matches)
            
            # Remove duplicates
            categories = list(set(categories))
            
            for category in categories[:10]:  # Limit to 10 categories
                data_rows.append([category.title(), 'Identified', f"Category found in analysis"])
            
            return data_rows
            
        except Exception as e:
            print(f"Error extracting categorical data: {e}")
            return []
    
    def _extract_source_data(self, cited_pages, page_scores):
        """Extract source information for Excel"""
        try:
            data_rows = []
            
            for i, (citation, score) in enumerate(zip(cited_pages, page_scores)):
                collection = citation.split(' from ')[1] if ' from ' in citation else 'Unknown'
                page_num = citation.split('Page ')[1].split(' from')[0] if 'Page ' in citation else str(i+1)
                
                data_rows.append([
                    f"Source {i+1}",
                    collection,
                    f"Page {page_num} (Score: {score:.3f})"
                ])
            
            return data_rows
            
        except Exception as e:
            print(f"Error extracting source data: {e}")
            return []
    
    def _map_data_to_custom_headers(self, rag_response, cited_pages, page_scores, custom_headers):
        """Map extracted data to custom headers for Excel export with context-aware sample data"""
        try:
            data_rows = []
            
            # Extract various types of data
            numerical_data = self._extract_numerical_data(rag_response)
            categorical_data = self._extract_categorical_data(rag_response, cited_pages)
            source_data = self._extract_source_data(cited_pages, page_scores)
            
            # Combine all available data
            all_data = []
            if numerical_data:
                all_data.extend(numerical_data)
            if categorical_data:
                all_data.extend(categorical_data)
            if source_data:
                all_data.extend(source_data)
            
            # Map data to custom headers
            for i, data_row in enumerate(all_data):
                mapped_row = []
                
                # Ensure we have enough data for all headers
                while len(mapped_row) < len(custom_headers):
                    if len(data_row) > len(mapped_row):
                        mapped_row.append(data_row[len(mapped_row)])
                    else:
                        # Fill with contextually relevant placeholder data
                        header = custom_headers[len(mapped_row)]
                        mapped_row.append(self._generate_contextual_sample_data(header, i, rag_response))
                
                # Truncate if we have too many values
                mapped_row = mapped_row[:len(custom_headers)]
                data_rows.append(mapped_row)
            
            # If no data was mapped, create contextually relevant sample data
            if not data_rows:
                data_rows = self._create_contextual_sample_data(custom_headers, rag_response)
            
            print(f"📊 [EXCEL] Mapped {len(data_rows)} rows to custom headers")
            return data_rows
            
        except Exception as e:
            print(f"Error mapping data to custom headers: {e}")
            return []

    def _generate_contextual_sample_data(self, header, index, rag_response):
        """Generate contextually relevant sample data based on header and content"""
        try:
            header_lower = header.lower()
            rag_lower = rag_response.lower()
            
            # Security context
            if any(word in rag_lower for word in ['threat', 'attack', 'security', 'vulnerability']):
                if 'threat' in header_lower or 'attack' in header_lower:
                    threats = ['Phishing', 'Malware', 'DDoS', 'Social Engineering', 'Ransomware']
                    return threats[index % len(threats)]
                elif 'frequency' in header_lower or 'count' in header_lower:
                    return str((index + 1) * 15) + '%'
                elif 'risk' in header_lower or 'severity' in header_lower:
                    risk_levels = ['Low', 'Medium', 'High', 'Critical']
                    return risk_levels[index % len(risk_levels)]
                elif 'impact' in header_lower:
                    impacts = ['Minimal', 'Moderate', 'Significant', 'Severe']
                    return impacts[index % len(impacts)]
                elif 'mitigation' in header_lower:
                    mitigations = ['Training', 'Firewall', 'Monitoring', 'Backup']
                    return mitigations[index % len(mitigations)]
            
            # Business context
            elif any(word in rag_lower for word in ['sales', 'revenue', 'business', 'performance']):
                if 'time' in header_lower or 'period' in header_lower:
                    periods = ['Q1 2024', 'Q2 2024', 'Q3 2024', 'Q4 2024']
                    return periods[index % len(periods)]
                elif 'sales' in header_lower or 'revenue' in header_lower:
                    return f"${(index + 1) * 10000:,}"
                elif 'growth' in header_lower:
                    return f"+{(index + 1) * 5}%"
                elif 'region' in header_lower:
                    regions = ['North', 'South', 'East', 'West']
                    return regions[index % len(regions)]
            
            # Technical context
            elif any(word in rag_lower for word in ['system', 'component', 'device', 'technology']):
                if 'component' in header_lower:
                    components = ['Server', 'Database', 'Network', 'Application']
                    return components[index % len(components)]
                elif 'status' in header_lower:
                    statuses = ['Active', 'Inactive', 'Maintenance', 'Error']
                    return statuses[index % len(statuses)]
                elif 'priority' in header_lower:
                    priorities = ['Low', 'Medium', 'High', 'Critical']
                    return priorities[index % len(priorities)]
                elif 'version' in header_lower:
                    return f"v{index + 1}.{index + 2}"
            
            # Generic fallback
            else:
                if any(word in header_lower for word in ['name', 'title', 'category', 'type']):
                    return f"Item {index + 1}"
                elif any(word in header_lower for word in ['value', 'score', 'number', 'count']):
                    return str((index + 1) * 10)
                elif any(word in header_lower for word in ['description', 'detail', 'info']):
                    return f"Sample description for {header}"
                else:
                    return f"Sample {header} {index + 1}"
            
        except Exception as e:
            print(f"Error generating contextual sample data: {e}")
            return f"Sample {header} {index + 1}"

    def _create_contextual_sample_data(self, custom_headers, rag_response):
        """Create contextually relevant sample data based on headers and content"""
        try:
            data_rows = []
            rag_lower = rag_response.lower()
            
            # Determine context and number of sample rows
            if any(word in rag_lower for word in ['threat', 'attack', 'security']):
                sample_count = 4  # Security threats
            elif any(word in rag_lower for word in ['sales', 'revenue', 'business']):
                sample_count = 4  # Business data
            elif any(word in rag_lower for word in ['system', 'component', 'device']):
                sample_count = 4  # Technical data
            else:
                sample_count = 5  # Generic data
            
            for i in range(sample_count):
                sample_row = []
                for header in custom_headers:
                    sample_row.append(self._generate_contextual_sample_data(header, i, rag_response))
                data_rows.append(sample_row)
            
            return data_rows
            
        except Exception as e:
            print(f"Error creating contextual sample data: {e}")
            return []

    def _create_summary_data(self, rag_response, cited_pages, page_scores):
        """Create summary data when no structured data is found"""
        try:
            data_rows = []
            
            # Add analysis summary
            data_rows.append(['Analysis Type', 'Comprehensive Review', 'AI-powered document analysis'])
            
            # Add source count
            data_rows.append(['Sources Analyzed', str(len(cited_pages)), f"From {len(set([p.split(' from ')[1] for p in cited_pages if ' from ' in p]))} collections"])
            
            # Add average relevance score
            if page_scores:
                avg_score = sum(page_scores) / len(page_scores)
                data_rows.append(['Average Relevance', f"{avg_score:.3f}", 'Based on AI relevance scoring'])
            
            # Add response length
            data_rows.append(['Response Length', f"{len(rag_response)} characters", 'Comprehensive analysis provided'])
            
            return data_rows
            
        except Exception as e:
            print(f"Error creating summary data: {e}")
            return [['Analysis', 'Completed', 'Data extracted successfully']]
    
    def _populate_data_sheet(self, sheet, structured_data, query):
        """Populate the data sheet with structured information"""
        try:
            # Add title
            sheet['A1'] = f"Data Export for Query: {query}"
            sheet['A1'].font = Font(bold=True, size=14)
            sheet['A1'].fill = PatternFill(start_color="2F5496", end_color="2F5496", fill_type="solid")
            sheet['A1'].font = Font(color="FFFFFF", bold=True)
            
            # Add headers
            headers = structured_data['headers']
            for col, header in enumerate(headers, 1):
                cell = sheet.cell(row=3, column=col, value=header)
                cell.font = Font(bold=True)
                cell.fill = PatternFill(start_color="D9E2F3", end_color="D9E2F3", fill_type="solid")
                cell.border = Border(
                    left=Side(style='thin'),
                    right=Side(style='thin'),
                    top=Side(style='thin'),
                    bottom=Side(style='thin')
                )
            
            # Add data
            data = structured_data['data']
            for row_idx, row_data in enumerate(data, 4):
                for col_idx, value in enumerate(row_data, 1):
                    cell = sheet.cell(row=row_idx, column=col_idx, value=value)
                    cell.border = Border(
                        left=Side(style='thin'),
                        right=Side(style='thin'),
                        top=Side(style='thin'),
                        bottom=Side(style='thin')
                    )
            
            # Auto-adjust column widths
            for column in sheet.columns:
                max_length = 0
                column_letter = column[0].column_letter
                for cell in column:
                    try:
                        if len(str(cell.value)) > max_length:
                            max_length = len(str(cell.value))
                    except:
                        pass
                adjusted_width = min(max_length + 2, 50)
                sheet.column_dimensions[column_letter].width = adjusted_width
            
        except Exception as e:
            print(f"Error populating data sheet: {e}")
    
    def _populate_summary_sheet(self, sheet, query, cited_pages, page_scores):
        """Populate the summary sheet with analysis overview"""
        try:
            # Add title
            sheet['A1'] = "Analysis Summary"
            sheet['A1'].font = Font(bold=True, size=16)
            sheet['A1'].fill = PatternFill(start_color="2F5496", end_color="2F5496", fill_type="solid")
            sheet['A1'].font = Font(color="FFFFFF", bold=True)
            
            # Add query information
            sheet['A3'] = "Query:"
            sheet['A3'].font = Font(bold=True)
            sheet['B3'] = query
            
            # Add analysis statistics
            sheet['A5'] = "Analysis Statistics:"
            sheet['A5'].font = Font(bold=True)
            
            sheet['A6'] = "Sources Analyzed:"
            sheet['B6'] = len(cited_pages)
            
            sheet['A7'] = "Collections Used:"
            collections = set([p.split(' from ')[1] for p in cited_pages if ' from ' in p])
            sheet['B7'] = len(collections)
            
            if page_scores:
                sheet['A8'] = "Average Relevance Score:"
                avg_score = sum(page_scores) / len(page_scores)
                sheet['B8'] = f"{avg_score:.3f}"
            
            sheet['A9'] = "Analysis Date:"
            sheet['B9'] = datetime.now().strftime('%B %d, %Y at %I:%M %p')
            
            # Add source details
            sheet['A11'] = "Source Details:"
            sheet['A11'].font = Font(bold=True)
            
            for i, (citation, score) in enumerate(zip(cited_pages, page_scores)):
                row = 12 + i
                sheet[f'A{row}'] = f"Source {i+1}:"
                sheet[f'B{row}'] = citation
                sheet[f'C{row}'] = f"Score: {score:.3f}"
            
            # Auto-adjust column widths
            for column in sheet.columns:
                max_length = 0
                column_letter = column[0].column_letter
                for cell in column:
                    try:
                        if len(str(cell.value)) > max_length:
                            max_length = len(str(cell.value))
                    except:
                        pass
                adjusted_width = min(max_length + 2, 50)
                sheet.column_dimensions[column_letter].width = adjusted_width
            
        except Exception as e:
            print(f"Error populating summary sheet: {e}")
    
    def _create_excel_charts(self, sheet, structured_data, query, custom_headers=None):
        """Create Excel charts based on the data with custom headers"""
        try:
            # Add title
            sheet['A1'] = "Data Visualizations"
            sheet['A1'].font = Font(bold=True, size=16)
            sheet['A1'].fill = PatternFill(start_color="2F5496", end_color="2F5496", fill_type="solid")
            sheet['A1'].font = Font(color="FFFFFF", bold=True)
            
            # Determine chart titles and axis labels based on custom headers
            if custom_headers and len(custom_headers) >= 2:
                # Use custom headers for chart configuration
                x_axis_title = custom_headers[0] if len(custom_headers) > 0 else "Categories"
                y_axis_title = custom_headers[1] if len(custom_headers) > 1 else "Values"
                
                # Create more descriptive chart title based on context
                if len(custom_headers) >= 3:
                    chart_title = f"Analysis: {x_axis_title} vs {y_axis_title} by {custom_headers[2]}"
                else:
                    chart_title = f"Analysis: {x_axis_title} vs {y_axis_title}"
                
                # Create bar chart with custom headers
                if len(structured_data['data']) > 1:
                    chart = BarChart()
                    chart.title = chart_title
                    chart.x_axis.title = x_axis_title
                    chart.y_axis.title = y_axis_title
                    
                    # Add chart to sheet
                    sheet.add_chart(chart, "A3")
                
                # Create pie chart with custom header if we have 3+ columns
                if len(structured_data['data']) > 2 and len(custom_headers) >= 3:
                    pie_chart = PieChart()
                    pie_chart.title = f"Distribution by {custom_headers[2]}"
                    
                    # Add pie chart to sheet
                    sheet.add_chart(pie_chart, "A15")
                elif len(structured_data['data']) > 2:
                    # Fallback pie chart
                    pie_chart = PieChart()
                    pie_chart.title = "Data Distribution"
                    sheet.add_chart(pie_chart, "A15")
            else:
                # Use default chart configuration
                if len(structured_data['data']) > 1:
                    chart = BarChart()
                    chart.title = f"Analysis Results for: {query[:30]}..."
                    chart.x_axis.title = "Categories"
                    chart.y_axis.title = "Values"
                    
                    # Add chart to sheet
                    sheet.add_chart(chart, "A3")
                
                # Create pie chart for source distribution
                if len(structured_data['data']) > 2:
                    pie_chart = PieChart()
                    pie_chart.title = "Data Distribution"
                    
                    # Add pie chart to sheet
                    sheet.add_chart(pie_chart, "A15")
            
        except Exception as e:
            print(f"Error creating Excel charts: {e}")
    
    def _prepare_doc_download(self, doc_filepath):
        """
        Prepare DOC file for download in Gradio
        """
        if doc_filepath and os.path.exists(doc_filepath):
            return doc_filepath
        else:
            return None
    
    def _prepare_excel_download(self, excel_filepath):
        """
        Prepare Excel file for download in Gradio
        """
        if excel_filepath and os.path.exists(excel_filepath):
            return excel_filepath
        else:
            return None
    
    def _generate_multi_page_response(self, query, img_paths, cited_pages, page_scores):
        """
        Enhanced RAG response generation with multi-page citations
        Implements comprehensive detail enhancement based on research strategies
        """
        try:
            # Strategy 1: Increase context by providing more detailed prompt
            detailed_prompt = f"""
            Please provide a comprehensive and detailed answer to the following query. 
            Use all available information from the provided document pages to give a thorough response.
            
            Query: {query}
            
            Instructions for detailed response:
            1. Provide extensive background information and context
            2. Include specific details, examples, and data points from the documents
            3. Explain concepts thoroughly with step-by-step breakdowns 
            4. Provide comprehensive analysis rather than simple answers when requested
            
            """
            
            # Generate base response with enhanced prompt
            rag_response = rag.get_answer_from_gemini(detailed_prompt, img_paths)
            
            # Strategy 2: Simple citation formatting without relevance scores
            citation_text = "📚 **Sources**:\n\n"
            
            # Group citations by collection for better organization
            collection_groups = {}
            for i, citation in enumerate(cited_pages):
                collection_name = citation.split(" from ")[1].split(" (")[0]
                if collection_name not in collection_groups:
                    collection_groups[collection_name] = []
                collection_groups[collection_name].append(citation)
            
            # Format citations by collection (without relevance scores)
            for collection_name, citations in collection_groups.items():
                citation_text += f"📁 **{collection_name}**:\n"
                for citation in citations:
                    # Remove relevance score from citation
                    clean_citation = citation.split(" (Relevance:")[0]
                    citation_text += f"  • {clean_citation}\n"
                citation_text += "\n"
            
            # Strategy 3: Check for different export requests
            csv_filepath = None
            doc_filepath = None
            excel_filepath = None
            
            # Check if user requested table format
            if self._detect_table_request(query):
                print("📊 Table request detected - generating CSV response")
                enhanced_rag_response, csv_filepath = self._generate_csv_table_response(query, rag_response, cited_pages, page_scores)
            else:
                enhanced_rag_response = rag_response
            
            # Check if user requested comprehensive report
            if self._detect_report_request(query):
                print("📄 Report request detected - generating DOC report")
                doc_filepath, doc_error = self._generate_comprehensive_doc_report(query, rag_response, cited_pages, page_scores)
                if doc_error:
                    print(f"⚠️ DOC report generation failed: {doc_error}")
            
            # Check if user requested charts/graphs or enhanced Excel export
            if self._detect_chart_request(query) or self._detect_table_request(query):
                print("📊 Chart/Excel request detected - generating enhanced Excel export")
                # Extract custom headers for Excel export
                excel_custom_headers = self._extract_custom_headers(query)
                excel_filepath, excel_error = self._generate_enhanced_excel_export(query, rag_response, cited_pages, page_scores, excel_custom_headers)
                if excel_error:
                    print(f"⚠️ Excel export generation failed: {excel_error}")
            
            # Strategy 4: Combine sections for clean response with export information
            export_info = ""
            
            if doc_filepath:
                export_info += f"""
📄 **Comprehensive Report Generated**:
• **Format**: Microsoft Word Document (.docx)
• **Content**: Executive summary, detailed analysis, methodology, findings, and appendices
• **Download**: Available below
"""
            
            if excel_filepath:
                export_info += f"""
📊 **Enhanced Excel Export Generated**:
• **Format**: Microsoft Excel (.xlsx)
• **Content**: Multiple sheets with data, summary, and charts
• **Features**: Formatted tables, auto-generated charts, source analysis
• **Download**: Available below
"""
            
            if csv_filepath:
                export_info += f"""
📋 **CSV Table Generated**:
• **Format**: Comma-Separated Values (.csv)
• **Content**: Structured data table
• **Download**: Available below
"""
            
            final_response = f"""
{enhanced_rag_response}

{citation_text}

{export_info}
"""
            
            return final_response, csv_filepath, doc_filepath, excel_filepath
            
        except Exception as e:
            print(f"Error generating multi-page response: {e}")
            # Fallback to simple response with enhanced prompt
            return rag.get_answer_from_gemini(detailed_prompt, img_paths), None, None, None
        
    def authenticate_user(self, username, password):
        """Authenticate user and create session"""
        user_info = self.db_manager.authenticate_user(username, password)
        if user_info:
            session_id = self.session_manager.create_session(user_info)
            return f"Welcome {user_info['username']} from {user_info['team']}!", session_id, user_info['team']
        else:
            return "Invalid username or password", None, None
    
    def logout_user(self, session_id):
        """Logout user and remove session"""
        if session_id:
            self.session_manager.remove_session(session_id)
        return "Logged out successfully", None, None
    

    def get_team_collections(self, session_id):
        """Get available collections for the user's team"""
        if not session_id:
            return "Please log in to view team collections"
        
        session = self.session_manager.get_session(session_id)
        if not session:
            return "Session expired. Please log in again."
        
        team = session['user_info']['team']
        collections = self.db_manager.get_team_collections(team)
        
        if not collections:
            return f"No collections found for {team}"
        
        return f"**{team} Collections:**\n" + "\n".join([f"- {coll}" for coll in collections])

def create_ui():
    app = PDFSearchApp()
    
    with gr.Blocks(theme=gr.themes.Ocean(), css="footer{display:none !important}") as demo:
        # Session state management
        session_state = gr.State(value=None)
        user_info_state = gr.State(value=None)
        
        gr.Markdown("# Collar Multimodal RAG Demo - Streamlined")
        gr.Markdown("Made by Collar - Document Upload and Query System")
        
        # Authentication Tab
        with gr.Tab("🔐 Authentication"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### Login")
                    username_input = gr.Textbox(label="Username", placeholder="Enter username")
                    password_input = gr.Textbox(label="Password", type="password", placeholder="Enter password")
                    login_btn = gr.Button("Login", variant="primary")
                    logout_btn = gr.Button("Logout")
                    auth_status = gr.Textbox(label="Authentication Status", interactive=False)
                    current_team = gr.Textbox(label="Current Team", interactive=False)
                
                with gr.Column(scale=1):
                    gr.Markdown("### Default Users")
                    gr.Markdown("""
                    **Team A:** admin_team_a / admin123_team_a
                    **Team B:** admin_team_b / admin123_team_b
                    """)
        
        # Document Management Tab
        with gr.Tab("📁 Document Management"):
            with gr.Column():
                gr.Markdown("### Upload Documents to Team Repository")
                folder_name_input = gr.Textbox(
                    label="Folder/Collection Name (Optional)", 
                    placeholder="Enter a name for this document collection"
                )
                max_pages_input = gr.Slider(
                    minimum=1,
                    maximum=10000,
                    value=20,
                    step=10,
                    label="Max pages to extract and index per document"
                )
                file_input = gr.Files(
                    label="Upload PPTs/PDFs (Multiple files supported)", 
                    file_count="multiple"
                )
                upload_btn = gr.Button("Upload to Repository", variant="primary")
                upload_status = gr.Textbox(label="Upload Status", interactive=False)
                
                gr.Markdown("### Team Collections")
                refresh_collections_btn = gr.Button("Refresh Collections")
                team_collections_display = gr.Textbox(
                    label="Available Collections", 
                    interactive=False, 
                    lines=5
                )
        
        # Enhanced Query Tab
        with gr.Tab("🔍 Advanced Query"):
            with gr.Column():
                gr.Markdown("### Multi-Page Document Search")
                
                query_input = gr.Textbox(
                    label="Enter your query", 
                    placeholder="Ask about any topic in your documents...",
                    lines=2
                )
                num_results = gr.Slider(
                    minimum=1,
                    maximum=10,
                    value=3,
                    step=1,
                    label="Number of pages to retrieve and cite"
                )
                search_btn = gr.Button("Search Documents", variant="primary")
                
                gr.Markdown("### Results")
                llm_answer = gr.Textbox(
                    label="AI Response with Citations", 
                    interactive=False,
                    lines=8
                )
                cited_pages_display = gr.Textbox(
                    label="Cited Pages", 
                    interactive=False,
                    lines=3
                )
                path = gr.Textbox(label="Document Paths", interactive=False)
                images = gr.Gallery(label="Retrieved Pages", show_label=True, columns=2, rows=2, height="auto")
                
                # Export Downloads Section
                gr.Markdown("### 📊 Export Downloads")
                
                with gr.Row():
                    with gr.Column(scale=1):
                        csv_download = gr.File(
                            label="📋 CSV Table",
                            interactive=False,
                            visible=True
                        )
                    with gr.Column(scale=1):
                        doc_download = gr.File(
                            label="📄 DOC Report",
                            interactive=False,
                            visible=True
                        )
                    with gr.Column(scale=1):
                        excel_download = gr.File(
                            label="📊 Excel Export",
                            interactive=False,
                            visible=True
                        )
        

        

                





        
        # Event handlers
        # Authentication events
        login_btn.click(
            fn=app.authenticate_user,
            inputs=[username_input, password_input],
            outputs=[auth_status, session_state, current_team]
        )
        
        logout_btn.click(
            fn=app.logout_user,
            inputs=[session_state],
            outputs=[auth_status, session_state, current_team]
        )
        
        # Document management events
        upload_btn.click(
            fn=app.upload_and_convert,
            inputs=[session_state, file_input, max_pages_input, session_state, folder_name_input],
            outputs=[upload_status]
        )
        
        refresh_collections_btn.click(
            fn=app.get_team_collections,
            inputs=[session_state],
            outputs=[team_collections_display]
        )
        
        # Query events
        search_btn.click(
            fn=app.search_documents,
            inputs=[session_state, query_input, num_results, session_state],
            outputs=[path, images, llm_answer, cited_pages_display, csv_download, doc_download, excel_download]
        )
        

        

        



        

        
    return demo

if __name__ == "__main__":
    demo = create_ui()
    #demo.launch(auth=("admin", "pass1234")) for with login page config
    demo.launch()