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 queue
import requests
import base64
from PIL import Image
import io
import schemdraw
import schemdraw.elements as elm
import matplotlib.pyplot as plt
from PIL import Image
import io
import schemdraw
import schemdraw.elements as elm
import matplotlib.pyplot as plt

from middleware import Middleware
from rag import Rag
from pathlib import Path
import subprocess
import getpass
# 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
            )
        ''')
        
        # Chat history table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS chat_history (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER,
                query TEXT NOT NULL,
                response TEXT NOT NULL,
                cited_pages TEXT,
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (user_id) REFERENCES users (id)
            )
        ''')
        
        # 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_chat_history(self, user_id: int, query: str, response: str, cited_pages: List[str]):
        """Save chat interaction to database"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cited_pages_json = json.dumps(cited_pages)
            cursor.execute(
                'INSERT INTO chat_history (user_id, query, response, cited_pages) VALUES (?, ?, ?, ?)',
                (user_id, query, response, cited_pages_json)
            )
            conn.commit()
            conn.close()
        except Exception as e:
            print(f"Error saving chat history: {e}")
    
    def get_chat_history(self, user_id: int, limit: int = 10) -> List[Dict]:
        """Get recent chat history for user"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT query, response, cited_pages, timestamp 
                FROM chat_history 
                WHERE user_id = ? 
                ORDER BY timestamp DESC 
                LIMIT ?
            ''', (user_id, limit))
            
            history = []
            for row in cursor.fetchall():
                history.append({
                    'query': row[0],
                    'response': row[1],
                    'cited_pages': json.loads(row[2]) if row[2] else [],
                    'timestamp': row[3]
                })
            
            conn.close()
            return history
        except Exception as e:
            print(f"Error getting chat history: {e}")
            return []
    
    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 []
    
    def clear_chat_history(self, user_id: int) -> bool:
        """Clear all chat history for a user"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('DELETE FROM chat_history WHERE user_id = ?', (user_id,))
            conn.commit()
            conn.close()
            return True
        except Exception as e:
            print(f"Error clearing chat history: {e}")
            return False

# 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)

    if is_ollama_running():
        print("Ollama server is already running.")
    else:
        print("Ollama server is not running. Starting it...")
        start_ollama()

    # --- 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.")

    
start_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)
            
            # Save chat history if user is logged in
            if user_info:
                self.db_manager.save_chat_history(
                    user_info['id'], 
                    query, 
                    rag_response, 
                    cited_pages
                )
            
            # 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_openai(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_openai(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_chat_history(self, session_id, limit=10):
        """Get chat history for logged-in user in a user-friendly format"""
        if not session_id:
            return "🔐 **Please log in to view chat history**"
        
        session = self.session_manager.get_session(session_id)
        if not session:
            return "⏰ **Session expired. Please log in again.**"
        
        user_info = session['user_info']
        history = self.db_manager.get_chat_history(user_info['id'], limit)
        
        if not history:
            return "📭 **No chat history found.**\n\nStart a conversation to see your chat history here!"
        
        # Format timestamp for better readability
        def format_timestamp(timestamp_str):
            try:
                # Parse the timestamp and format it nicely
                dt = datetime.fromisoformat(timestamp_str.replace('Z', '+00:00'))
                return dt.strftime("%B %d, %Y at %I:%M %p")
            except:
                return timestamp_str
        
        # Truncate response for better display
        def truncate_response(response, max_length=300):
            if len(response) <= max_length:
                return response
            return response[:max_length] + "..."
        
        history_text = f"""
# 💬 Chat History for {user_info['username']} ({user_info['team']})

📊 **Showing last {len(history)} conversations**

---
"""
        
        for i, entry in enumerate(reversed(history), 1):  # Show newest first
            # Format the conversation entry
            conversation_entry = f"""
## 🗨️ Conversation #{len(history) - i + 1}

**❓ Your Question:**
{entry['query']}

**🤖 AI Response:**
{truncate_response(entry['response'])}

**📄 Sources Referenced:**
{', '.join(entry['cited_pages']) if entry['cited_pages'] else 'No specific pages cited'}

**📅 Date:** {format_timestamp(entry['timestamp'])}

---
"""
            history_text += conversation_entry
        
        # Add summary at the end
        history_text += f"""
## 📈 Summary
• **Total Conversations:** {len(history)}
• **Date Range:** {format_timestamp(history[-1]['timestamp'])} to {format_timestamp(history[0]['timestamp'])}
• **Team:** {user_info['team']}
• **User:** {user_info['username']}
"""
        
        return history_text
    
    def clear_chat_history(self, session_id):
        """Clear chat history for logged-in user"""
        if not session_id:
            return "🔐 **Please log in to manage chat history**"
        
        session = self.session_manager.get_session(session_id)
        if not session:
            return "⏰ **Session expired. Please log in again.**"
        
        user_info = session['user_info']
        success = self.db_manager.clear_chat_history(user_info['id'])
        
        if success:
            return "🗑️ **Chat history cleared successfully!**\n\nYour conversation history has been removed."
        else:
            return "❌ **Error clearing chat history.**\n\nPlease try again or contact support."
    
    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 delete(self, state, choice, session_id=None):
        """Delete collection with team-based access control"""
        if session_id:
            session = self.session_manager.get_session(session_id)
            if not session:
                return "Session expired. Please log in again."
            
            team = session['user_info']['team']
            # Only allow deletion if collection belongs to user's team
            team_collections = self.db_manager.get_team_collections(team)
            if choice not in team_collections:
                return f"Access denied. Collection {choice} does not belong to {team}"
        
        # Delete file in pages, then use middleware to delete collection 
        client = MilvusClient(
                uri="http://localhost:19530",
                token="root:Milvus"
            )
        path = f"pages/{choice}"
        if os.path.exists(path):
            shutil.rmtree(path)
            # Call milvus manager to delete collection
            client.drop_collection(collection_name=choice)
            return f"Deleted {choice}"
        else:
            return "Directory not found"

        

            


    def describe_image_with_gemma3(self, image):
        """Describe image using Gemma3 vision model via Ollama"""
        try:
            print("🔍 [CIRCUIT] Starting image description with Gemma3...")
            
            if image is None:
                print("❌ [CIRCUIT] No image provided")
                return "No image provided"
            
            print("📸 [CIRCUIT] Converting image to base64...")
            # Convert PIL image to base64
            buffered = io.BytesIO()
            image.save(buffered, format="PNG")
            img_str = base64.b64encode(buffered.getvalue()).decode()
            print("✅ [CIRCUIT] Image converted successfully")
            
            # Prepare request for Ollama Gemma3
            print("🤖 [CIRCUIT] Preparing request for Gemma3 model...")
            payload = {
                "model": "gemma3:4b",
                "prompt": "Just generate a netlist of circuit components of the image with explanations ONLY, NO OTHER TEXT",
                "images": [img_str],
                "stream": False
            }
            
            print("🚀 [CIRCUIT] Sending request to Ollama Gemma3...")
            # Send request to Ollama
            response = requests.post("http://localhost:11434/api/generate", json=payload, timeout=1200)
            
            if response.status_code == 200:
                result = response.json()
                description = result.get('response', 'No description generated')
                print(f"✅ [CIRCUIT] Image description completed successfully")
                print(f"📝 [CIRCUIT] Description length: {len(description)} characters")
                return description
            else:
                error_msg = f"Error: {response.status_code} - {response.text}"
                print(f"❌ [CIRCUIT] {error_msg}")
                return error_msg
                
        except Exception as e:
            error_msg = f"Error describing image: {str(e)}"
            print(f"❌ [CIRCUIT] {error_msg}")
            return error_msg

    def generate_circuit_with_deepseek(self, image_description, max_retries=3):
        """Generate netlist and circuit diagram using DeepSeek R1 with error handling and retry logic"""
        previous_error = None
        consecutive_failures = 0
        
        for attempt in range(max_retries):
            try:
                print(f"🔧 [CIRCUIT] Starting circuit generation with DeepSeek R1 (Attempt {attempt + 1}/{max_retries})...")
                
                if not image_description or image_description == "No image provided":
                    print("❌ [CIRCUIT] No image description available")
                    return "No image description available"
                
                print("📝 [CIRCUIT] Preparing prompt for DeepSeek R1...")
                
                # Use retry prompt if this is not the first attempt
                if attempt == 0:
                    # Generate unique filename for this attempt
                    unique_filename = self._generate_unique_filename()
                    
                    # Parse complex circuit description if available
                    circuit_data = self._parse_complex_circuit_description(image_description)
                    
                    # Use specialized prompt for complex circuits if parsing was successful
                    if circuit_data and circuit_data.get('complexity_level') in ['complex', 'very_complex']:
                        print(f"Using specialized prompt for {circuit_data['complexity_level']} circuit")
                        prompt = self._generate_complex_circuit_prompt(circuit_data, unique_filename)
                        if not prompt:
                            # Fallback to standard prompt if specialized prompt generation fails
                            prompt = f"""Generate a complex circuit diagram using the python schemdraw library based on this detailed description. 

COMPLEX CIRCUIT REQUIREMENTS:
1. **Component Mapping**: Map ALL components from the description to schemdraw equivalents:
   - Resistors: elm.Resistor with proper values
   - Capacitors: elm.Capacitor with proper values  
   - Inductors: elm.Inductor with proper values
   - Diodes: elm.Diode, elm.LED, elm.Zener with proper types
   - Transistors: elm.Transistor, elm.BjtNpn, elm.BjtPnp, elm.FetN, elm.FetP
   - ICs: elm.RBox with proper labels and pin configurations
   - Power sources: elm.SourceV, elm.Battery, elm.SourceSin, elm.SourceSquare
   - Switches: elm.Switch, elm.SwitchSpdt
   - Connectors: elm.Connector, elm.Dot for connection points

2. **Complex Topology Handling**:
   - Use elm.Dot for wire junctions and connection points
   - Use elm.Line for explicit wire connections
   - Use elm.Label for power rails and voltage/current labels
   - Use elm.Text for component labels and values
   - Use elm.Node for connection nodes
   - Handle multiple power rails (VCC, GND, VDD, etc.)
   - Support feedback loops and control paths
   - Handle parallel and series connections properly

3. **Advanced Positioning**:
   - Use .up(), .down(), .left(), .right() for basic positioning
   - Use .to() for precise connections: .to(d.elements[0].start)
   - Use .at() for absolute positioning when needed
   - Use .move() for relative positioning
   - Arrange components in logical blocks and sections
   - Use consistent spacing and alignment

4. **Component Labeling**:
   - Label ALL components with their values and designators
   - Use .label() method for component values
   - Use elm.Text for additional labels and annotations
   - Include voltage/current ratings where applicable
   - Add pin numbers for ICs and connectors

5. **Circuit Organization**:
   - Group related components together
   - Use clear signal flow from left to right or top to bottom
   - Separate power supply sections from signal processing
   - Use consistent naming conventions
   - Minimize wire crossings and clutter

IMPORTANT REQUIREMENTS:
1. Use ONLY ASCII characters - replace Ω with 'Ohm', μ with 'u', ° with 'deg'
2. Use ONLY components available in schemdraw.elements library
3. If a component is not in schemdraw.elements, use elm.RBox and label it appropriately
4. Do NOT use matplotlib or any other plotting library
5. Generate a complete, executable Python script
6. ALWAYS use d.save() to save the diagram, NEVER use d.draw()
7. Save the output as a PNG file with the EXACT filename: {unique_filename}
8. Handle all connections properly using schemdraw's native positioning methods
9. Create a functional circuit that matches the description - all components must be properly connected
10. INCLUDE ALL COMPONENTS mentioned in the description - do not miss any components
11. Use .to() method for precise connections and circuit completion
12. Support complex topologies with multiple power rails and signal paths
13. NEVER use d.element - this is INVALID and will cause errors
14. NEVER use d.last_end, d.last_start, d.end, d.start, d.position - these are INVALID attributes
15. CRITICAL: If you use d.element, the circuit will fail validation and not be generated

Description of the circuit: {image_description}

CORRECT SCHEMDRAW API USAGE:
- Use d += elm.Component() to add components
- Use .up(), .down(), .left(), .right() for positioning
- Use .to() to connect to specific points: .to(d.elements[0].start)
- Use .label() to add labels: .label('10V')
- Use .at() for absolute positioning: .at((x, y))
- Use d.save() to save the diagram
- Use elm.Dot for connection points
- NEVER use d.element - this is INVALID and will cause errors
- ALWAYS use d.elements[-1] instead of d.element
- NEVER use d.last_end, d.last_start, d.end, d.start, d.position - these are INVALID attributes
- Use elm.Line for explicit wire connections
- Use elm.Text for additional labels
- DO NOT use: d.last_end, d.last_start, d.end, d.start, d.position, d.element

COMPLEX CIRCUIT EXAMPLE (for reference only):
```python
import schemdraw
import schemdraw.elements as elm

d = schemdraw.Drawing()
# Power supply section
d += elm.SourceV().up().label('12V').at((0, 0))
d += elm.Resistor().right().label('1KOhm')
d += elm.Capacitor().down().label('100uF')
d += elm.Line().left().to(d.elements[0].start)  # Close main loop

# Signal processing section
d += elm.Dot().at((4, 0))
d += elm.Transistor().up().label('Q1')
d += elm.Resistor().right().label('10KOhm')
d += elm.Line().down().to(d.elements[-2].start)  # Close secondary loop
d += elm.Line().left().to(d.elements[0].start)   # Ensure complete closure
d.save('{unique_filename}')
```

IMPORTANT: Always use .to(d.elements[0].start) to close the circuit loop back to the power source!

CRITICAL REQUIREMENTS:
- Create a circuit that accurately represents the complex description provided
- Use appropriate components and values that match the actual circuit described
- INCLUDE ALL COMPONENTS listed above - missing components will cause validation failure
- Ensure all components are properly connected and labeled
- Handle complex topologies with multiple power rails and signal paths
- Use proper component positioning and wire routing
- Support feedback loops, control paths, and complex connections
- Arrange components logically with clear signal flow
- Use consistent labeling and naming conventions
- Minimize wire clutter while maintaining circuit clarity

CRITICAL CIRCUIT CLOSURE REQUIREMENTS:
- ALWAYS close the circuit loop using .to() method: d += elm.Line().to(d.elements[0].start)
- Ensure ALL components are connected in a complete loop
- Use explicit Line() elements to connect components when needed
- Start with a power source (elm.SourceV, elm.Battery)
- End with a connection back to the power source
- Use proper positioning to create logical circuit flow
- For complex circuits, use multiple .to() connections to ensure complete closure
"""
                    else:
                        # Use standard prompt for simple circuits
                        prompt = f"""Generate a circuit diagram using the python schemdraw library based on this description. 
                    
IMPORTANT REQUIREMENTS:
1. Use ONLY ASCII characters - replace Ω with 'Ohm', μ with 'u', ° with 'deg'
2. Use ONLY components available in schemdraw.elements library
3. If a component is not in schemdraw.elements, use a RBOX element (schemdraw.elements.twoterm.RBox) and label it with the component name
4. Do NOT use matplotlib or any other plotting library
5. Generate a complete, executable Python script
6. Use d.save() to save the diagram, NOT d.draw()
7. Save the output as a PNG file with the EXACT filename: {unique_filename}
8. Handle all connections properly using schemdraw's native positioning methods
9. Create a CLOSED LOOP circuit that matches the description - all components must form a complete loop
10. INCLUDE ALL COMPONENTS mentioned in the description - do not miss any components
11. DO NOT use any grounding elements (elm.Ground, elm.GroundChassis, etc.) - create a complete closed loop circuit
12. Use .to() method to explicitly close the circuit loop back to the starting point

Description of the circuit: {image_description}

CORRECT USAGE EXAMPLE (for reference only):
import schemdraw
import schemdraw.elements as elm

d = schemdraw.Drawing()
d += elm.SourceV().up().label('10V')
d += elm.Resistor().right().label('100KOhm')
d += elm.Capacitor().down().label('0.1uF')
d += elm.Line().left().to(d.elements[0].start)  # Clean connection back to voltage source
d.save('{unique_filename}')

IMPORTANT: Always use .to(d.elements[0].start) to close the circuit loop back to the power source!

CRITICAL REQUIREMENTS:
- Do NOT copy the example circuit above
- Create a completely different circuit that accurately represents the description provided
- Use different components, values, and layout that match the actual circuit described in the image
- INCLUDE ALL COMPONENTS listed above - missing components will cause validation failure
- Ensure all components are properly connected and labeled
- ENSURE COMPLETE CIRCUIT CONNECTIVITY - all components must form a connected, working circuit
- Include power sources (voltage/current sources) and ground connections where appropriate
- Use explicit Line() elements to connect components when needed
- Create logical circuit flow with proper component sequencing
- MINIMIZE WIRE CLUTTER - use direct component connections instead of unnecessary Line() elements
- Use net labels (VoltageLabel, CurrentLabel) for power rails instead of long wires
- Arrange components in clean, symmetrical layouts with consistent spacing
- Use horizontal and vertical connections only - avoid diagonal wires
- ENSURE COMPLETE CIRCUIT CONNECTIVITY - all components must form a connected, working circuit
- Include power sources (voltage/current sources) and ground connections where appropriate
- Use explicit Line() elements to connect components when needed
- Create a logical circuit flow with proper component sequencing
- MINIMIZE UNNECESSARY WIRES - use net labels and direct connections instead of excessive Line() elements
- Use horizontal and vertical wire orientations only - avoid diagonal connections
- Limit wire junctions to maximum 3 connections per point
- Arrange components symmetrically and maintain consistent spacing

COMMON ERRORS TO AVOID:
- Do NOT use: elm.Tip, elm.DCSourceV, elm.SpiceNetlist
- Do NOT use: matplotlib, pyplot, or any plotting libraries
- Do NOT use Unicode characters in labels or component names
- Do NOT use components not in schemdraw.elements
- Do NOT use invalid assignment syntax like "light_bulb = d += elm.Lamp()" - use "d += elm.Lamp()" only
- Do NOT use any grounding elements (elm.Ground, elm.GroundChassis, elm.GroundSignal) - create closed loop circuits only
- Do NOT use excessive Line() elements - minimize unnecessary wires and use direct connections
- Do NOT use redundant wire patterns (up().down(), left().right(), etc.) - use efficient routing
- Do NOT use any other filename - use exactly: {unique_filename}
- Do NOT copy the example circuit - create your own unique design
- Do NOT miss any components from the description
- DO NOT use: elm.Lightbulb, use elm.Lamp instead!

CRITICAL CIRCUIT CLOSURE REQUIREMENTS:
- ALWAYS close the circuit loop using .to() method: d += elm.Line().to(d.elements[0].start)
- Ensure ALL components are connected in a complete loop
- Use explicit Line() elements to connect components when needed
- Start with a power source (elm.SourceV, elm.Battery)
- End with a connection back to the power source
- Use proper positioning to create logical circuit flow

Generate ONLY the Python code, no explanations or markdown formatting."""
                else:
                    # Use retry prompt with previous error information
                    prompt = self._create_retry_prompt(image_description, previous_error)
                
                # Send request to DeepSeek R1 via Ollama
                print("🤖 [CIRCUIT] Preparing request for Reasoning model...")
                payload = {
                    "model": "qwen3-coder:latest",
                    "prompt": prompt,
                    "stream": False,
                    #"think": True,
                    "temperature": 0.5,
                }
                
                print("🚀 [CIRCUIT] Sending request to Reasoning Model...")
                response = requests.post("http://localhost:11434/api/generate", json=payload, timeout=3000)
                
                if response.status_code == 200:
                    result = response.json()
                    generated_code = result.get('response', '')
                    print(f"✅ [CIRCUIT] DeepSeek R1 response received successfully")
                    print(f"📝 [CIRCUIT] Generated code length: {len(generated_code)} characters")
                    
                    # Extract Python code from markdown blocks if present
                    print("🔧 [CIRCUIT] Extracting Python code from response...")
                    extracted_code = self._extract_python_code(generated_code)
                    print(f"📝 [CIRCUIT] Extracted code length: {len(extracted_code)} characters")
                    
                    # Fix circuit structure and enhance connections
                    print("🔧 [CIRCUIT] Fixing circuit structure and enhancing connections...")
                    enhanced_code = self._fix_circuit_structure(extracted_code)
                    
                    # Validate the enhanced code
                    if not self._validate_circuit_code(enhanced_code):
                        print("⚠️ [CIRCUIT] Enhanced code validation failed, will retry...")
                        if attempt < max_retries - 1:
                            continue
                        else:
                            return "Error: Enhanced code failed validation after all retries"
                    
                    # Validate circuit connectivity

                    
                    # Execute the enhanced code
                    print("⚙️ [CIRCUIT] Executing enhanced circuit code...")
                    result = self._execute_generated_circuit_code(enhanced_code)
                    
                    # Check if execution was successful
                    if result and result.endswith('.png'):
                        print(f"✅ [CIRCUIT] Circuit generation successful on attempt {attempt + 1}")
                        consecutive_failures = 0  # Reset failure counter on success
                        
                        # Check if this was the final attempt
                        if attempt == max_retries - 1:
                            print("✅ [CIRCUIT] Circuit generated successfully")
                            return f"{result} (Note: Circuit generated successfully)"
                        
                        return result
                    else:
                        print(f"⚠️ [CIRCUIT] Circuit execution failed: {result}")
                        consecutive_failures += 1
                        previous_error = result
                        
                        # Circuit breaker: if too many consecutive failures, provide partial result
                        if consecutive_failures >= 2 and attempt == max_retries - 1:
                            print("⚠️ [CIRCUIT] Multiple consecutive failures detected, providing partial result...")
                            return f"Partial circuit generated (Note: Some components may be missing due to generation difficulties)"
                        
                        if attempt < max_retries - 1:
                            print(f"🔄 [CIRCUIT] Retrying... (Attempt {attempt + 2}/{max_retries})")
                            continue
                        else:
                            return f"Error: Circuit generation failed after {max_retries} attempts. Last error: {result}"
                else:
                    error_msg = f"Error: {response.status_code} - {response.text}"
                    print(f"❌ [CIRCUIT] {error_msg}")
                    previous_error = error_msg
                    if attempt < max_retries - 1:
                        print(f"🔄 [CIRCUIT] Retrying... (Attempt {attempt + 2}/{max_retries})")
                        continue
                    else:
                        return error_msg
                        
            except Exception as e:
                error_msg = f"Error generating circuit: {str(e)}"
                print(f"❌ [CIRCUIT] {error_msg}")
                previous_error = error_msg
                if attempt < max_retries - 1:
                    print(f"🔄 [CIRCUIT] Retrying... (Attempt {attempt + 2}/{max_retries})")
                    continue
                else:
                    return error_msg
        
        return f"Error: Circuit generation failed after {max_retries} attempts"

    def _create_retry_prompt(self, image_description, previous_error):
        """Create an enhanced prompt for retry attempts with error feedback"""
        # Generate unique filename for retry attempts
        unique_filename = self._generate_unique_filename()
        
        prompt = f"""The previous attempt to generate a circuit diagram failed. Please fix the issues and try again.

PREVIOUS ERROR: {previous_error}

IMPORTANT REQUIREMENTS (must follow exactly):
1. Use ONLY ASCII characters - replace Ω with 'Ohm', μ with 'u', ° with 'deg'
2. Use ONLY components available in schemdraw.elements library
3. If a component is not in schemdraw.elements, use a Rbox element (schemdraw.elements.twoterm.RBox) and label it with the component name
4. Do NOT use matplotlib or any other plotting library
5. Generate a complete, executable Python script
6. Use d.save() to save the diagram, NOT d.draw()
7. Save the output as a PNG file with the EXACT filename: {unique_filename}
8. Handle all connections properly using schemdraw's native positioning methods
9. Create a CLOSED LOOP circuit that matches the description - all components must form a complete loop
10. INCLUDE ALL COMPONENTS mentioned in the description - do not miss any components
11. DO NOT use any grounding elements (elm.Ground, elm.GroundChassis, etc.) - create a complete closed loop circuit
12. Use .to() method to explicitly close the circuit loop back to the starting point

Description of the circuit: {image_description}

CORRECT USAGE EXAMPLE (for reference only - create your own unique circuit):
```python
import schemdraw
import schemdraw.elements as elm

d = schemdraw.Drawing()
d += elm.SourceV().up().label('10V')
d += elm.Resistor().right().label('100KOhm')
d += elm.Capacitor().down().label('0.1uF')
d += elm.Line().left().to(d.elements[0].start)  # Close the loop back to voltage source
d.save('{unique_filename}')
```

IMPORTANT: Always use .to(d.elements[0].start) to close the circuit loop back to the power source!

CRITICAL REQUIREMENTS:
- Create a circuit that accurately represents the description provided
- Use different components, values, and layout that match the actual circuit described in the image
- INCLUDE ALL COMPONENTS listed above - missing components will cause validation failure
- Ensure all components are properly connected and labeled

COMMON ERRORS TO AVOID:
- Do NOT use: elm.Tip, elm.DCSourceV, elm.SpiceNetlist
- Do NOT use: matplotlib, pyplot, or any plotting libraries
- Do NOT use Unicode characters in labels or component names
- Do NOT use components not in schemdraw.elements
- Do NOT use invalid assignment syntax like "light_bulb = d += elm.Lamp()" - use "d += elm.Lamp()" only
- Do NOT use any other filename - use exactly: {unique_filename}
- Do NOT miss any components from the description

CRITICAL CIRCUIT CLOSURE REQUIREMENTS:
- ALWAYS close the circuit loop using .to() method: d += elm.Line().to(d.elements[0].start)
- Ensure ALL components are connected in a complete loop
- Use explicit Line() elements to connect components when needed
- Start with a power source (elm.SourceV, elm.Battery)
- End with a connection back to the power source
- Use proper positioning to create logical circuit flow

Generate ONLY the Python code, no explanations or markdown formatting."""
        return prompt

    def _cleanup_previous_circuit_files(self):
        """Clean up previous circuit diagram files to ensure fresh generation"""
        try:
            print("🧹 [CIRCUIT] Cleaning up previous circuit diagram files...")
            circuit_files = []
            
            # Find all PNG files that might be circuit diagrams
            for file in os.listdir('.'):
                if file.endswith('.png') and any(keyword in file.lower() for keyword in ['circuit', 'diagram', 'schematic']):
                    circuit_files.append(file)
            
            # Remove previous circuit diagram files
            for file in circuit_files:
                try:
                    os.remove(file)
                    print(f"🗑️ [CIRCUIT] Removed previous circuit file: {file}")
                except Exception as e:
                    print(f"⚠️ [CIRCUIT] Failed to remove {file}: {str(e)}")
            
            print(f"✅ [CIRCUIT] Cleaned up {len(circuit_files)} previous circuit files")
            
        except Exception as e:
            print(f"⚠️ [CIRCUIT] Error during cleanup: {str(e)}")

    def _generate_unique_filename(self):
        """Generate a unique filename for the circuit diagram"""
        import time
        timestamp = int(time.time())
        return f"circuit_diagram_{timestamp}.png"

    def _preprocess_circuit_image(self, image):
        """Preprocess circuit image for better component detection"""
        try:
            print("Preprocessing circuit image...")
            
            # Convert to RGB if needed
            if image.mode != 'RGB':
                image = image.convert('RGB')
            
            # Enhance image quality
            from PIL import ImageEnhance, ImageFilter
            
            # Increase contrast for better component visibility
            enhancer = ImageEnhance.Contrast(image)
            image = enhancer.enhance(1.5)
            
            # Sharpen image for clearer component boundaries
            image = image.filter(ImageFilter.SHARPEN)
            
            # Increase brightness slightly
            enhancer = ImageEnhance.Brightness(image)
            image = enhancer.enhance(1.2)
            
            print("Image preprocessing completed")
            return image
            
        except Exception as e:
            print(f"Image preprocessing failed: {str(e)}")
            return image  # Return original image if preprocessing fails

    def _parse_complex_circuit_description(self, image_description):
        """Parse complex circuit description and extract structured component information"""
        try:
            print("🔍 [CIRCUIT] Parsing complex circuit description...")
            
            # Initialize structured data
            circuit_data = {
                'components': [],
                'connections': [],
                'power_rails': [],
                'signal_paths': [],
                'circuit_function': '',
                'complexity_level': 'simple'
            }
            
            # Enhanced component detection for complex circuits
            import re
            
            # Enhanced component detection with comprehensive patterns
            component_patterns = [
                # Switches (DPDT and safety switches)
                r'\bSW\d+\b',  # SW1, SW2, SW3
                r'\bDPDT\b',   # DPDT switches
                r'\bswitch\b',
                r'\bsafety\s*switch\b',
                r'\barming\s*arm\b',
                r'\bSAKLAR\s*PENGAMAN\b',
                
                # Power sources (batteries and voltage sources)
                r'\bBAT\d+\b',  # BAT1, BAT2
                r'\bbattery\b',
                r'\b9V\b',
                r'\b12V\b',
                r'\bvoltage\s*source\b',
                r'\bpower\s*supply\b',
                r'\bVCC\b', r'\bGND\b', r'\bVDD\b', r'\bVSS\b',
                
                # Resistors (with specific values)
                r'\bR\d+\b',    # R1, R2, R3, R4, R5
                r'\bresistor\b',
                r'\b1k\b', r'\b2k\b', r'\b100\b', r'\b10k\b', r'\b100k\b',  # Common values
                r'\bohm\b', r'\bΩ\b',
                
                # LEDs (indicators and status lights)
                r'\bLED\s*D\d+\b',  # LED D1, LED D2, LED D3
                r'\bled\b',
                r'\bblue\b',
                r'\bindicator\b',
                r'\bstatus\s*light\b',
                r'\bIDIKATOR\b', r'\bINDIKATOR\b',
                
                # Active components (SCR, transistors, ICs)
                r'\bSCR\b',
                r'\bU\d+\b',    # U1
                r'\bSilicon\s*Controlled\s*Rectifier\b',
                r'\bthyristor\b',
                r'\btransistor\b',
                r'\bBJT\b', r'\bFET\b', r'\bMOSFET\b',
                r'\bopamp\b', r'\boperational\s*amplifier\b',
                r'\bIC\b', r'\bintegrated\s*circuit\b',
                
                # Special components (initiator, coils)
                r'\bL\d+\b',    # L1
                r'\binisiator\b',
                r'\binitiator\b',
                r'\bcoil\b',
                r'\b12V\s*inisiator\b',
                r'\binductor\b',
                
                # General components
                r'\bcapacitor\b', r'\bcondenser\b',
                r'\bdiode\b', r'\brectifier\b',
                r'\bwire\b', r'\bconnection\b',
                r'\bterminal\b', r'\bnode\b',
                r'\bground\b', r'\bearth\b',
                
                # Circuit sections and labels
                r'\binput\s*section\b',
                r'\bcontrol\s*section\b',
                r'\boutput\s*section\b',
                r'\bpower\s*rail\b',
                r'\bsignal\s*path\b'
            ]
            
            # Extract components from description
            for pattern in component_patterns:
                matches = re.findall(pattern, image_description, re.IGNORECASE)
                circuit_data['components'].extend(matches)
            
            # Remove duplicates and clean up
            circuit_data['components'] = list(set(circuit_data['components']))
            circuit_data['components'] = [comp for comp in circuit_data['components'] if len(comp) > 1]
            
            # Parse components section if available (fallback)
            if 'COMPONENTS:' in image_description and not circuit_data['components']:
                components_section = image_description.split('COMPONENTS:')[1].split('CONNECTIONS:')[0]
                for line in components_section.strip().split('\n'):
                    if line.strip().startswith('-'):
                        component_info = line.strip()[1:].strip()
                        circuit_data['components'].append(component_info)
            
            # Enhanced connection detection with comprehensive patterns
            connection_patterns = [
                # Power connections
                r'\bpositive\s+terminal\b',
                r'\bnegative\s+terminal\b',
                r'\bconnected\s+to\b',
                r'\bconnected\s+between\b',
                r'\bconnected\s+together\b',
                r'\bconnected\s+via\b',
                r'\bconnected\s+through\b',
                
                # Component terminals
                r'\banode\b',
                r'\bcathode\b',
                r'\bgate\b',
                r'\bcollector\b',
                r'\bemitter\b',
                r'\bbase\b',
                r'\bdrain\b',
                r'\bsource\b',
                r'\bterminal\b',
                r'\bpin\b',
                
                # Ground and power
                r'\bground\b',
                r'\bcommon\s+ground\b',
                r'\bearth\b',
                r'\bVCC\b',
                r'\bGND\b',
                r'\bVDD\b',
                r'\bVSS\b',
                r'\bpower\s+rail\b',
                r'\bvoltage\s+rail\b',
                
                # Switch connections
                r'\boutput\s+throw\b',
                r'\binput\s+pole\b',
                r'\bswitch\s+position\b',
                r'\bswitch\s+state\b',
                r'\barming\s+position\b',
                r'\bsafety\s+position\b',
                
                # Physical connections
                r'\bone\s+end\b',
                r'\bother\s+end\b',
                r'\bwire\b',
                r'\bline\b',
                r'\bconnection\b',
                r'\bjunction\b',
                r'\bnode\b',
                r'\bpoint\b',
                
                # Signal flow
                r'\bsignal\s+path\b',
                r'\bcurrent\s+flow\b',
                r'\bvoltage\s+path\b',
                r'\bcontrol\s+signal\b',
                r'\btrigger\s+signal\b',
                r'\boutput\s+signal\b',
                
                # Circuit topology
                r'\bseries\s+connection\b',
                r'\bparallel\s+connection\b',
                r'\bbranch\b',
                r'\bloop\b',
                r'\bcircuit\s+path\b',
                r'\breturn\s+path\b'
            ]
            
            # Extract connections from description
            for pattern in connection_patterns:
                matches = re.findall(pattern, image_description, re.IGNORECASE)
                circuit_data['connections'].extend(matches)
            
            # Remove duplicates
            circuit_data['connections'] = list(set(circuit_data['connections']))
            
            # SPECIFIC POWER RAIL AND POWER SUPPLY DETECTION
            power_rail_patterns = [
                # Standard power rails
                r'\bVCC\b', r'\bGND\b', r'\bVDD\b', r'\bVSS\b', r'\bVEE\b', r'\bVBB\b',
                r'\bpower\s+rail\b', r'\bvoltage\s+rail\b', r'\bpositive\s+rail\b', 
                r'\bnegative\s+rail\b', r'\bground\s+rail\b',
                r'\b12V\s+rail\b', r'\b5V\s+rail\b', r'\b3\.3V\s+rail\b', r'\b9V\s+rail\b',
                
                # Power supplies (count as power rails)
                r'\bpower\s+supply\b', r'\bvoltage\s+supply\b', r'\bcurrent\s+supply\b',
                r'\bBAT\d+\b', r'\bbattery\b', r'\b9V\b', r'\b12V\b', r'\b5V\b', r'\b3\.3V\b',
                r'\bvoltage\s+source\b', r'\bcurrent\s+source\b', r'\bSourceV\b', r'\bSourceI\b',
                
                # Power distribution
                r'\bpower\s+distribution\b', r'\bvoltage\s+distribution\b',
                r'\bpower\s+bus\b', r'\bvoltage\s+bus\b', r'\bpower\s+line\b', r'\bvoltage\s+line\b'
            ]
            
            for pattern in power_rail_patterns:
                matches = re.findall(pattern, image_description, re.IGNORECASE)
                circuit_data['power_rails'].extend(matches)
            
            # Remove duplicates from power rails
            circuit_data['power_rails'] = list(set(circuit_data['power_rails']))
            
            # Parse connections section if available (fallback)
            if 'CONNECTIONS:' in image_description and not circuit_data['connections']:
                connections_section = image_description.split('CONNECTIONS:')[1].split('CIRCUIT FUNCTION:')[0]
                for line in connections_section.strip().split('\n'):
                    if line.strip().startswith('-'):
                        connection_info = line.strip()[1:].strip()
                        circuit_data['connections'].append(connection_info)
            
            # Parse circuit function section
            if 'CIRCUIT FUNCTION:' in image_description:
                function_section = image_description.split('CIRCUIT FUNCTION:')[1]
                circuit_data['circuit_function'] = function_section.strip()
            
            # Determine complexity level
            component_count = len(circuit_data['components'])
            connection_count = len(circuit_data['connections'])
            
            if component_count > 15 or connection_count > 20:
                circuit_data['complexity_level'] = 'very_complex'
            elif component_count > 10 or connection_count > 15:
                circuit_data['complexity_level'] = 'complex'
            elif component_count > 5 or connection_count > 10:
                circuit_data['complexity_level'] = 'moderate'
            else:
                circuit_data['complexity_level'] = 'simple'
            
            print(f"📊 [CIRCUIT] Circuit complexity: {circuit_data['complexity_level']}")
            print(f"📋 [CIRCUIT] Components found: {component_count}")
            print(f"🔗 [CIRCUIT] Connections found: {connection_count}")
            print(f"⚡ [CIRCUIT] Power rails and supplies found: {len(circuit_data['power_rails'])}")
            if circuit_data['power_rails']:
                print(f"   - Power rails/supplies: {', '.join(circuit_data['power_rails'])}")
            
            return circuit_data
            
        except Exception as e:
            print(f"❌ [CIRCUIT] Error parsing complex circuit description: {str(e)}")
            return None

    def _generate_complex_circuit_prompt(self, circuit_data, unique_filename):
        """Generate a specialized prompt for complex circuit generation"""
        try:
            print("Generating specialized prompt for complex circuit...")
            
            complexity_level = circuit_data.get('complexity_level', 'simple')
            components = circuit_data.get('components', [])
            connections = circuit_data.get('connections', [])
            power_rails = circuit_data.get('power_rails', [])
            circuit_function = circuit_data.get('circuit_function', '')
            
            # Base prompt template
            prompt = f"""Generate a {complexity_level} circuit diagram using the python schemdraw library.

CIRCUIT ANALYSIS:
- Complexity Level: {complexity_level}
- Component Count: {len(components)}
- Connection Count: {len(connections)}
- Power Rails: {len(power_rails)} ({', '.join(power_rails) if power_rails else 'None detected'})
- Circuit Function: {circuit_function}

COMPONENTS TO INCLUDE:
"""
            
            # Add component details
            for i, component in enumerate(components[:10]):  # Limit to first 10 for prompt length
                prompt += f"- Component {i+1}: {component}\n"
            
            if len(components) > 10:
                prompt += f"- ... and {len(components) - 10} more components\n"
            
            prompt += f"""
POWER RAILS AND SUPPLIES TO IMPLEMENT:
"""
            
            # Add power rail details
            if power_rails:
                for i, rail in enumerate(power_rails):
                    prompt += f"- Power Rail/Supply {i+1}: {rail}\n"
            else:
                prompt += "- Power Rails/Supplies: Use standard VCC/GND rails and power supplies as needed\n"
            
            prompt += f"""
CONNECTIONS TO IMPLEMENT:
"""
            
            # Add connection details
            for i, connection in enumerate(connections[:10]):  # Limit to first 10 for prompt length
                prompt += f"- Connection {i+1}: {connection}\n"
            
            if len(connections) > 10:
                prompt += f"- ... and {len(connections) - 10} more connections\n"
            
            # Add complexity-specific instructions
            if complexity_level == 'very_complex':
                prompt += """
VERY COMPLEX CIRCUIT REQUIREMENTS:
- Use modular design with clear sections
- Implement multiple power rails (VCC, GND, VDD, etc.)
- Use elm.Dot for wire junctions and connection points
- Use elm.Label for power rails and voltage/current labels
- Organize components in logical blocks
- Use absolute positioning (.at()) for precise placement
- Minimize wire crossings and clutter
- Support feedback loops and control paths
- NEVER use d.element - this is INVALID and will cause errors
- ALWAYS use d.elements[-1] instead of d.element
- NEVER use d.last_end, d.last_start, d.end, d.start, d.position - these are INVALID attributes

SPECIALIZED COMPONENT HANDLING:
- DPDT switches: Use elm.Switch for double-pole double-throw switches
- SCR/Thyristor: Use elm.SCR for Silicon Controlled Rectifiers
- Multiple batteries: Use elm.Battery with proper labeling (BAT1, BAT2)
- Indicator LEDs: Use elm.LED with color specifications
- Initiator/Coil: Use elm.Inductor for coils and initiators
- Safety switches: Use elm.Switch with safety labels
- Power distribution: Use elm.Label for multiple voltage rails
- Ground connections: Use elm.Ground for common ground points

CIRCUIT ORGANIZATION:
- Input section: Safety switches and indicators (left side)
- Control section: Logic and power supplies (middle)
- Output section: Initiator and final controls (right side)
- Use elm.Text for section labels and component descriptions
"""
            elif complexity_level == 'complex':
                prompt += """
COMPLEX CIRCUIT REQUIREMENTS:
- Use clear signal flow from input to output
- Implement proper power distribution
- Use elm.Dot for connection points
- Group related components together
- Use consistent spacing and alignment
- Support multiple signal paths
"""
            else:
                prompt += """
STANDARD CIRCUIT REQUIREMENTS:
- Use logical component arrangement
- Implement proper connections
- Use clear labeling
- Maintain circuit clarity
"""
            
            # Add standard requirements
            prompt += f"""
STANDARD REQUIREMENTS:
- Use ONLY ASCII characters
- Use ONLY schemdraw.elements components
- Generate complete, executable Python script
- Use d.save() with filename: {unique_filename}
- Use proper positioning methods (.up(), .down(), .left(), .right(), .to())
- Label all components appropriately
- Handle all connections properly

CRITICAL CIRCUIT CLOSURE REQUIREMENTS:
- ALWAYS close the circuit loop using .to() method: d += elm.Line().to(d.elements[0].start)
- Ensure ALL components are connected in a complete loop
- Use explicit Line() elements to connect components when needed
- Start with a power source (elm.SourceV, elm.Battery)
- End with a connection back to the power source
- Use proper positioning to create logical circuit flow

Generate ONLY the Python code, no explanations."""
            
            return prompt
            
        except Exception as e:
            print(f"❌ [CIRCUIT] Error generating complex circuit prompt: {str(e)}")
            return None

    def _fix_component_naming_issues(self, code):
        """Fix common component naming issues in generated code"""
        try:
            print("🔧 [CIRCUIT] Fixing component naming issues...")
            
            # Fix IC -> Ic issue
            fixed_code = code.replace('elm.IC', 'elm.Ic')
            fixed_code = fixed_code.replace('elm.IC(', 'elm.Ic(')
            
            # Fix other common naming issues
            fixed_code = fixed_code.replace('elm.IC)', 'elm.Ic)')
            
            # Check if any fixes were made
            if fixed_code != code:
                print("✅ [CIRCUIT] Fixed component naming issues")
            else:
                print("✅ [CIRCUIT] No component naming issues found")
            
            return fixed_code
            
        except Exception as e:
            print(f"❌ [CIRCUIT] Error fixing component naming issues: {str(e)}")
            return code

    def _execute_generated_circuit_code(self, generated_code):
        """Execute the generated circuit code and return the diagram file"""
        temp_script = None
        try:
            # Clean up previous circuit files first
            self._cleanup_previous_circuit_files()
            
            # Extract the expected filename from the generated code
            expected_filename = None
            import re
            save_match = re.search(r"d\.save\(['\"]([^'\"]+)['\"]\)", generated_code)
            if save_match:
                expected_filename = save_match.group(1)
                print(f"🎯 [CIRCUIT] Expected filename from code: {expected_filename}")
            
            print("🔧 [CIRCUIT] Normalizing Unicode characters in generated code...")
            # Normalize Unicode characters to ASCII equivalents for better compatibility
            import unicodedata
            normalized_code = unicodedata.normalize('NFD', generated_code)
            # Replace common Unicode characters with ASCII equivalents
            normalized_code = normalized_code.replace('Ω', 'Ohm')
            normalized_code = normalized_code.replace('μ', 'u')
            normalized_code = normalized_code.replace('°', 'deg')
            normalized_code = normalized_code.replace('±', '+/-')
            normalized_code = normalized_code.replace('≤', '<=')
            normalized_code = normalized_code.replace('≥', '>=')
            normalized_code = normalized_code.replace('≠', '!=')
            print("✅ [CIRCUIT] Unicode normalization completed")
            
            print("📄 [CIRCUIT] Creating temporary Python script...")
            # Create a temporary file for the generated code with UTF-8 encoding
            with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False, encoding='utf-8') as f:
                f.write(normalized_code)
                temp_script = f.name
            print(f"📁 [CIRCUIT] Temporary script created: {temp_script}")
            
            print("⚙️ [CIRCUIT] Setting up execution environment...")
            # Execute the generated script with UTF-8 environment
            env = os.environ.copy()
            env['PYTHONIOENCODING'] = 'utf-8'
            
            print("🚀 [CIRCUIT] Executing generated Python script...")
            result = subprocess.run(['python', temp_script], 
                                  capture_output=True, text=True, timeout=60,
                                  env=env, encoding='utf-8')
            
            if result.returncode == 0:
                print("✅ [CIRCUIT] Script executed successfully")
                print("🔍 [CIRCUIT] Searching for generated PNG files...")
                
                # First, look for the expected filename
                if expected_filename and os.path.exists(expected_filename):
                    print(f"✅ [CIRCUIT] Found expected file: {expected_filename}")
                    return expected_filename
                
                # Look for generated PNG files
                generated_files = []
                for file in os.listdir('.'):
                    if file.endswith('.png'):
                        generated_files.append(file)
                
                if generated_files:
                    # Prefer files with 'circuit' in the name
                    circuit_files = [f for f in generated_files if 'circuit' in f.lower()]
                    if circuit_files:
                        selected_file = circuit_files[0]
                        print(f"✅ [CIRCUIT] Found generated circuit diagram: {selected_file}")
                        return selected_file
                    else:
                        # Use the first PNG file found
                        selected_file = generated_files[0]
                        print(f"✅ [CIRCUIT] Found generated diagram: {selected_file}")
                        return selected_file
                else:
                    print("❌ [CIRCUIT] No PNG files found after successful execution")
                    return "Error: No PNG files generated despite successful script execution"
            else:
                print(f"❌ [CIRCUIT] Script execution failed with return code: {result.returncode}")
                print(f"📝 [CIRCUIT] Error output: {result.stderr}")
                print(f"📝 [CIRCUIT] Standard output: {result.stdout}")
                
                # Provide more specific error messages
                error_msg = result.stderr.strip()
                if "ModuleNotFoundError" in error_msg:
                    return f"Error: Missing required module - {error_msg}"
                elif "AttributeError: module 'schemdraw.elements' has no attribute 'IC'. Did you mean: 'Ic'?" in error_msg:
                    return f"Error: Use 'elm.Ic' instead of 'elm.IC' for integrated circuits - {error_msg}"
                elif "AttributeError" in error_msg:
                    return f"Error: Invalid component or method used - {error_msg}"
                elif "SyntaxError" in error_msg:
                    return f"Error: Syntax error in generated code - {error_msg}"
                elif "ImportError" in error_msg:
                    return f"Error: Import error - {error_msg}"
                elif "d.draw()" in error_msg:
                    # Handle d.draw() gracefully - it's now allowed but may not work as expected
                    return f"Warning: d.draw() was used but may not generate a file. Consider using d.save() for better results."
                elif "Duplicate `at` parameter in element" in error_msg:
                    return f"Warning: Duplicate positioning parameters detected - {error_msg}"
                else:
                    return f"Error: Script execution failed - {error_msg}"
                
        except subprocess.TimeoutExpired:
            print("❌ [CIRCUIT] Script execution timed out")
            return "Error: Script execution timed out (60 seconds)"
        except Exception as e:
            print(f"❌ [CIRCUIT] Exception during code execution: {str(e)}")
            return f"Error: Exception during code execution - {str(e)}"
        finally:
            # Clean up the temporary script
            if temp_script and os.path.exists(temp_script):
                try:
                    os.unlink(temp_script)
                    print("🧹 [CIRCUIT] Temporary script cleaned up")
                except Exception as e:
                    print(f"⚠️ [CIRCUIT] Failed to clean up temporary script: {str(e)}")

    def _validate_circuit_code(self, code):
        """Validate the generated circuit code for common issues"""
        try:
            print("🔍 [CIRCUIT] Validating generated code...")
            
            # Check for required imports
            if 'import schemdraw' not in code:
                print("❌ [CIRCUIT] Missing schemdraw import")
                return False
            
            # Check for forbidden components
            forbidden_components = [
                'elm.Tip', 'elm.DCSourceV', 'elm.SpiceNetlist', 'elm.SpiceNetlistElement',
                'matplotlib', 'pyplot', 'plt', 'import matplotlib', 'from matplotlib'
            ]
            
            for component in forbidden_components:
                if component in code:
                    print(f"❌ [CIRCUIT] Forbidden component found: {component}")
                    return False
            
            # Check for invalid assignment syntax (e.g., light_bulb = d += elm.Lamp())
            import re
            invalid_assignment_patterns = [
                r'\w+\s*=\s*d\s*\+=',  # variable = d +=
                r'\w+\s*=\s*d\.add\(',  # variable = d.add(
                r'\w+\s*=\s*d\.append\(',  # variable = d.append(
            ]
            for pattern in invalid_assignment_patterns:
                if re.search(pattern, code):
                    print(f"❌ [CIRCUIT] Invalid assignment syntax detected: {pattern}")
                    return False
            
            # Check for grounding elements (not allowed for closed loop circuits)
            grounding_elements = ['elm.Ground', 'elm.GroundChassis', 'elm.GroundSignal', 'elm.Ground']
            for ground_element in grounding_elements:
                if ground_element in code:
                    print(f"❌ [CIRCUIT] Grounding element found: {ground_element} - closed loop circuits should not have grounding elements")
                    return False
            
            # Check for closed loop circuit structure
            if not self._validate_closed_loop_circuit(code):
                print("❌ [CIRCUIT] Circuit is not a complete closed loop")
                return False
            
            # Check for forbidden methods (but ignore d.draw() as it's now allowed)
            # Note: d.draw() is now allowed to pass validation
            if 'd.draw()' in code:
                print("⚠️ [CIRCUIT] d.draw() found - allowing to pass validation")
                # Don't fail validation for d.draw() anymore
            
            # Check for Unicode characters
            unicode_chars = ['Ω', 'μ', '°', '±', '≤', '≥', '≠', '∞', '∑', '∏', '∫', '∂']
            for char in unicode_chars:
                if char in code:
                    print(f"❌ [CIRCUIT] Unicode character found: {char}")
                    return False
            
            # Check for proper save method
            if 'd.save(' not in code:
                print("❌ [CIRCUIT] Missing d.save() method")
                return False
            
            # Check for basic structure
            if 'schemdraw.Drawing()' not in code:
                print("❌ [CIRCUIT] Missing schemdraw.Drawing() initialization")
                return False
            
            # Check if the circuit is just a copy of the example
            example_components = ['100KOhm', '0.1uF', '10V']
            example_count = sum(1 for component in example_components if component in code)
            if example_count >= 2:  # If 2 or more example values are used
                print("⚠️ [CIRCUIT] Circuit appears to be copying example values too closely")
                # Don't fail validation, but warn about potential copying
            
            # Check for minimum circuit complexity (should have at least 3 components)
            component_patterns = [
                'elm.Resistor', 'elm.Capacitor', 'elm.Inductor', 'elm.Diode', 
                'elm.SourceV', 'elm.SourceI', 'elm.Ground', 'elm.Line', 'elm.Dot',
                'elm.Rect', 'elm.RBox', 'elm.Circle', 'elm.Transistor', 'elm.OpAmp',
                'elm.Switch', 'elm.LED', 'elm.Motor', 'elm.Relay', 'elm.Crystal',
                'elm.Transformer', 'elm.Potentiometer', 'elm.Thermistor', 'elm.Varistor',
                'elm.Fuse', 'elm.Connector', 'elm.Ic', 'elm.Battery', 'elm.CurrentLabel',
                'elm.VoltageLabel', 'elm.Node', 'elm.Dot2', 'elm.Contact', 'elm.Arrow',
                'elm.Text', 'elm.Lamp'
            ]
            component_count = sum(1 for pattern in component_patterns if pattern in code)
            if component_count < 3:
                print("⚠️ [CIRCUIT] Circuit appears too simple - may be copying example")
                # Don't fail validation, but warn about potential copying
            
            # Check for component labeling (should have labels for most components)
            label_count = code.count('.label(')
            if component_count > 0 and label_count < component_count * 0.5:  # Less than 50% labeled
                print("⚠️ [CIRCUIT] Many components are not labeled - consider adding labels")
                # Don't fail validation, but warn about missing labels
            
            print("✅ [CIRCUIT] Code validation passed")
            return True
            
        except Exception as e:
            print(f"❌ [CIRCUIT] Error during code validation: {str(e)}")
            return False

    def _validate_closed_loop_circuit(self, code):
        """Validate that the circuit forms a complete closed loop without grounding elements"""
        try:
            print("🔍 [CIRCUIT] Validating closed loop circuit structure...")
            
            # Extract component lines
            lines = code.split('\n')
            component_lines = []
            
            for line in lines:
                line = line.strip()
                if line.startswith('d += elm.') and not line.startswith('d += elm.Ground'):
                    component_lines.append(line)
            
            if len(component_lines) < 3:
                print("❌ [CIRCUIT] Circuit must have at least 3 components for a closed loop")
                return False
            
            # Check for power source
            power_sources = ['elm.SourceV', 'elm.SourceI', 'elm.Battery', 'elm.SourceSin', 'elm.SourceSquare']
            has_power = any(source in code for source in power_sources)
            if not has_power:
                print("❌ [CIRCUIT] Closed loop circuit must have a power source")
                return False
            
            # Check for proper connection methods (up, down, left, right, to)
            connection_methods = ['.up()', '.down()', '.left()', '.right()', '.to(']
            has_connections = any(method in code for method in connection_methods)
            if not has_connections:
                print("❌ [CIRCUIT] Circuit components must be properly connected using directional methods")
                return False
            
            # Check for loop completion (should have a .to() method or complete path)
            if '.to(' not in code:
                # Check if the last component connects back to form a loop
                # This is a simplified check - in practice, the LLM should use .to() method
                print("⚠️ [CIRCUIT] Consider using .to() method to explicitly close the circuit loop")
            
            print("✅ [CIRCUIT] Closed loop circuit validation passed")
            return True
            
        except Exception as e:
            print(f"❌ [CIRCUIT] Error validating closed loop circuit: {str(e)}")
            return False

    def _extract_python_code(self, response_text):
        """Extract Python code from AI model response, handling markdown code blocks"""
        try:
            print("🔍 [CIRCUIT] Analyzing response for code blocks...")
            
            # Check if response contains markdown code blocks
            if '```python' in response_text:
                print("📦 [CIRCUIT] Found Python code block, extracting...")
                # Extract code between ```python and ```
                start_marker = '```python'
                end_marker = '```'
                
                start_idx = response_text.find(start_marker)
                if start_idx != -1:
                    # Find the end of the code block
                    code_start = start_idx + len(start_marker)
                    end_idx = response_text.find(end_marker, code_start)
                    
                    if end_idx != -1:
                        extracted_code = response_text[code_start:end_idx].strip()
                        print("✅ [CIRCUIT] Successfully extracted Python code from markdown block")
                        return extracted_code
                    else:
                        print("⚠️ [CIRCUIT] Found start marker but no end marker, using rest of text")
                        return response_text[code_start:].strip()
                else:
                    print("⚠️ [CIRCUIT] No start marker found")
                    return response_text
            
            # Check for other code block formats
            elif '```' in response_text:
                print("📦 [CIRCUIT] Found generic code block, extracting...")
                # Extract code between ``` and ```
                start_marker = '```'
                end_marker = '```'
                
                start_idx = response_text.find(start_marker)
                if start_idx != -1:
                    code_start = start_idx + len(start_marker)
                    end_idx = response_text.find(end_marker, code_start)
                    
                    if end_idx != -1:
                        extracted_code = response_text[code_start:end_idx].strip()
                        # Remove language identifier if present
                        if extracted_code.startswith('python'):
                            extracted_code = extracted_code[6:].strip()
                        print("✅ [CIRCUIT] Successfully extracted code from generic block")
                        return extracted_code
                    else:
                        print("⚠️ [CIRCUIT] Found start marker but no end marker, using rest of text")
                        return response_text[code_start:].strip()
                else:
                    print("⚠️ [CIRCUIT] No start marker found")
                    return response_text
            
            else:
                print("📝 [CIRCUIT] No code blocks found, using response as-is")
                return response_text
                
        except Exception as e:
            print(f"❌ [CIRCUIT] Error extracting Python code: {str(e)}")
            return response_text

    def process_circuit_image(self, image):
        """Main function to process uploaded circuit image"""
        try:
            print("=" * 60)
            print("🚀 [CIRCUIT] Starting circuit diagram generation process")
            print("=" * 60)
            
            if image is None:
                print("❌ [CIRCUIT] No image uploaded")
                return "No image uploaded", None
            
            print("📸 [CIRCUIT] Image uploaded successfully")
            
            # Step 1: Describe image with Gemma3
            print("\n" + "=" * 40)
            print("🔍 STEP 1: Image Description with Gemma3")
            print("=" * 40)
            description = self.describe_image_with_gemma3(image)
            
            # Step 2: Generate circuit with DeepSeek R1
            print("\n" + "=" * 40)
            print("🔧 STEP 2: Circuit Generation with DeepSeek R1")
            print("=" * 40)
            circuit_result = self.generate_circuit_with_deepseek(description)
            
            # Step 3: Return results
            print("\n" + "=" * 40)
            print("📊 STEP 3: Finalizing Results")
            print("=" * 40)
            
            if circuit_result and (circuit_result.endswith('.png') or 'circuit_diagram_' in circuit_result):
                print(f"✅ [CIRCUIT] Circuit diagram generated successfully: {circuit_result}")
                print("=" * 60)
                print("🎉 [CIRCUIT] Process completed successfully!")
                print("=" * 60)
                
                # Check if there's a note about missing components
                if "(Note:" in circuit_result:
                    # Extract the actual filename and the note
                    filename = circuit_result.split(' (Note:')[0]
                    note = circuit_result.split('(Note:')[1].rstrip(')')
                    return f"Image Description: {description}\n\nCircuit Generated: {filename}\n\n{note}", filename
                else:
                    return f"Image Description: {description}\n\nCircuit Generated: {circuit_result}", circuit_result
            else:
                print(f"⚠️ [CIRCUIT] Circuit generation failed: {circuit_result}")
                print("=" * 60)
                print("❌ [CIRCUIT] Process completed with errors")
                print("=" * 60)
                
                # Provide more detailed error information
                error_details = ""
                if "Error:" in circuit_result:
                    error_details = f"\n\nError Details:\n{circuit_result}"
                
                return f"Image Description: {description}\n\nCircuit Generation Failed{error_details}", None
                
        except Exception as e:
            error_msg = f"Error processing circuit image: {str(e)}"
            print(f"❌ [CIRCUIT] {error_msg}")
            print("=" * 60)
            print("💥 [CIRCUIT] Process failed!")
            print("=" * 60)
            return error_msg, None

    def _enhance_circuit_connections(self, code):
        """Enhance circuit connections to ensure proper closure and connectivity"""
        try:
            print("🔧 [CIRCUIT] Enhancing circuit connections for proper closure...")
            
            lines = code.split('\n')
            component_lines = []
            connection_lines = []
            
            # Separate component lines from connection lines
            for i, line in enumerate(lines):
                line = line.strip()
                if line.startswith('d += elm.') and not line.startswith('d += elm.Ground'):
                    component_lines.append((i, line))
                elif line.startswith('d += elm.Line') or line.startswith('d += elm.Dot'):
                    connection_lines.append((i, line))
            
            if len(component_lines) < 2:
                print("⚠️ [CIRCUIT] Not enough components to enhance connections")
                return code
            
            # Check if circuit already has proper closure
            has_closure = any('.to(' in line for _, line in component_lines + connection_lines)
            
            if not has_closure:
                print("🔗 [CIRCUIT] Adding circuit closure connection...")
                
                # Find the last component line
                last_component_idx, last_component_line = component_lines[-1]
                
                # Create closure connection
                closure_line = f"d += elm.Line().to(d.elements[0].start)"
                
                # Insert closure line after the last component
                lines.insert(last_component_idx + 1, closure_line)
                
                print("✅ [CIRCUIT] Added circuit closure connection")
            
            # Check for disconnected components and add connections
            enhanced_code = self._add_missing_connections(lines)
            
            return enhanced_code
            
        except Exception as e:
            print(f"❌ [CIRCUIT] Error enhancing circuit connections: {str(e)}")
            return code

    def _add_missing_connections(self, lines):
        """Add missing connections between components"""
        try:
            print("🔗 [CIRCUIT] Adding missing connections between components...")
            
            # Find all component lines
            component_indices = []
            for i, line in enumerate(lines):
                if line.strip().startswith('d += elm.') and not line.strip().startswith('d += elm.Ground'):
                    component_indices.append(i)
            
            if len(component_indices) < 2:
                return '\n'.join(lines)
            
            # Check for gaps in connections
            enhanced_lines = lines.copy()
            insertions = 0
            
            for i in range(len(component_indices) - 1):
                current_idx = component_indices[i] + insertions
                next_idx = component_indices[i + 1] + insertions
                
                # Check if there's a connection between these components
                has_connection = False
                for j in range(current_idx + 1, next_idx):
                    if j < len(enhanced_lines) and enhanced_lines[j].strip().startswith('d += elm.Line'):
                        has_connection = True
                        break
                
                if not has_connection:
                    # Add a connection line
                    connection_line = "d += elm.Line().right()"
                    enhanced_lines.insert(next_idx, connection_line)
                    insertions += 1
                    print(f"🔗 [CIRCUIT] Added connection between components {i+1} and {i+2}")
            
            return '\n'.join(enhanced_lines)
            
        except Exception as e:
            print(f"❌ [CIRCUIT] Error adding missing connections: {str(e)}")
            return '\n'.join(lines)

    def _validate_circuit_connectivity(self, code):
        """Validate that all components are properly connected"""
        try:
            print("🔍 [CIRCUIT] Validating circuit connectivity...")
            
            lines = code.split('\n')
            component_count = 0
            connection_count = 0
            
            for line in lines:
                line = line.strip()
                if line.startswith('d += elm.') and not line.startswith('d += elm.Ground'):
                    component_count += 1
                elif line.startswith('d += elm.Line') or line.startswith('d += elm.Dot'):
                    connection_count += 1
            
            # Basic connectivity check
            if component_count < 2:
                print("❌ [CIRCUIT] Circuit needs at least 2 components")
                return False
            
            if connection_count < 1:
                print("❌ [CIRCUIT] Circuit needs at least 1 connection")
                return False
            
            # Check for proper closure
            has_closure = '.to(' in code
            if not has_closure:
                print("⚠️ [CIRCUIT] Circuit may not be properly closed")
            
            print(f"✅ [CIRCUIT] Circuit connectivity validation passed - {component_count} components, {connection_count} connections")
            return True
            
        except Exception as e:
            print(f"❌ [CIRCUIT] Error validating circuit connectivity: {str(e)}")
            return False

    def _fix_circuit_structure(self, code):
        """Fix common circuit structure issues"""
        try:
            print("🔧 [CIRCUIT] Fixing circuit structure issues...")
            
            lines = code.split('\n')
            fixed_lines = []
            
            for line in lines:
                line = line.strip()
                
                # Fix common positioning issues
                if 'd += elm.' in line:
                    # Ensure proper positioning methods are used
                    if not any(method in line for method in ['.up()', '.down()', '.left()', '.right()', '.to(', '.at(']):
                        # Add basic positioning if missing
                        if 'elm.SourceV' in line or 'elm.Battery' in line:
                            line = line.rstrip() + '.up()'
                        elif 'elm.Resistor' in line or 'elm.Capacitor' in line:
                            line = line.rstrip() + '.right()'
                        elif 'elm.LED' in line or 'elm.Diode' in line:
                            line = line.rstrip() + '.down()'
                
                # Fix component naming issues
                line = line.replace('elm.IC', 'elm.Ic')
                line = line.replace('elm.IC(', 'elm.Ic(')
                
                fixed_lines.append(line)
            
            # Ensure proper circuit closure
            fixed_code = '\n'.join(fixed_lines)
            enhanced_code = self._enhance_circuit_connections(fixed_code)
            
            print("✅ [CIRCUIT] Circuit structure fixes applied")
            return enhanced_code
            
        except Exception as e:
            print(f"❌ [CIRCUIT] Error fixing circuit structure: {str(e)}")
            return code

    def _generate_robust_circuit_template(self, components, unique_filename):
        """Generate a robust circuit template with proper connections"""
        try:
            print("🔧 [CIRCUIT] Generating robust circuit template...")
            
            template = f"""import schemdraw
import schemdraw.elements as elm

d = schemdraw.Drawing()

# Power source
d += elm.SourceV().up().label('12V').at((0, 0))

# Main circuit components
"""
            
            # Add components with proper positioning
            for i, component in enumerate(components[:5]):  # Limit to 5 components for template
                if 'resistor' in component.lower():
                    template += f"d += elm.Resistor().right().label('R{i+1}')\n"
                elif 'capacitor' in component.lower():
                    template += f"d += elm.Capacitor().down().label('C{i+1}')\n"
                elif 'led' in component.lower():
                    template += f"d += elm.LED().right().label('LED{i+1}')\n"
                elif 'switch' in component.lower():
                    template += f"d += elm.Switch().up().label('SW{i+1}')\n"
                elif 'battery' in component.lower() or 'power' in component.lower():
                    template += f"d += elm.Battery().up().label('BAT{i+1}')\n"
                else:
                    template += f"d += elm.RBox().right().label('{component}')\n"
            
            # Add proper closure
            template += f"""
# Close the circuit loop
d += elm.Line().left().to(d.elements[0].start)

# Save the diagram
d.save('{unique_filename}')
"""
            
            print("✅ [CIRCUIT] Robust circuit template generated")
            return template
            
        except Exception as e:
            print(f"❌ [CIRCUIT] Error generating robust circuit template: {str(e)}")
            return None

    def _create_validated_circuit_template(self, image_description, unique_filename):
        """Create a validated circuit template based on image description"""
        try:
            print("🔧 [CIRCUIT] Creating validated circuit template...")
            
            # Extract components from description
            components = self._extract_components_from_description(image_description)
            
            if not components:
                print("⚠️ [CIRCUIT] No specific components found, using generic template")
                return self._generate_generic_validated_template(unique_filename)
            
            # Create template with extracted components
            template = f"""import schemdraw
import schemdraw.elements as elm

d = schemdraw.Drawing()

# Power source - always start with power
d += elm.SourceV().up().label('12V').at((0, 0))

# Circuit components based on image description
"""
            
            # Add components with proper validation
            component_count = 0
            for component in components[:6]:  # Limit to 6 components for template
                component_count += 1
                component_type = component.get('type', 'RBox')
                value = component.get('value', str(component_count))
                
                if component_type.lower() == 'resistor':
                    template += f"d += elm.Resistor().right().label('R{component_count}')\n"
                elif component_type.lower() == 'capacitor':
                    template += f"d += elm.Capacitor().down().label('C{component_count}')\n"
                elif component_type.lower() == 'led':
                    template += f"d += elm.LED().right().label('LED{component_count}')\n"
                elif component_type.lower() == 'diode':
                    template += f"d += elm.Diode().right().label('D{component_count}')\n"
                elif component_type.lower() == 'switch':
                    template += f"d += elm.Switch().up().label('SW{component_count}')\n"
                elif component_type.lower() == 'transistor':
                    template += f"d += elm.Transistor().up().label('Q{component_count}')\n"
                elif component_type.lower() == 'battery':
                    template += f"d += elm.Battery().up().label('BAT{component_count}')\n"
                elif component_type.lower() == 'sourcev':
                    template += f"d += elm.SourceV().up().label('V{component_count}')\n"
                elif component_type.lower() == 'ic':
                    template += f"d += elm.Ic().right().label('IC{component_count}')\n"
                else:
                    template += f"d += elm.RBox().right().label('{component_type}{component_count}')\n"
            
            # Add proper closure and validation
            template += f"""
# Ensure circuit closure - critical for proper operation
d += elm.Line().left().to(d.elements[0].start)

# Save the validated circuit diagram
d.save('{unique_filename}')
"""
            
            print(f"✅ [CIRCUIT] Validated circuit template created with {component_count} components")
            return template
            
        except Exception as e:
            print(f"❌ [CIRCUIT] Error creating validated circuit template: {str(e)}")
            return self._generate_generic_validated_template(unique_filename)

    def _generate_generic_validated_template(self, unique_filename):
        """Generate a generic but validated circuit template"""
        try:
            print("🔧 [CIRCUIT] Generating generic validated template...")
            
            template = f"""import schemdraw
import schemdraw.elements as elm

d = schemdraw.Drawing()

# Power source - essential for circuit operation
d += elm.SourceV().up().label('12V').at((0, 0))

# Basic circuit components with proper connections
d += elm.Resistor().right().label('R1')
d += elm.LED().down().label('LED1')
d += elm.Capacitor().left().label('C1')

# Critical: Close the circuit loop for proper current flow
d += elm.Line().up().to(d.elements[0].start)

# Save the validated circuit
d.save('{unique_filename}')
"""
            
            print("✅ [CIRCUIT] Generic validated template generated")
            return template
            
        except Exception as e:
            print(f"❌ [CIRCUIT] Error generating generic template: {str(e)}")
            return None

    def _extract_components_from_description(self, image_description):
        """Extract component information from the image description"""
        try:
            components = []
            
            # Enhanced component patterns based on circuit validation best practices
            component_patterns = [
                (r'resistor[s]?\s+(\w+)', 'Resistor'),
                (r'capacitor[s]?\s+(\w+)', 'Capacitor'),
                (r'led[s]?\s+(\w+)', 'LED'),
                (r'diode[s]?\s+(\w+)', 'Diode'),
                (r'switch[s]?\s+(\w+)', 'Switch'),
                (r'transistor[s]?\s+(\w+)', 'Transistor'),
                (r'bjt[s]?\s+(\w+)', 'Transistor'),
                (r'battery[s]?\s+(\w+)', 'Battery'),
                (r'voltage\s+source[s]?\s+(\w+)', 'SourceV'),
                (r'power\s+supply[s]?\s+(\w+)', 'SourceV'),
                (r'ic[s]?\s+(\w+)', 'Ic'),
                (r'integrated\s+circuit[s]?\s+(\w+)', 'Ic'),
                (r'inductor[s]?\s+(\w+)', 'Inductor'),
                (r'relay[s]?\s+(\w+)', 'Relay'),
                (r'motor[s]?\s+(\w+)', 'Motor'),
                (r'fuse[s]?\s+(\w+)', 'Fuse'),
                (r'connector[s]?\s+(\w+)', 'Connector'),
            ]
            
            import re
            for pattern, component_type in component_patterns:
                matches = re.findall(pattern, image_description.lower())
                for match in matches:
                    components.append({
                        'type': component_type,
                        'value': match,
                        'description': f"{component_type} {match}"
                    })
            
            # Remove duplicates while preserving order
            seen = set()
            unique_components = []
            for component in components:
                key = f"{component['type']}_{component['value']}"
                if key not in seen:
                    seen.add(key)
                    unique_components.append(component)
            
            return unique_components
            
        except Exception as e:
            print(f"❌ [CIRCUIT] Error extracting components from description: {str(e)}")
            return []



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 - Production Ready")
        gr.Markdown("Made by Collar - Enhanced with Team Management & Chat History")
        
        # 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
                        )
        
        # Chat History Tab
        with gr.Tab("💬 Chat History"):
            with gr.Column():
                gr.Markdown("### 📚 Conversation History")
                gr.Markdown("View and manage your previous conversations with the AI assistant.")
                
                with gr.Row():
                    with gr.Column(scale=2):
                        history_limit = gr.Slider(
                            minimum=5,
                            maximum=50,
                            value=10,
                            step=5,
                            label="Number of recent conversations to display"
                        )
                    with gr.Column(scale=1):
                        refresh_history_btn = gr.Button("🔄 Refresh History", variant="secondary")
                        clear_history_btn = gr.Button("🗑️ Clear History", variant="stop")
                
                chat_history_display = gr.Markdown(
                    label="Recent Conversations", 
                    value="💬 **Welcome to Chat History!**\n\nLog in and start a conversation to see your chat history here."
                )
        
        # Data Management Tab
        with gr.Tab("⚙️ Data Management"):
            with gr.Column():
                gr.Markdown("### Collection Management")
                choice = gr.Dropdown(
                    choices=app.display_file_list(),
                    label="Select Collection to Delete"
                )
                delete_button = gr.Button("Delete Collection", variant="stop")
                delete_status = gr.Textbox(label="Deletion Status", interactive=False)
                




        # Circuit Diagram Generation Tab
        with gr.Tab("⚡ Circuit Diagram Generator"):
            with gr.Column():
                gr.Markdown("### Circuit Diagram Generation")
                gr.Markdown("Upload a circuit image to generate a netlist and circuit diagram using AI models.")
                
                circuit_image_input = gr.Image(
                    type="pil", 
                    label="Upload Circuit Image",
                    height=300
                )
                generate_circuit_btn = gr.Button("Generate Circuit Diagram", variant="primary")
                
                gr.Markdown("### Results")
                circuit_output = gr.Textbox(
                    label="Processing Results", 
                    interactive=False,
                    lines=8
                )
                circuit_diagram_output = gr.Image(
                    label="Generated Circuit Diagram",
                    height=400
                )
        
        # 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]
        )
        

        
        # Chat history events
        refresh_history_btn.click(
            fn=app.get_chat_history,
            inputs=[session_state, history_limit],
            outputs=[chat_history_display]
        )
        
        clear_history_btn.click(
            fn=app.clear_chat_history,
            inputs=[session_state],
            outputs=[chat_history_display]
        )
        
        # Data management events
        delete_button.click(
            fn=app.delete,
            inputs=[session_state, choice, session_state],
            outputs=[delete_status]
        )


        
        # Circuit generation events
        generate_circuit_btn.click(
            fn=app.process_circuit_image,
            inputs=[circuit_image_input],
            outputs=[circuit_output, circuit_diagram_output]
        )
        
    return demo

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