"""
Document Base Module - Backward Compatibility Layer

This module has been refactored into modular components:
- app/models/document.py (SQLAlchemy models)
- app/services/document_db_service.py (ORM database operations)  
- app/services/document_service.py (main orchestrator)
- app/utils/content_extractor.py
- app/utils/image_extractor.py
- app/utils/preview_generator.py
- app/utils/chroma_vector_db.py
- app/utils/ai_analyzer.py
- app/utils/document_ui_helper.py

The original classes are kept here for backward compatibility.
For new code, prefer importing from the new module locations.
"""

import os
import hashlib
import sqlite3
from datetime import datetime
from pathlib import Path
from typing import Any, List, Dict, Tuple, Optional
import pymupdf
import PyPDF2
import pandas as pd
from PIL import Image, ImageDraw, ImageFont
import pytesseract
import docx
from openai import OpenAI
import chromadb

from app.services.llm_service import LLMService
from chromadb.config import Settings
from chromadb.utils import embedding_functions
import re
import json
import io

# Note: New modular components are available at:
# - app.models.document (DocumentFile, DocumentFolder, ExtractedImage)
# - app.services.document_db_service (DocumentDbService)
# - app.services.document_service (DocumentService)
# - app.utils.content_extractor (ContentExtractor)
# - app.utils.image_extractor (ImageExtractor)
# - app.utils.preview_generator (PreviewGenerator)
# - app.utils.chroma_vector_db (ChromaVectorDB)
# - app.utils.ai_analyzer (AIAnalyzer)
# - app.utils.document_ui_helper (DocumentUIHelper)
#
# Import them directly from those modules for new code.
# The original classes below are kept for backward compatibility.


def _use_app_document_db():
    """True when running inside app and pooled/WAL document DB is available."""
    try:
        from app.core.database import get_document_engine, init_document_database  # noqa: F401
        from app.services.document_db_service import DocumentDbService  # noqa: F401
        return True
    except ImportError:
        return False


class DatabaseManager:
    """Manages file_sys_metadata.db. When run inside the app, uses pooled WAL engine from app.core.database."""

    def __init__(self, db_file: str = "file_sys_metadata.db"):
        self.db_file = db_file
        if _use_app_document_db():
            from app.services.document_db_service import DocumentDbService
            self._service = DocumentDbService(db_file=db_file)
            self._use_app = True
        else:
            self._use_app = False
            self.init_database()

    def init_database(self):
        if getattr(self, '_use_app', False):
            from app.core.database import init_document_database
            init_document_database(self.db_file)
            return
        conn = sqlite3.connect(self.db_file)
        cursor = conn.cursor()
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS files (
                id TEXT PRIMARY KEY,
                filename TEXT NOT NULL,
                original_filename TEXT NOT NULL,
                file_type TEXT,
                file_path TEXT NOT NULL,
                folder_path TEXT DEFAULT '',
                relative_path TEXT DEFAULT '',
                upload_date TEXT,
                file_size INTEGER,
                content_preview TEXT,
                full_content TEXT,
                image_path TEXT,
                thumbnail_path TEXT
            )
        ''')
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS folders (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                folder_name TEXT NOT NULL,
                folder_path TEXT NOT NULL,
                parent_path TEXT,
                full_path TEXT,
                file_count INTEGER DEFAULT 0,
                created_date TEXT,
                agent_id TEXT
            )
        ''')
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS extracted_images (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                file_id TEXT NOT NULL,
                image_path TEXT NOT NULL,
                image_type TEXT,
                page_number INTEGER,
                upload_date TEXT,
                agent_id TEXT,
                FOREIGN KEY (file_id) REFERENCES files(id) ON DELETE CASCADE
            )
        ''')
        conn.commit()
        conn.close()
        try:
            from app.core.database import ensure_document_db_indexes
            ensure_document_db_indexes(self.db_file)
        except ImportError:
            self._ensure_document_indexes_standalone()

    def _ensure_document_indexes_standalone(self) -> None:
        """Create file metadata indexes when running outside the FastAPI app stack."""
        specs = (
            ("ix_files_folder_path", "files", "folder_path", ("folder_path",)),
            ("ix_files_upload_date", "files", "upload_date DESC", ("upload_date",)),
            ("ix_folders_full_path", "folders", "full_path", ("full_path",)),
            ("ix_extracted_images_file_id", "extracted_images", "file_id", ("file_id",)),
            (
                "ix_extracted_images_file_page",
                "extracted_images",
                "file_id, page_number, id",
                ("file_id", "page_number", "id"),
            ),
        )
        conn = sqlite3.connect(self.db_file)
        try:
            cursor = conn.cursor()
            col_cache: dict = {}

            def table_columns(table: str) -> set:
                if table not in col_cache:
                    cursor.execute(f"PRAGMA table_info({table})")
                    col_cache[table] = {row[1] for row in cursor.fetchall()}
                return col_cache[table]

            for idx_name, table, col_expr, required in specs:
                if not table_columns(table).issuperset(required):
                    continue
                cursor.execute(
                    f"CREATE INDEX IF NOT EXISTS {idx_name} ON {table} ({col_expr})"
                )
            conn.commit()
        finally:
            conn.close()

    def save_folder(self, folder_name: str, folder_path: str, parent_path: str = "", full_path: str = "", agent_id: str = None):
        if getattr(self, '_use_app', False):
            self._service.save_folder(folder_name, folder_path, parent_path, full_path, agent_id=agent_id)
            return
        conn = sqlite3.connect(self.db_file)
        cursor = conn.cursor()
        
        # Check if exists for this agent
        cursor.execute('SELECT id FROM folders WHERE folder_path = ? AND (agent_id = ? OR agent_id IS NULL OR agent_id = "")', 
                      (folder_path, agent_id))
        if cursor.fetchone():
            conn.close()
            return

        cursor.execute('''
            INSERT INTO folders 
            (folder_name, folder_path, parent_path, full_path, created_date, agent_id)
            VALUES (?, ?, ?, ?, ?, ?)
        ''', (folder_name, folder_path, parent_path, full_path, datetime.now().isoformat(), agent_id))
        
        conn.commit()
        conn.close()

    def get_all_folders(self) -> List[tuple]:
        if getattr(self, '_use_app', False):
            return self._service.get_all_folders()
        conn = sqlite3.connect(self.db_file)
        cursor = conn.cursor()
        cursor.execute('''SELECT folder_name, folder_path, file_count, parent_path, full_path FROM folders ORDER BY full_path''')
        results = cursor.fetchall()
        conn.close()
        return results

    def get_files_in_folder(self, folder_path: str) -> List[tuple]:
        if getattr(self, '_use_app', False):
            return self._service.get_files_in_folder(folder_path)
        conn = sqlite3.connect(self.db_file)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT * FROM files 
            WHERE folder_path = ?
            ORDER BY original_filename
        ''', (folder_path,))
        
        results = cursor.fetchall()
        conn.close()
        return results

    def save_file(self, file_id: str, original_filename: str, file_type: str,
                  file_path: str, content: str, folder_path: str = "",
                  relative_path: str = "", image_path: str = "", thumbnail_path: str = ""):
        if getattr(self, '_use_app', False):
            self._service.save_file(file_id, original_filename, file_type, file_path, content,
                                    folder_path, relative_path, image_path, thumbnail_path)
            return
        conn = sqlite3.connect(self.db_file)
        cursor = conn.cursor()
        file_size = os.path.getsize(file_path)
        preview = content[:500] if len(content) > 500 else content
        
        cursor.execute('''
            INSERT OR REPLACE INTO files 
            (id, filename, original_filename, file_type, file_path, folder_path, relative_path,
             upload_date, file_size, content_preview, full_content, image_path, thumbnail_path)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        ''', (file_id, os.path.basename(file_path), original_filename, file_type, 
              file_path, folder_path, relative_path, datetime.now().isoformat(), file_size, preview, "",
              image_path, thumbnail_path))
        
        if folder_path:
            cursor.execute('''
                UPDATE folders 
                SET file_count = (SELECT COUNT(*) FROM files WHERE folder_path = ?)
                WHERE folder_path = ?
            ''', (folder_path, folder_path))
        
        conn.commit()
        conn.close()

    def get_file(self, file_id: str) -> Optional[tuple]:
        if getattr(self, '_use_app', False):
            return self._service.get_file(file_id)
        conn = sqlite3.connect(self.db_file)
        cursor = conn.cursor()
        cursor.execute('SELECT * FROM files WHERE id = ?', (file_id,))
        result = cursor.fetchone()
        conn.close()
        return result

    def get_all_files(self) -> List[tuple]:
        if getattr(self, '_use_app', False):
            return self._service.get_all_files()
        conn = sqlite3.connect(self.db_file)
        cursor = conn.cursor()
        cursor.execute('SELECT * FROM files ORDER BY upload_date DESC')
        results = cursor.fetchall()
        conn.close()
        return results

    def search_files(
        self,
        keyword: str,
        folder_path: Optional[str] = None,
        limit: Optional[int] = None,
    ) -> List[tuple]:
        """SQL-level keyword search (uses indexes, omits full_content, bounded by LIMIT)."""
        if getattr(self, '_use_app', False):
            return self._service.search_files(keyword, folder_path=folder_path, limit=limit)
        kw = (keyword or "").strip()
        if not kw:
            return []
        like = f"%{kw.lower()}%"
        sql = (
            "SELECT id, filename, original_filename, file_type, file_path, "
            "folder_path, relative_path, upload_date, file_size, content_preview, "
            "image_path, thumbnail_path "
            "FROM files WHERE ("
            "LOWER(original_filename) LIKE ? "
            "OR LOWER(folder_path) LIKE ? "
            "OR LOWER(content_preview) LIKE ?"
            ")"
        )
        params: List[Any] = [like, like, like]
        if folder_path:
            sql += " AND LOWER(folder_path) LIKE ?"
            params.append(f"%{folder_path.strip().lower()}%")
        sql += " ORDER BY upload_date DESC"
        if limit is not None:
            sql += " LIMIT ?"
            params.append(limit)
        conn = sqlite3.connect(self.db_file)
        try:
            cursor = conn.cursor()
            cursor.execute(sql, params)
            return cursor.fetchall()
        finally:
            conn.close()

    def save_extracted_image(self, file_id: str, image_path: str, image_type: str, page_number: int, agent_id: str = None):
        if getattr(self, '_use_app', False):
            self._service.save_extracted_image(file_id, image_path, image_type, page_number, agent_id=agent_id)
            return
        conn = sqlite3.connect(self.db_file)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT INTO extracted_images 
            (file_id, image_path, image_type, page_number, upload_date, agent_id)
            VALUES (?, ?, ?, ?, ?, ?)
        ''', (file_id, image_path, image_type, page_number, datetime.now().isoformat(), agent_id))
        
        conn.commit()
        conn.close()

    def get_extracted_images(self, file_id: str) -> List[tuple]:
        if getattr(self, '_use_app', False):
            return self._service.get_extracted_images(file_id)
        conn = sqlite3.connect(self.db_file)
        cursor = conn.cursor()
        cursor.execute('''SELECT id, image_path, image_type, page_number FROM extracted_images WHERE file_id = ? ORDER BY page_number, id''', (file_id,))
        results = cursor.fetchall()
        conn.close()
        return results

    def delete_file(self, file_id: str):
        if getattr(self, '_use_app', False):
            self._service.delete_file(file_id)
            return
        conn = sqlite3.connect(self.db_file)
        cursor = conn.cursor()
        
        cursor.execute('SELECT folder_path FROM files WHERE id = ?', (file_id,))
        result = cursor.fetchone()
        folder_path = result[0] if result else None
        
        cursor.execute('DELETE FROM extracted_images WHERE file_id = ?', (file_id,))
        cursor.execute('DELETE FROM files WHERE id = ?', (file_id,))
        
        if folder_path:
            cursor.execute('''
                UPDATE folders 
                SET file_count = (SELECT COUNT(*) FROM files WHERE folder_path = ?)
                WHERE folder_path = ?
            ''', (folder_path, folder_path))
        
        conn.commit()
        conn.close()

    def clear_all(self):
        if getattr(self, '_use_app', False):
            self._service.clear_all()
            return
        conn = sqlite3.connect(self.db_file)
        cursor = conn.cursor()
        cursor.execute('DELETE FROM extracted_images')
        cursor.execute('DELETE FROM files')
        cursor.execute('DELETE FROM folders')
        conn.commit()
        conn.close()




class ContentExtractor:
    
    @staticmethod
    def _mostly_whitespace(s: str) -> bool:
        """Check if string is mostly whitespace"""
        if not s:
            return True
        non_ws = len(re.sub(r"\s+", "", s))
        return non_ws < max(40, int(0.10 * len(s)))
    
    @staticmethod
    def _page_to_image(page: "pymupdf.Page", dpi: int = 300) -> Image.Image:
        """Render a page to a PIL Image"""
        mat = pymupdf.Matrix(dpi / 72, dpi / 72)
        pix = page.get_pixmap(matrix=mat, alpha=False)
        img_bytes = pix.tobytes("png")
        return Image.open(io.BytesIO(img_bytes))
    
    @staticmethod
    def _ocr_image(img: Image.Image) -> str:
        """Perform OCR on image using Tesseract"""
        return pytesseract.image_to_string(img, lang="eng")
    
    @staticmethod
    def extract_from_pdf(file, force_ocr: bool = False, ocr_dpi: int = 300) -> str:
        """
        Extract text from PDF using PyMuPDF's native text first; 
        fall back to Tesseract OCR when needed.
        """
        try:
            # Handle file input (path string or file object)
            if isinstance(file, str):
                # File path provided
                doc = pymupdf.open(file)
            else:
                # File object provided - read bytes
                if hasattr(file, 'read'):
                    data = file.read()
                    if hasattr(file, 'seek'):
                        file.seek(0)  # Reset file pointer for potential reuse
                else:
                    data = file
                
                if not data:
                    return "Error: Empty file"
                
                # Open from bytes
                doc = pymupdf.open(stream=data, filetype="pdf")
            
            # Handle password-protected PDFs
            if doc.needs_pass:
                # Try common passwords or return error
                common_passwords = ["", "password", "123456"]
                authenticated = False
                for pwd in common_passwords:
                    if doc.authenticate(pwd):
                        authenticated = True
                        break
                
                if not authenticated:
                    return "Error: PDF is password protected and requires authentication"
            
            all_text = []
            
            # Process each page
            for i in range(doc.page_count):
                page = doc.load_page(i)
                method = "text"
                page_text = page.get_text("text") if not force_ocr else ""
                
                # Use OCR if forced or if native text extraction yields poor results
                if force_ocr or not page_text or ContentExtractor._mostly_whitespace(page_text):
                    try:
                        img = ContentExtractor._page_to_image(page, dpi=ocr_dpi)
                        page_text = ContentExtractor._ocr_image(img)
                        method = "ocr"
                    except Exception as ocr_error:
                        method = "ocr_failed"
                        page_text = page_text or ""
                        print(f"OCR failed for page {i + 1}: {str(ocr_error)}")
                
                all_text.append(page_text)
            
            doc.close()
            
            full_text = "\n\n".join(all_text)
            return full_text.strip() if full_text.strip() else "No text extracted from PDF"
            
        except Exception as e:
            return f"Error reading PDF: {str(e)}"
    
    @staticmethod
    def extract_from_excel(file) -> str:
        try:
            df = pd.read_excel(file, sheet_name=None)
            content = []
            for sheet_name, sheet_data in df.items():
                content.append(f"Sheet: {sheet_name}\n")
                content.append(sheet_data.to_string())
            return "\n\n".join(content)
        except Exception as e:
            return f"Error reading Excel: {str(e)}"
    
    @staticmethod
    def extract_from_csv(file) -> str:
        try:
            df = pd.read_csv(file)
            return df.to_string()
        except Exception as e:
            return f"Error reading CSV: {str(e)}"
    
    @staticmethod
    def extract_from_docx(file) -> str:
        """
        Extract text from DOCX file with enhanced error handling
        """
        try:
            doc = docx.Document(file)
            text = []
            
            # Extract text from paragraphs
            for paragraph in doc.paragraphs:
                if paragraph.text.strip():
                    text.append(paragraph.text)
            
            # Extract text from tables
            for table in doc.tables:
                for row in table.rows:
                    row_text = []
                    for cell in row.cells:
                        if cell.text.strip():
                            row_text.append(cell.text.strip())
                    if row_text:
                        text.append(" | ".join(row_text))
            
            result = "\n".join(text)
            return result if result.strip() else "No text extracted from DOCX"
            
        except Exception as e:
            return f"Error reading DOCX: {str(e)}"
    
    @staticmethod
    def extract_from_doc(file) -> str:
        """
        Extract text from legacy .doc (Word 97-2003) file using pypandoc.
        Requires pandoc binary to be installed (e.g. apt-get install pandoc, or choco install pandoc).
        """
        try:
            import tempfile
            import pypandoc  # type: ignore[import-untyped]
            path = None
            if isinstance(file, str) and os.path.isfile(file):
                path = file
            else:
                data = file.read() if hasattr(file, "read") else file
                if hasattr(file, "seek"):
                    file.seek(0)
                if not data:
                    return "Error: Empty file"
                fd, path = tempfile.mkstemp(suffix=".doc")
                try:
                    os.write(fd, data)
                    os.close(fd)
                except Exception:
                    os.close(fd)
                    raise
            try:
                text = pypandoc.convert_file(path, "plain", extra_args=["--wrap=none"])
                result = (text or "").strip()
                return result if result else "No text extracted from DOC"
            finally:
                if path and not (isinstance(file, str) and file == path):
                    try:
                        os.unlink(path)
                    except OSError:
                        pass
        except ImportError:
            return "Error: pypandoc is required for .doc processing. Install with: pip install pypandoc_binary (includes pandoc), or pip install pypandoc and install the pandoc binary from https://pandoc.org/installing.html"
        except Exception as e:
            err_msg = str(e)
            if "pandoc" in err_msg.lower() and ("not found" in err_msg.lower() or "was found" in err_msg.lower()):
                return "Error reading DOC: The pandoc binary is required. Install with: pip install pypandoc_binary (includes pandoc), or install pandoc from https://pandoc.org/installing.html"
            return f"Error reading DOC: {err_msg}"
    
    @staticmethod
    def extract_from_txt(file) -> str:
        try:
            if isinstance(file, str):
                with open(file, 'r', encoding='utf-8') as f:
                    return f.read()
            else:
                if hasattr(file, 'read'):
                    content = file.read()
                    if isinstance(content, bytes):
                        return content.decode('utf-8')
                    return content
                return file.decode('utf-8')
        except UnicodeDecodeError:
            # Try different encodings
            try:
                if isinstance(file, str):
                    with open(file, 'r', encoding='latin-1') as f:
                        return f.read()
                else:
                    if hasattr(file, 'seek'):
                        file.seek(0)
                    content = file.read() if hasattr(file, 'read') else file
                    return content.decode('latin-1')
            except Exception as e:
                return f"Error reading TXT with alternative encoding: {str(e)}"
        except Exception as e:
            return f"Error reading TXT: {str(e)}"
    
    @staticmethod
    def extract_from_image(file, dpi: int = 300) -> str:
        """
        Extract text from image using OCR with enhanced preprocessing
        """
        try:
            # Open image
            if isinstance(file, str):
                image = Image.open(file)
            else:
                if hasattr(file, 'seek'):
                    file.seek(0)
                image = Image.open(file)
            
            # Convert to RGB if necessary
            if image.mode not in ('RGB', 'L'):
                image = image.convert('RGB')
            
            # Perform OCR
            text = pytesseract.image_to_string(image, lang="eng")
            
            result = text.strip()
            return result if result else "No text detected in image"
            
        except Exception as e:
            # Fallback: Try converting image to PDF and extract
            try:
                if isinstance(file, str):
                    image = Image.open(file)
                else:
                    if hasattr(file, 'seek'):
                        file.seek(0)
                    image = Image.open(file)
                
                # Convert image to PDF in memory
                pdf_bytes = io.BytesIO()
                image_rgb = image.convert('RGB')
                image_rgb.save(pdf_bytes, format='PDF')
                pdf_bytes.seek(0)
                
                # Use PDF extraction with OCR
                return ContentExtractor.extract_from_pdf(pdf_bytes, force_ocr=True, ocr_dpi=dpi)
                
            except Exception as fallback_error:
                return f"Error reading image: {str(e)} | Fallback error: {str(fallback_error)}"


class ImageExtractor:
    def __init__(self, output_folder: str = "extracted_images"):
        self.output_folder = output_folder
        Path(output_folder).mkdir(parents=True, exist_ok=True)
        self._image_counter = {}
    
    def _get_unique_image_id(self, file_id: str, base_name: str) -> str:
        """Generate unique image ID to avoid duplicates"""
        key = f"{file_id}_{base_name}"
        if key not in self._image_counter:
            self._image_counter[key] = 0
        self._image_counter[key] += 1
        return f"{base_name}_{self._image_counter[key]}"
    
    def extract_from_pdf(self, file_path: str, file_id: str, high_quality: bool = True) -> List[Dict]:
        """
        Extract images from PDF with fallback to high-quality page snapshots.
        
        Args:
            file_path: Path to PDF file
            file_id: Unique identifier for the file
            high_quality: If True, use higher DPI for page snapshots
        """
        extracted_images = []
        
        try:
            doc = pymupdf.open(file_path)
            
            # Extract embedded images
            for page_num in range(len(doc)):
                page = doc[page_num]
                image_list = page.get_images()
                
                for img_index, img in enumerate(image_list):
                    try:
                        xref = img[0]
                        base_image = doc.extract_image(xref)
                        image_bytes = base_image["image"]
                        image_ext = base_image["ext"]
                        
                        unique_id = self._get_unique_image_id(file_id, f"p{page_num + 1}_img{img_index + 1}")
                        img_filename = f"{file_id}_{unique_id}.{image_ext}"
                        img_path = os.path.join(self.output_folder, img_filename)
                        
                        with open(img_path, "wb") as img_file:
                            img_file.write(image_bytes)
                        
                        extracted_images.append({
                            'path': img_path,
                            'page': page_num + 1,
                            'type': 'embedded_image',
                            'index': img_index + 1
                        })
                    except Exception as e:
                        print(f"Error extracting embedded image on page {page_num + 1}: {e}")
                        continue
            
            # Extract page snapshots with enhanced quality
            dpi_scale = 3 if high_quality else 2  # 3x = ~216 DPI, 2x = ~144 DPI
            
            for page_num in range(len(doc)):
                try:
                    page = doc[page_num]
                    mat = pymupdf.Matrix(dpi_scale, dpi_scale)
                    pix = page.get_pixmap(matrix=mat, alpha=False)
                    
                    # Convert to PIL Image
                    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                    
                    unique_id = self._get_unique_image_id(file_id, f"page{page_num + 1}")
                    page_img_filename = f"{file_id}_{unique_id}.png"
                    page_img_path = os.path.join(self.output_folder, page_img_filename)
                    
                    # Save with optimization
                    img.save(page_img_path, "PNG", optimize=True)
                    
                    extracted_images.append({
                        'path': page_img_path,
                        'page': page_num + 1,
                        'type': 'page_snapshot',
                        'index': 0,
                        'dpi': dpi_scale * 72  # Approximate DPI
                    })
                except Exception as e:
                    print(f"Error extracting page snapshot for page {page_num + 1}: {e}")
                    continue
            
            doc.close()
            
        except Exception as e:
            print(f"Error opening PDF for image extraction: {e}")
            # Try alternative approach with lower quality if high quality failed
            if high_quality:
                print("Retrying with standard quality...")
                return self.extract_from_pdf(file_path, file_id, high_quality=False)
        
        return extracted_images
    
    def extract_from_docx(self, file_path: str, file_id: str) -> List[Dict]:
        """
        Extract images from DOCX file with enhanced error handling
        """
        extracted_images = []
        
        try:
            doc = docx.Document(file_path)
            img_index = 0
            
            # Extract images from relationships
            for rel in doc.part.rels.values():
                if "image" in rel.target_ref:
                    try:
                        img_data = rel.target_part.blob
                        ext = rel.target_ref.split('.')[-1]
                        img_index += 1
                        
                        unique_id = self._get_unique_image_id(file_id, f"docx_img{img_index}")
                        img_filename = f"{file_id}_{unique_id}.{ext}"
                        img_path = os.path.join(self.output_folder, img_filename)
                        
                        with open(img_path, "wb") as img_file:
                            img_file.write(img_data)
                        
                        # Verify image is valid
                        try:
                            img_verify = Image.open(img_path)
                            img_verify.verify()
                            
                            extracted_images.append({
                                'path': img_path,
                                'page': 1,
                                'type': 'docx_image',
                                'index': img_index,
                                'size': len(img_data)
                            })
                        except Exception as verify_error:
                            print(f"Image verification failed for {img_filename}: {verify_error}")
                            # Remove invalid image file
                            if os.path.exists(img_path):
                                os.remove(img_path)
                            
                    except Exception as e:
                        print(f"Error extracting DOCX image {img_index}: {e}")
                        continue
        
        except Exception as e:
            print(f"Error opening DOCX for image extraction: {e}")
        
        return extracted_images
    
    def extract_from_image_file(self, file_path: str, file_id: str) -> List[Dict]:
        """
        Process standalone image files (for consistency with other extractors)
        """
        extracted_images = []
        
        try:
            # Verify and copy image to output folder
            img = Image.open(file_path)
            img.verify()
            
            # Reopen after verify
            img = Image.open(file_path)
            
            ext = file_path.split('.')[-1].lower()
            unique_id = self._get_unique_image_id(file_id, "image")
            img_filename = f"{file_id}_{unique_id}.{ext}"
            img_path = os.path.join(self.output_folder, img_filename)
            
            # Save optimized copy
            if img.mode not in ('RGB', 'L'):
                img = img.convert('RGB')
            img.save(img_path, format='PNG', optimize=True)
            
            extracted_images.append({
                'path': img_path,
                'page': 1,
                'type': 'standalone_image',
                'index': 1,
                'original_path': file_path
            })
            
        except Exception as e:
            print(f"Error processing image file: {e}")
        
        return extracted_images

class PreviewGenerator:
    def __init__(self, output_folder: str = "document_images"):
        self.output_folder = output_folder
        Path(output_folder).mkdir(parents=True, exist_ok=True)
    
    def generate_preview(self, file_path: str, file_type: str, file_id: str) -> Tuple[Optional[str], Optional[str]]:
        generators = {
            'pdf': self._generate_pdf_preview,
            'png': self._generate_image_preview,
            'jpg': self._generate_image_preview,
            'jpeg': self._generate_image_preview,
        }
        
        generator = generators.get(file_type)
        if generator:
            return generator(file_path, file_id)
        return None, None
    
    def _generate_pdf_preview(self, file_path: str, file_id: str) -> Tuple[Optional[str], Optional[str]]:
        try:
            doc = pymupdf.open(file_path)
            page = doc[0]
            
            mat = pymupdf.Matrix(2, 2)
            pix = page.get_pixmap(matrix=mat)
            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
            
            image_path = os.path.join(self.output_folder, f"{file_id}_preview.png")
            img.save(image_path, "PNG", optimize=True)
            
            thumbnail = img.copy()
            thumbnail.thumbnail((400, 400), Image.Resampling.LANCZOS)
            thumbnail_path = os.path.join(self.output_folder, f"{file_id}_thumb.png")
            thumbnail.save(thumbnail_path, "PNG", optimize=True)
            
            doc.close()
            return image_path, thumbnail_path
        except:
            return None, None
    
    def _generate_image_preview(self, file_path: str, file_id: str) -> Tuple[Optional[str], Optional[str]]:
        try:
            img = Image.open(file_path)
            
            if img.mode not in ('RGB', 'RGBA'):
                img = img.convert('RGB')
            
            image_path = os.path.join(self.output_folder, f"{file_id}_preview.png")
            img.save(image_path, "PNG", optimize=True)
            
            thumbnail = img.copy()
            thumbnail.thumbnail((400, 400), Image.Resampling.LANCZOS)
            thumbnail_path = os.path.join(self.output_folder, f"{file_id}_thumb.png")
            thumbnail.save(thumbnail_path, "PNG", optimize=True)
            
            return image_path, thumbnail_path
        except:
            return None, None


class ChromaVectorDB:
    def __init__(self, db_path: str = "chroma_db", openai_api_key: str = None):
        self.db_path = db_path
        self.openai_api_key = openai_api_key
        self.client = None
        self.collection = None
        self.init_chromadb()
    
    def init_chromadb(self):
        openai_ef = embedding_functions.OpenAIEmbeddingFunction(
            api_key=self.openai_api_key,
            model_name="text-embedding-3-large"
        )
        
        self.client = chromadb.PersistentClient(
            path=self.db_path,
            settings=Settings(
                anonymized_telemetry=False,
                allow_reset=True
            )
        )
        
        try:
            self.collection = self.client.get_collection(
                name="documents",
                embedding_function=openai_ef
            )
        except:
            self.collection = self.client.create_collection(
                name="documents",
                embedding_function=openai_ef,
                metadata={"hnsw:space": "cosine"}
            )
    
    def chunk_text(self, text: str, chunk_size: int = 1000, overlap: int = 200) -> List[str]:
        chunks = []
        start = 0
        
        while start < len(text):
            end = start + chunk_size
            chunk = text[start:end]
            chunks.append(chunk)
            start += (chunk_size - overlap)
        
        return chunks
    
    def add_document(self, file_id: str, content: str, metadata: dict):
        chunks = self.chunk_text(content)
        
        ids = []
        documents = []
        metadatas = []
        
        for idx, chunk in enumerate(chunks):
            chunk_id = f"{file_id}_chunk_{idx}"
            ids.append(chunk_id)
            documents.append(chunk)
            
            chunk_metadata = metadata.copy()
            chunk_metadata.update({
                'chunk_index': idx,
                'total_chunks': len(chunks),
                'file_id': file_id
            })
            metadatas.append(chunk_metadata)
        
        self.collection.add(
            ids=ids,
            documents=documents,
            metadatas=metadatas
        )
        
        return len(chunks)
    
    def search(self, query: str, n_results: int = 5, folder_path: str = None) -> dict:
        where_filter = None
        if folder_path:
            where_filter = {"folder_path": folder_path}
        
        results = self.collection.query(
            query_texts=[query],
            n_results=n_results,
            where=where_filter
        )
        
        return results
    
    def get_document(self, file_id: str) -> Optional[str]:
        try:
            results = self.collection.get(
                where={"file_id": file_id},
                include=["documents", "metadatas"]
            )
            
            if results['documents']:
                chunks_with_meta = list(zip(results['documents'], results['metadatas']))
                chunks_with_meta.sort(key=lambda x: x[1].get('chunk_index', 0))
                return "\n".join([chunk for chunk, _ in chunks_with_meta])
            
        except Exception as e:
            print(f"Error getting document: {e}")
        
        return None
    
    def delete_document(self, file_id: str):
        try:
            results = self.collection.get(
                where={"file_id": file_id}
            )
            
            if results['ids']:
                self.collection.delete(ids=results['ids'])
        except Exception as e:
            print(f"Error deleting document: {e}")
    
    def count(self) -> int:
        return self.collection.count()
    
    def clear_all(self):
        try:
            self.client.delete_collection("documents")
            self.init_chromadb()
        except Exception as e:
            print(f"Error clearing collection: {e}")


class AIAnalyzer:
    def __init__(self, api_key: str, claude_api_key: str = ""):
        self.api_key = api_key
        self._llm = (
            LLMService(api_key=api_key or "", claude_api_key=claude_api_key or "")
            if (api_key or "").strip() or (claude_api_key or "").strip()
            else None
        )
    
    def extract_folder_from_query(self, query: str, available_folders: List[str]) -> Optional[str]:
        if not self._llm or not available_folders:
            return None
        
        try:
            folder_list = "\n".join([f"- {folder}" for folder in available_folders])
            
            extraction_prompt = f"""User query: "{query}"

Available folders:
{folder_list}

Determine if user asks to search within a specific folder.
Respond with ONLY JSON:
{{
    "folder_name": "exact folder name or null"
}}"""

            response_text = self._llm.chat_completion(
                messages=[
                    {"role": "system", "content": "Query parser. Respond with valid JSON only."},
                    {"role": "user", "content": extraction_prompt}
                ],
                temperature=0.2,
                max_tokens=200
            ).strip()
            response_text = re.sub(r'```json\s*|\s*```', '', response_text).strip()
            
            parsed = json.loads(response_text)
            folder_name = parsed.get('folder_name')
            
            if folder_name and folder_name in available_folders:
                return folder_name
                
        except Exception as e:
            print(f"Folder extraction error: {e}")
        
        return None
    
    def analyze_document(self, content: str, filename: str) -> str:
        if not self._llm:
            return "API key not set"
        
        try:
            prompt = f"""Analyze this document:

Document: {filename}
Content: {content[:4000]}

Provide:
1. Summary
2. Key insights
3. Important data
4. Conclusions"""
            
            return self._llm.chat_completion(
                messages=[
                    {"role": "system", "content": "Document analyst"},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.7
            )
        except Exception as e:
            return f"Error: {str(e)}"
    
    def summarize_search_results(self, query: str, results: List[Dict], folder_detected: str = None) -> str:
        if not self._llm or not results:
            return ""
        
        try:
            doc_summaries = []
            for idx, result in enumerate(results, 1):
                filename = result.get('filename', 'Unknown')
                folder = result.get('folder_path', 'Root')
                file_type = result.get('file_type', 'Unknown')
                relevance = result.get('relevance', 0)
                content_preview = result.get('content', '')[:2000]
                
                doc_summaries.append(
                    f"""DOCUMENT {idx}:
File: {filename}
Folder: {folder}
Type: {file_type}
Relevance: {relevance:.1f}%

Content:
{content_preview}
"""
                )
            
            combined_context = "\n\n".join(doc_summaries)
            
            folder_context = ""
            if folder_detected:
                folder_context = f"Auto-detected folder: '{folder_detected}'\n\n"
            
            summary_prompt = f"""{folder_context}Query: "{query}"

Documents found:
{combined_context}

Create structured summary grouped by folder with key findings."""

            return self._llm.chat_completion(
                messages=[
                    {"role": "system", "content": "Document analyst"},
                    {"role": "user", "content": summary_prompt}
                ],
                temperature=0.6,
                max_tokens=4000
            )
            
        except Exception as e:
            print(f"Summary error: {e}")
            return ""
    
    def find_in_document(self, query: str, content: str, filename: str) -> str:
        if not self._llm:
            return None
        
        try:
            lines = content.split('\n')
            matching_sections = []
            
            query_lower = query.lower()
            for i, line in enumerate(lines):
                if query_lower in line.lower():
                    start = max(0, i-8)
                    end = min(len(lines), i+9)
                    context = '\n'.join(lines[start:end])
                    matching_sections.append(context)
                    
                    if len(matching_sections) >= 5:
                        break
            
            if not matching_sections:
                return None
            
            matches_text = "\n\n---\n\n".join(matching_sections)
            
            summary_prompt = f"""Found "{query}" in "{filename}".

Matching sections:
{matches_text}

Create structured summary with key findings."""

            return self._llm.chat_completion(
                messages=[
                    {"role": "system", "content": "Document analyst"},
                    {"role": "user", "content": summary_prompt}
                ],
                temperature=0.6,
                max_tokens=2000
            )
            
        except Exception as e:
            print(f"AI summary error: {e}")
            return None


class DocumentAssistant:
    def __init__(self, 
                 upload_folder: str = "uploaded_files",
                 image_folder: str = "document_images",
                 extracted_folder: str = "extracted_images",
                 db_file: str = "file_sys_metadata.db",
                 chroma_path: str = "file_vectors_db",
                 openai_api_key: str = None,
                 claude_api_key: str = None,
                 agent_id: str = None):
        
        #Path(upload_folder).mkdir(exist_ok=True)
        Path(image_folder).mkdir(parents=True, exist_ok=True)
        Path(extracted_folder).mkdir(parents=True, exist_ok=True)
        
        self.upload_folder = upload_folder
        self.image_folder = image_folder
        self.extracted_folder = extracted_folder
        self.db = DatabaseManager(db_file)
        self.content_extractor = ContentExtractor()
        self.image_extractor = ImageExtractor(extracted_folder)
        self.preview_generator = PreviewGenerator(image_folder)
        self.vector_db = ChromaVectorDB(chroma_path, openai_api_key) if openai_api_key else None
        self.ai_analyzer = (
            AIAnalyzer(openai_api_key or "", claude_api_key or "")
            if openai_api_key
            else None
        )
        self.openai_client = OpenAI(api_key=openai_api_key) if openai_api_key else None
        self.agent_id = agent_id
    
    def scan_folder_path(self, folder_path: str) -> Dict:
        """Scan folder and upload all files preserving structure"""
        if not os.path.exists(folder_path) or not os.path.isdir(folder_path):
            return {'success': False, 'message': 'Invalid folder path'}
        
        supported_extensions = ['pdf', 'docx', 'doc', 'xlsx', 'xls', 'csv', 'txt', 'png', 'jpg', 'jpeg']
        
        uploaded_count = 0
        skipped_count = 0
        total_chunks = 0
        
        root_folder_name = os.path.basename(os.path.abspath(folder_path))
        
        # Save root folder
        self.db.save_folder(root_folder_name, root_folder_name, "", folder_path)
        
        for root, dirs, files in os.walk(folder_path):
            # Get relative path from root folder
            rel_dir = os.path.relpath(root, folder_path)
            
            # Create folder structure in database
            if rel_dir != '.':
                current_path_parts = rel_dir.split(os.sep)
                current_folder_path = root_folder_name
                
                for i, part in enumerate(current_path_parts):
                    parent_path = current_folder_path
                    current_folder_path = os.path.join(current_folder_path, part).replace(os.sep, '/')
                    
                    full_sys_path = os.path.join(folder_path, *current_path_parts[:i+1])
                    self.db.save_folder(part, current_folder_path, parent_path, full_sys_path)
                
                folder_for_file = current_folder_path
            else:
                folder_for_file = root_folder_name
            
            for filename in files:
                # Skip macOS AppleDouble/resource-fork files (._*); they are not real documents
                if filename.startswith('._'):
                    skipped_count += 1
                    continue
                file_ext = filename.split('.')[-1].lower()
                
                if file_ext not in supported_extensions:
                    continue
                
                file_full_path = os.path.join(root, filename)
                relative_file_path = os.path.relpath(file_full_path, folder_path)
                
                try:
                    with open(file_full_path, 'rb') as f:
                        file_bytes = io.BytesIO(f.read())
                        file_bytes.name = filename
                        
                        content, _ = self.process_file(file_bytes, file_ext)
                        
                        if content and not content.startswith("Error") and not content.startswith("Unsupported"):
                            file_id = self.get_file_hash(content, filename + relative_file_path)
                            
                            if self.db.get_file(file_id):
                                skipped_count += 1
                                continue
                            
                            # saved_path = os.path.join(self.upload_folder, f"{file_id}.{file_ext}")
                            
                            # file_bytes.seek(0)
                            # with open(saved_path, 'wb') as dest:
                            #     dest.write(file_bytes.read())

                            saved_path = file_full_path 
                            
                            image_path, thumbnail_path = self.preview_generator.generate_preview(
                                saved_path, file_ext, file_id
                            )
                            
                            # Reset image counter for each file to avoid duplicate IDs
                            self.image_extractor._image_counter = {}
                            
                            extracted_images = []
                            if file_ext == 'pdf':
                                extracted_images = self.image_extractor.extract_from_pdf(saved_path, file_id)
                            elif file_ext == 'docx':
                                extracted_images = self.image_extractor.extract_from_docx(saved_path, file_id)
                            
                            self.db.save_file(file_id, filename, file_ext, saved_path, 
                                            content, folder_for_file, relative_file_path, 
                                            image_path or "", thumbnail_path or "", self.agent_id)
                            
                            for img_info in extracted_images:
                                self.db.save_extracted_image(file_id, img_info['path'], 
                                                            img_info['type'], img_info['page'])
                            
                            if self.vector_db:
                                chunks = self.vector_db.add_document(
                                    file_id, 
                                    content,
                                    {
                                        "filename": filename,
                                        "file_type": file_ext,
                                        "folder_path": folder_for_file,
                                        "relative_path": relative_file_path,
                                        "upload_date": datetime.now().isoformat()
                                    }
                                )
                                total_chunks += chunks
                            
                            uploaded_count += 1
                        else:
                            skipped_count += 1
                            
                except Exception as e:
                    print(f"Error processing {filename}: {e}")
                    skipped_count += 1
        
        return {
            'success': True,
            'uploaded': uploaded_count,
            'skipped': skipped_count,
            'total_chunks': total_chunks,
            'folder_name': root_folder_name
        }
    
    def _get_all_files_in_folder_recursive(self, folder_path: str) -> List[tuple]:
        """
        Get all files in folder and all its subfolders recursively
        """
        all_files = []
        
        # Get direct files in this folder
        files = self.db.get_files_in_folder(folder_path)
        all_files.extend(files)
        
        # Get all subfolders
        all_folders = self.db.get_all_folders()
        
        for folder_info in all_folders:
            folder_name, current_path, file_count, parent_path, full_path = folder_info
            
            # Check if this folder is a subfolder of the target folder
            if current_path.startswith(folder_path + '/') or current_path == folder_path:
                subfolder_files = self.db.get_files_in_folder(current_path)
                all_files.extend(subfolder_files)
        
        return all_files
    
    def get_file_hash(self, content: str, filename: str) -> str:
        agent_part = str(getattr(self, 'agent_id', ''))
        return hashlib.md5(f"{filename}{agent_part}{content[:1000]}".encode()).hexdigest()
    
    def process_file(self, uploaded_file, file_type: str) -> Tuple[str, str]:
        if hasattr(uploaded_file, 'seek'):
            uploaded_file.seek(0)
        
        extractors = {
            'pdf': self.content_extractor.extract_from_pdf,
            'xlsx': self.content_extractor.extract_from_excel,
            'xls': self.content_extractor.extract_from_excel,
            'csv': self.content_extractor.extract_from_csv,
            'png': self.content_extractor.extract_from_image,
            'jpg': self.content_extractor.extract_from_image,
            'jpeg': self.content_extractor.extract_from_image,
            'docx': self.content_extractor.extract_from_docx,
            'doc': self.content_extractor.extract_from_doc,
            'txt': self.content_extractor.extract_from_txt
        }
        
        extractor = extractors.get(file_type)
        if extractor:
            content = extractor(uploaded_file)
            return content, file_type
        return "Unsupported file type", None
    
    def upload_document(self, uploaded_file, folder_path: str = "") -> Dict:
        if not self.vector_db:
            return {'success': False, 'message': 'OpenAI API key not set'}
        
        file_type = uploaded_file.name.split('.')[-1].lower()
        content, _ = self.process_file(uploaded_file, file_type)
        
        if not content or content.startswith("Error") or content.startswith("Unsupported"):
            return {'success': False, 'message': content}
        
        file_id = self.get_file_hash(content, uploaded_file.name)
        
        if self.db.get_file(file_id):
            return {'success': False, 'message': 'File already exists'}
        
        # file_extension = uploaded_file.name.split('.')[-1].lower()
        # filename = f"{file_id}.{file_extension}"
        # file_path = os.path.join(self.upload_folder, filename)
        
        # uploaded_file.seek(0)
        # with open(file_path, "wb") as f:
        #     f.write(uploaded_file.read())

        import tempfile
        file_extension = uploaded_file.name.split('.')[-1].lower()
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_extension}")
        uploaded_file.seek(0)
        temp_file.write(uploaded_file.read())
        temp_file.close()
        file_path = temp_file.name
        
        image_path, thumbnail_path = self.preview_generator.generate_preview(
            file_path, file_type, file_id
        )
        
        # Reset image counter for each file
        self.image_extractor._image_counter = {}
        
        extracted_images = []
        if file_type == 'pdf':
            extracted_images = self.image_extractor.extract_from_pdf(file_path, file_id)
        elif file_type == 'docx':
            extracted_images = self.image_extractor.extract_from_docx(file_path, file_id)
        
        self.db.save_file(file_id, uploaded_file.name, file_type, file_path, 
                         content, folder_path, uploaded_file.name, image_path or "", thumbnail_path or "", self.agent_id)
        
        for img_info in extracted_images:
            self.db.save_extracted_image(file_id, img_info['path'], 
                                        img_info['type'], img_info['page'])
        
        metadata = {
            "filename": uploaded_file.name,
            "file_type": file_type,
            "folder_path": folder_path,
            "upload_date": datetime.now().isoformat()
        }
        
        chunks_created = self.vector_db.add_document(file_id, content, metadata)
        
        return {
            'success': True, 
            'message': 'Upload successful',
            'file_id': file_id,
            'chunks_created': chunks_created,
            'images_extracted': len(extracted_images),
            'folder_path': folder_path
        }
    
    def _detect_folder_from_query(self, query: str) -> Optional[str]:
        folders = self.db.get_all_folders()
        if not folders:
            return None
        
        available_folders = [folder[1] for folder in folders]
        
        if self.ai_analyzer:
            detected_folder = self.ai_analyzer.extract_folder_from_query(query, available_folders)
            if detected_folder:
                return detected_folder
        
        return None
    
    def search_semantic(self, query: str, folder_path: str = None, n_results: int = 10) -> List[Dict]:
        if not self.vector_db:
            return []
        
        if not folder_path:
            folder_path = self._detect_folder_from_query(query)
        
        results = self.vector_db.search(query, n_results, folder_path)
        
        formatted_results = []
        
        if results['ids'][0]:
            for idx, (doc_id, document, metadata, distance) in enumerate(zip(
                results['ids'][0],
                results['documents'][0],
                results['metadatas'][0],
                results['distances'][0]
            )):
                relevance = max(0, 100 - (distance * 100))
                
                file_id = metadata.get('file_id')
                file_info = self.db.get_file(file_id)
                
                if file_info:
                    formatted_results.append({
                        'id': file_id,
                        'chunk_id': doc_id,
                        'filename': file_info[2],
                        'file_type': file_info[3],
                        'file_path': file_info[4],
                        'folder_path': file_info[5],
                        'relative_path': file_info[6],
                        'upload_date': file_info[7],
                        'relevance': relevance,
                        'distance': distance,
                        'chunk_text': document,
                        'content': document,
                        'chunk_index': metadata.get('chunk_index', 0),
                        'total_chunks': metadata.get('total_chunks', 1),
                        'image_path': file_info[11],
                        'thumbnail_path': file_info[12]
                    })
        
        return formatted_results
    
    def handle_query(self, query: str) -> Dict:
        query = query.strip()
        print(f"Handling query: {query}")
        if query.lower().startswith('@file'):
            search_query = query[5:].strip()
            results = self.search_semantic(search_query, n_results=10)
            
            folder_detected = self._detect_folder_from_query(search_query)
            
            folder_groups = {}
            for result in results:
                folder = result['folder_path'] or 'Root'
                if folder not in folder_groups:
                    folder_groups[folder] = []
                
                images = self.db.get_extracted_images(result['id'])
                
                folder_groups[folder].append({
                    'filename': result['filename'],
                    'folder_path': result['folder_path'],
                    'relative_path': result.get('relative_path', ''),
                    'relevance': result['relevance'],
                    'summary': result['chunk_text'][:300],
                    'file_id': result['id'],
                    'thumbnail': result.get('thumbnail_path'),
                    'all_images': [{'path': img[1], 'page': img[3], 'type': img[2]} for img in images],
                    'file_path': result['file_path'],
                    'file_type': result.get('file_type', 'Unknown')
                })
            
            overall_summary = ""
            if self.ai_analyzer:
                overall_summary = self.ai_analyzer.summarize_search_results(
                    search_query, results, folder_detected
                )
            
            return {
                'type': 'document_search',
                'query': search_query,
                'folder_detected': folder_detected,
                'folder_groups': folder_groups,
                'results': results,
                'overall_summary': overall_summary
            }
        
        elif query.lower().startswith('@analyze'):
            doc_query = query[8:].strip()
            
            # Detect folder from query
            folder_path = self._detect_folder_from_query(doc_query)
            
            # Get all files in detected folder (including subfolders)
            if folder_path:
                files = self._get_all_files_in_folder_recursive(folder_path)
            else:
                files = self.db.get_all_files()
            
            # Find file by name or ID
            file_info = None
            
            # First try exact filename match
            for f in files:
                if doc_query.lower() == f[2].lower() or f[0] == doc_query:
                    file_info = f
                    break
            
            # Then try partial filename match
            if not file_info:
                for f in files:
                    if doc_query.lower() in f[2].lower():
                        file_info = f
                        break
            
            # Finally try semantic search
            if not file_info and self.vector_db:
                search_results = self.vector_db.search(
                    doc_query, 
                    n_results=5,
                    folder_path=folder_path
                )
                
                if search_results['ids'][0]:
                    # Get the best matching file
                    best_file_id = search_results['metadatas'][0][0].get('file_id')
                    file_info = self.db.get_file(best_file_id)
            
            if not file_info:
                return {
                    'type': 'error',
                    'error': f'Document not found: {doc_query}',
                    'folder_detected': folder_path,
                    'files_searched': len(files)
                }
            
            # Get full content and images
            content = self.vector_db.get_document(file_info[0]) if self.vector_db else ""
            images = self.db.get_extracted_images(file_info[0])
            
            # Perform AI analysis
            analysis = None
            if self.ai_analyzer and content:
                analysis = self.ai_analyzer.analyze_document(content, file_info[2])
            
            return {
                'type': 'document_analysis',
                'filename': file_info[2],
                'file_type': file_info[3],
                'folder_path': file_info[5],
                'relative_path': file_info[6],
                'folder_detected': folder_path,
                'analysis': analysis or 'AI analysis not available',
                'all_images': [{'path': img[1], 'page': img[3], 'type': img[2]} for img in images],
                'file_path': file_info[4],
                'file_size': file_info[8],
                'upload_date': file_info[7]
            }
        
        elif query.lower().startswith('@find'):
            search_query = query[5:].strip()
            
            # Detect folder from query
            folder_path = self._detect_folder_from_query(search_query)
            
            # Get all files in detected folder or all files
            if folder_path:
                all_files = self._get_all_files_in_folder_recursive(folder_path)
            else:
                all_files = self.db.get_all_files()
            
            if not all_files:
                return {
                    'type': 'document_find',
                    'query': search_query,
                    'folder_detected': folder_path,
                    'results': [],
                    'message': 'No files found'
                }
            
            # Use semantic search for better relevance scoring
            file_results = []
            
            for file_info in all_files:
                if not file_info:
                    continue
                
                file_id = file_info[0]
                
                # Perform semantic search on this specific file
                semantic_results = self.vector_db.search(
                    search_query, 
                    n_results=3,  # Get top 3 chunks from this file
                    folder_path=None  # Don't filter by folder in vector search
                )
                
                # Calculate relevance for this file
                file_relevance = 0
                file_chunks = []
                
                if semantic_results['ids'][0]:
                    for chunk_id, distance, doc, metadata in zip(
                        semantic_results['ids'][0],
                        semantic_results['distances'][0],
                        semantic_results['documents'][0],
                        semantic_results['metadatas'][0]
                    ):
                        # Only consider chunks from this file
                        if metadata.get('file_id') == file_id:
                            # Convert distance to relevance percentage (0-100%)
                            chunk_relevance = max(0, 100 - (distance * 100))
                            file_chunks.append({
                                'text': doc,
                                'relevance': chunk_relevance
                            })
                
                # Calculate average relevance for this file
                if file_chunks:
                    file_relevance = sum(c['relevance'] for c in file_chunks) / len(file_chunks)
                    
                    # Also do keyword matching for exact phrase detection
                    content = self.vector_db.get_document(file_id)
                    if content:
                        search_lower = search_query.lower()
                        content_lower = content.lower()
                        
                        exact_matches = content_lower.count(search_lower)
                        
                        # Boost relevance if exact matches found
                        if exact_matches > 0:
                            file_relevance = min(100, file_relevance + (exact_matches * 5))
                        
                        file_results.append({
                            'file_info': file_info,
                            'relevance': file_relevance,
                            'exact_matches': exact_matches,
                            'best_chunks': file_chunks[:3],  # Top 3 chunks
                            'content': content
                        })
            
            # Sort by relevance
            file_results.sort(key=lambda x: x['relevance'], reverse=True)
            
            if not file_results or file_results[0]['relevance'] < 10:
                return {
                    'type': 'document_find',
                    'query': search_query,
                    'folder_detected': folder_path,
                    'results': [],
                    'message': 'No relevant matches found'
                }
            
            # Get the best match
            best = file_results[0]
            file_info = best['file_info']
            content = best['content']
            
            # Extract matching sections with context
            lines = content.split('\n')
            matching_sections = []
            matching_pages = set()
            search_lower = search_query.lower()
            
            for i, line in enumerate(lines):
                if search_lower in line.lower():
                    start = max(0, i - 8)
                    end = min(len(lines), i + 9)
                    context_lines = lines[start:end]
                    
                    context_text = []
                    for j, ctx_line in enumerate(context_lines):
                        if start + j == i:
                            context_text.append(f">>> {ctx_line} <<<")
                        else:
                            context_text.append(ctx_line)
                    
                    context = '\n'.join(context_text)
                    
                    if len(matching_sections) < 10:
                        matching_sections.append(context)
                    
                    page_num = (i // 50) + 1
                    matching_pages.add(page_num)
            
            # Get images
            all_images = self.db.get_extracted_images(file_info[0])
            page_images = {}
            
            for img in all_images:
                page = img[3]
                if page not in page_images:
                    page_images[page] = []
                page_images[page].append({
                    'path': img[1],
                    'type': img[2],
                    'is_matching': page in matching_pages
                })
            
            # AI summary
            ai_summary = None
            if self.ai_analyzer:
                # Use best chunks for AI summary
                chunks_text = '\n\n'.join([c['text'] for c in best['best_chunks']])
                ai_summary = self.ai_analyzer.find_in_document(
                    search_query, 
                    chunks_text, 
                    file_info[2]
                )
            
            best_result = {
                'filename': file_info[2],
                'file_id': file_info[0],
                'folder_path': file_info[5],
                'relative_path': file_info[6],
                'match_count': best['exact_matches'],
                'relevance_score': round(best['relevance'], 1),  # Percentage
                'matches': matching_sections,
                'matching_pages': sorted(matching_pages),
                'page_images': page_images,
                'file_path': file_info[4],
                'search_term': search_query,
                'ai_summary': ai_summary,
                'top_chunks': best['best_chunks']  # Include semantic chunks
            }
            
            return {
                'type': 'document_find',
                'query': search_query,
                'folder_detected': folder_path,
                'results': [best_result],
                'total_files_searched': len(all_files)
            }
        
        else:
            print("Performing general semantic search")
            results = self.search_semantic(query, n_results=5)
            folder_detected = self._detect_folder_from_query(query)
            
            folder_groups = {}
            for result in results:
                folder = result['folder_path'] or 'Root'
                if folder not in folder_groups:
                    folder_groups[folder] = []
                
                images = self.db.get_extracted_images(result['id'])
                
                folder_groups[folder].append({
                    'filename': result['filename'],
                    'folder_path': result['folder_path'],
                    'relative_path': result.get('relative_path', ''),
                    'relevance': result['relevance'],
                    'summary': result['chunk_text'][:300],
                    'file_id': result['id'],
                    'thumbnail': result.get('thumbnail_path'),
                    'all_images': [{'path': img[1], 'page': img[3], 'type': img[2]} for img in images],
                    'file_path': result['file_path'],
                    'file_type': result.get('file_type', 'Unknown')
                })
            
            overall_summary = ""
            if self.ai_analyzer:
                overall_summary = self.ai_analyzer.summarize_search_results(
                    query, results, folder_detected
                )
            
            return {
                'type': 'document_search',
                'query': query,
                'folder_detected': folder_detected,
                'folder_groups': folder_groups,
                'results': results,
                'overall_summary': overall_summary
            }
    
    def get_statistics(self) -> Dict:
        folders = self.db.get_all_folders()
        files = self.db.get_all_files()
        chunks = self.vector_db.count() if self.vector_db else 0
        
        return {
            'total_folders': len(folders),
            'total_files': len(files),
            'total_chunks': chunks,
            'avg_chunks_per_file': chunks / len(files) if files else 0
        }
    
    def delete_document(self, file_id: str):
        file_info = self.db.get_file(file_id)
        if file_info:
            # if os.path.exists(file_info[4]):
            #     os.remove(file_info[4])
            
            if file_info[11] and os.path.exists(file_info[11]):
                os.remove(file_info[11])
            if file_info[12] and os.path.exists(file_info[12]):
                os.remove(file_info[12])
            
            extracted_imgs = self.db.get_extracted_images(file_id)
            for img_info in extracted_imgs:
                if os.path.exists(img_info[1]):
                    os.remove(img_info[1])
            
            if self.vector_db:
                self.vector_db.delete_document(file_id)
            
            self.db.delete_file(file_id)
    
    def clear_all(self):
        if self.vector_db:
            self.vector_db.clear_all()
        
        self.db.clear_all()
        
        for folder in [self.upload_folder, self.image_folder, self.extracted_folder]:
            if os.path.exists(folder):
                for file in os.listdir(folder):
                    try:
                        os.remove(os.path.join(folder, file))
                    except:
                        pass
    
    def get_all_folders(self) -> List[tuple]:
        return self.db.get_all_folders()
    
    def get_files_in_folder(self, folder_path: str) -> List[Dict]:
        files = self.db.get_files_in_folder(folder_path)
        
        results = []
        for file_info in files:
            results.append({
                'id': file_info[0],
                'filename': file_info[2],
                'file_type': file_info[3],
                'file_path': file_info[4],
                'folder_path': file_info[5],
                'relative_path': file_info[6],
                'upload_date': file_info[7],
                'size': file_info[8],
                'preview': file_info[9]
            })
        
        return results

    def get_file_path(self, file_id: str) -> Optional[str]:
        """Get original file path for viewing/downloading"""
        file_info = self.db.get_file(file_id)
        if file_info and os.path.exists(file_info[4]):
            return file_info[4]
        return None

class DocumentUIHelper:
    @staticmethod
    def _generate_unique_key(base_key, *args):
        """Generate unique key for widgets to avoid duplicates in history"""
        combined = f"{base_key}_{'_'.join(map(str, args))}"
        return hashlib.md5(combined.encode()).hexdigest()[:12]
    

    @staticmethod
    def render_document_search(st_module, results, avatar_url, is_history=False):
        """Render @file search results - matching chat tab UI exactly"""
        asst_msg = st_module.chat_message("assistant", avatar=avatar_url)
        
        asst_msg.markdown(f"### 🔍 Results: {results['query']}")
        
        # Show folder detection
        if results.get('folder_detected'):
            asst_msg.success(f"🔍 Auto-detected folder: **{results['folder_detected']}**")
        
        # Show AI summary
        if results.get('overall_summary'):
            with asst_msg.expander("📝 AI Summary", expanded=True):
                asst_msg.markdown(results['overall_summary'])
        
        # Show document count
        asst_msg.caption(f"📊 {len(results.get('results', []))} documents found")
        asst_msg.markdown("---")
        
        # Group by folder
        for folder_name, docs in results.get('folder_groups', {}).items():
            asst_msg.markdown(f"### 📁 {folder_name} ({len(docs)} files)")
            
            for doc_idx, doc in enumerate(docs):
                # Generate unique key for this document
                doc_key = DocumentUIHelper._generate_unique_key("doc_search", folder_name, doc_idx, doc['file_id'])
                
                with asst_msg.expander(f"📄 {doc['filename']} • {doc['relevance']:.1f}% match", expanded=False):
                    if doc.get('relative_path'):
                        asst_msg.caption(f"📍 {doc['relative_path']}")
                    
                    asst_msg.markdown("**📋 Preview:**")
                    asst_msg.write(doc['summary'])
                    
                    # Metrics
                    col1, col2, col3 = asst_msg.columns(3)
                    with col1:
                        st_module.metric("Relevance", f"{doc['relevance']:.1f}%")
                    with col2:
                        st_module.metric("Type", doc.get('file_type', 'Unknown').upper())
                    with col3:
                        st_module.metric("Images", len(doc.get('all_images', [])))
                    
                    asst_msg.markdown("---")
                    
                    # Action buttons - 4 columns now
                    col_btn2, col_btn4 = asst_msg.columns(2)
                    
                    # with col_btn1:
                    #     analyze_key = f"analyze_{doc_key}"
                    #     if st_module.button("🤖 Analyze", key=analyze_key):
                    #         st_module.session_state['trigger_analyze'] = {
                    #             'filename': doc['filename'],
                    #             'file_id': doc['file_id']
                    #         }
                    #         st_module.rerun()
                    
                    with col_btn2:
                        images_key = f"images_{doc_key}"
                        if st_module.button("🖼️ Images", key=images_key):
                            toggle_key = f"show_images_{doc['file_id']}"
                            st_module.session_state[toggle_key] = not st_module.session_state.get(toggle_key, False)
                            st_module.rerun()
                    
                    # with col_btn3:
                    #     view_key = f"view_{doc_key}"
                    #     if st_module.button("👁️ View", key=view_key):
                    #         st_module.session_state['view_file'] = doc['file_path']
                    #         st_module.rerun()
                    
                    with col_btn4:
                        if os.path.exists(doc['file_path']):
                            with open(doc['file_path'], 'rb') as f:
                                dl_key = f"dl_{doc_key}"
                                st_module.download_button("⬇️ Download", f, file_name=doc['filename'], key=dl_key)
                    
                    # Show images if toggled - ALL IN SINGLE ROW
                    if st_module.session_state.get(f"show_images_{doc['file_id']}", False):
                        if doc.get('all_images'):
                            asst_msg.markdown("**🖼️ Extracted Images:**")
                            
                            # Single scrollable row for all images
                            num_images = len(doc['all_images'])
                            cols = asst_msg.columns(min(num_images, 10))  # Max 10 columns or less
                            
                            for i, img in enumerate(doc['all_images']):
                                col_idx = i % min(num_images, 10)
                                with cols[col_idx]:
                                    if os.path.exists(img['path']):
                                        st_module.image(img['path'], use_container_width=True)
                                        st_module.caption(f"P{img['page']}")
    
    @staticmethod
    def render_document_analysis(st_module, analysis, avatar_url, is_history=False):
        """Render @analyze results - matching chat tab UI exactly"""
        asst_msg = st_module.chat_message("assistant", avatar=avatar_url)
        
        asst_msg.markdown(f"### 📊 Analysis: {analysis['filename']}")
        asst_msg.caption(f"📁 Folder: {analysis.get('folder_path', 'Root')}")
        
        if analysis.get('relative_path'):
            asst_msg.caption(f"📍 Path: {analysis['relative_path']}")
        
        if analysis.get('folder_detected'):
            asst_msg.success(f"🔍 Auto-detected folder: **{analysis['folder_detected']}**")
        
        # Show analysis
        with asst_msg.expander("🤖 Detailed Analysis", expanded=True):
            asst_msg.write(analysis.get('analysis', 'No analysis available'))
        
        # Metrics
        col1, col2 = asst_msg.columns(2)
        with col1:
            st_module.metric("File Type", analysis.get('file_type', 'Unknown').upper())
        with col2:
            st_module.metric("Total Images", len(analysis.get('all_images', [])))
        

        # Show all images in a single row (no page grouping)
        if analysis.get('all_images'):
            asst_msg.markdown("### 🖼️ Document Images")
            
            # Add unique container key to prevent duplicate rendering
            container_key = DocumentUIHelper._generate_unique_key("img_container", analysis.get('filename', 'doc'))
            
            # Create a single row with all images
            num_images = len(analysis['all_images'])
            cols = asst_msg.columns(num_images)
            
            for i, img in enumerate(analysis['all_images']):
                with cols[i]:
                    if isinstance(img, dict):
                        img_path = img.get('path', '')
                        img_type = img.get('type', 'Image')
                        page_num = img.get('page', 1)
                        if img_path and os.path.exists(img_path):
                            st_module.image(img_path, use_container_width=True)
                            st_module.caption(f"P{page_num} - {img_type}")
    @staticmethod
    def render_document_find(st_module, results, avatar_url, is_history=False):
        """Render @find results - matching chat tab UI exactly"""
        asst_msg = st_module.chat_message("assistant", avatar=avatar_url)
        
        if not results.get('results'):
            asst_msg.warning(f"❌ No matches found for: **{results['query']}**")
            if results.get('folder_detected'):
                asst_msg.info(f"Searched in folder: {results['folder_detected']}")
            return
        
        doc = results['results'][0]
        
        asst_msg.success(f"✅ Best Match: **{doc['filename']}**")
        asst_msg.caption(f"📁 {doc.get('folder_path', 'Root')}")
        
        if doc.get('relative_path'):
            asst_msg.caption(f"📍 {doc['relative_path']}")
        
        if results.get('folder_detected'):
            asst_msg.info(f"🔍 Auto-detected folder: {results['folder_detected']}")
        
        # Metrics
        col1, col2, col3, col4 = asst_msg.columns(4)
        with col1:
            st_module.metric("Relevance", f"{doc['relevance_score']}%")
        with col2:
            st_module.metric("Exact Matches", doc['match_count'])
        with col3:
            st_module.metric("Pages", len(doc.get('matching_pages', [])))
        with col4:
            if results.get('total_files_searched'):
                st_module.metric("Files Searched", results['total_files_searched'])
        
        # Show matching pages
        if doc.get('matching_pages'):
            pages_str = ', '.join(map(str, doc['matching_pages'][:10]))
            if len(doc['matching_pages']) > 10:
                pages_str += f" ... +{len(doc['matching_pages']) - 10} more"
            asst_msg.info(f"📄 Found on pages: {pages_str}")
        
        asst_msg.markdown("---")
        
        # AI Summary
        if doc.get('ai_summary'):
            with asst_msg.expander("🤖 AI Summary", expanded=True):
                asst_msg.markdown(doc['ai_summary'])
        
        # Matching sections (collapsible)
        asst_msg.markdown("### 🔍 Matching Content")
        for i, match in enumerate(doc.get('matches', [])[:5]):
            match_key = DocumentUIHelper._generate_unique_key("match", doc.get('filename', 'doc'), i)
            with asst_msg.expander(f"Match {i+1}", expanded=(i < 2)):
                st_module.code(match, language="text")
        
        # Show images with matching indicator (collapsible)
        if doc.get('page_images'):
            asst_msg.markdown("### 🖼️ Document Images")
            
            matching_pages = set(doc.get('matching_pages', []))
            
            # Collect all images from all pages
            all_imgs = []
            for page_num in sorted(doc['page_images'].keys()):
                for img in doc['page_images'][page_num]:
                    img_data = img.copy()
                    img_data['page_num'] = page_num
                    all_imgs.append(img_data)
            
            # Display all images in a single row
            if all_imgs:
                num_images = len(all_imgs)
                cols = asst_msg.columns(num_images)
                for i, img in enumerate(all_imgs):
                    with cols[i]:
                        if os.path.exists(img['path']):
                            st_module.image(img['path'], use_container_width=True)
                            caption = f"P{img['page_num']} - {img['type']}"
                            if img['page_num'] in matching_pages:
                                caption += " ⭐"
                            st_module.caption(caption)
        
        # Download button
        asst_msg.markdown("---")
        if os.path.exists(doc['file_path']):
            with open(doc['file_path'], 'rb') as f:
                dl_key = DocumentUIHelper._generate_unique_key("dl_find", doc.get('filename', 'doc'))
                asst_msg.download_button(
                    "⬇️ Download Document", 
                    f, 
                    file_name=doc['filename'], 
                    use_container_width=True, 
                    key=dl_key
                )