app/utils/document_processor.py

# ============================================================================
#  DOCUMENT PROCESSOR
#  Handles PDF/text extraction, chunking, token management
# ============================================================================

import os
import re
import base64
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass
from enum import Enum

from app.utils.logger import logger


class DocumentType(str, Enum):
    PDF = "pdf"
    DOCX = "docx"
    TXT = "txt"
    MD = "md"
    JSON = "json"
    UNKNOWN = "unknown"


@dataclass
class DocumentChunk:
    """Represents a chunk of document content"""
    chunk_id: int
    content: str
    token_count: int
    source_document: str
    source_page: Optional[int] = None
    metadata: Optional[Dict[str, Any]] = None


@dataclass
class ProcessedDocument:
    """Represents a fully processed document"""
    document_id: str
    filename: str
    document_type: DocumentType
    total_tokens: int
    chunks: List[DocumentChunk]
    skipped: bool = False
    skip_reason: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None


class TokenCounter:
    """Estimate token count for text (approximate for LLMs)"""
    
    # Approximate tokens per character ratio (varies by model)
    # GPT/Groq models: ~4 chars per token on average
    CHARS_PER_TOKEN = 4
    
    @staticmethod
    def estimate(text: str) -> int:
        """Estimate token count for text"""
        if not text:
            return 0
        # Count words and apply multiplier for subword tokens
        words = len(text.split())
        chars = len(text)
        # Use both estimations and take average
        word_estimate = words * 1.3  # Words + subword tokens
        char_estimate = chars / TokenCounter.CHARS_PER_TOKEN
        return int((word_estimate + char_estimate) / 2)
    
    @staticmethod
    def truncate_to_token_limit(text: str, max_tokens: int) -> Tuple[str, int]:
        """Truncate text to fit within token limit"""
        estimated = TokenCounter.estimate(text)
        if estimated <= max_tokens:
            return text, estimated
        
        # Truncate proportionally
        ratio = max_tokens / estimated
        target_chars = int(len(text) * ratio * 0.9)  # 10% buffer
        truncated = text[:target_chars]
        
        # Try to end at a sentence boundary
        last_period = truncated.rfind('.')
        last_newline = truncated.rfind('\n')
        cut_point = max(last_period, last_newline)
        if cut_point > target_chars * 0.7:
            truncated = truncated[:cut_point + 1]
        
        return truncated, TokenCounter.estimate(truncated)


class DocumentProcessor:
    """Process documents for LLM context injection"""
    
    # Maximum tokens for document context (leaving room for prompt/response)
    MAX_CONTEXT_TOKENS = 6000  # Out of ~8k input limit
    MAX_CHUNK_TOKENS = 1500
    OVERLAP_TOKENS = 100  # Overlap between chunks for continuity
    
    # Supported file extensions
    SUPPORTED_EXTENSIONS = {
        '.pdf': DocumentType.PDF,
        '.docx': DocumentType.DOCX,
        '.doc': DocumentType.DOCX,
        '.txt': DocumentType.TXT,
        '.md': DocumentType.MD,
        '.json': DocumentType.JSON,
    }
    
    @staticmethod
    def get_document_type(filename: str) -> DocumentType:
        """Determine document type from filename"""
        ext = os.path.splitext(filename.lower())[1]
        return DocumentProcessor.SUPPORTED_EXTENSIONS.get(ext, DocumentType.UNKNOWN)
    
    @staticmethod
    def is_processable(filename: str) -> bool:
        """Check if document can be processed"""
        return DocumentProcessor.get_document_type(filename) != DocumentType.UNKNOWN
    
    @staticmethod
    async def extract_text_from_pdf(file_path: str) -> Tuple[str, List[Dict]]:
        """Extract text from PDF file"""
        try:
            # Try PyPDF2 first
            try:
                from PyPDF2 import PdfReader
                reader = PdfReader(file_path)
                pages = []
                page_info = []
                
                for i, page in enumerate(reader.pages):
                    text = page.extract_text() or ""
                    pages.append(text)
                    page_info.append({
                        "page_number": i + 1,
                        "char_count": len(text)
                    })
                
                return "\n\n".join(pages), page_info
                
            except ImportError:
                # Fallback to pdfplumber
                try:
                    import pdfplumber
                    pages = []
                    page_info = []
                    
                    with pdfplumber.open(file_path) as pdf:
                        for i, page in enumerate(pdf.pages):
                            text = page.extract_text() or ""
                            pages.append(text)
                            page_info.append({
                                "page_number": i + 1,
                                "char_count": len(text)
                            })
                    
                    return "\n\n".join(pages), page_info
                    
                except ImportError:
                    logger.warning("No PDF library installed. Install: pip install PyPDF2 or pdfplumber")
                    return "", []
                    
        except Exception as e:
            logger.error(f"PDF extraction error: {e}")
            return "", []
    
    @staticmethod
    async def extract_text_from_docx(file_path: str) -> Tuple[str, List[Dict]]:
        """Extract text from DOCX file"""
        try:
            from docx import Document
            doc = Document(file_path)
            
            paragraphs = []
            for para in doc.paragraphs:
                if para.text.strip():
                    paragraphs.append(para.text)
            
            # Also extract tables
            for table in doc.tables:
                for row in table.rows:
                    row_text = " | ".join(cell.text for cell in row.cells)
                    paragraphs.append(row_text)
            
            text = "\n".join(paragraphs)
            return text, [{"section": "main", "char_count": len(text)}]
            
        except ImportError:
            logger.warning("python-docx not installed. Install: pip install python-docx")
            return "", []
        except Exception as e:
            logger.error(f"DOCX extraction error: {e}")
            return "", []
    
    @staticmethod
    async def extract_text_from_file(file_path: str, doc_type: DocumentType) -> Tuple[str, List[Dict]]:
        """Extract text from file based on type"""
        
        if doc_type == DocumentType.PDF:
            return await DocumentProcessor.extract_text_from_pdf(file_path)
        
        elif doc_type == DocumentType.DOCX:
            return await DocumentProcessor.extract_text_from_docx(file_path)
        
        elif doc_type in [DocumentType.TXT, DocumentType.MD, DocumentType.JSON]:
            try:
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                    text = f.read()
                return text, [{"section": "main", "char_count": len(text)}]
            except Exception as e:
                logger.error(f"Text file read error: {e}")
                return "", []
        
        else:
            return "", []
    
    @staticmethod
    def create_chunks(text: str, doc_id: str, filename: str, 
                      max_chunk_tokens: int = None,
                      overlap_tokens: int = None) -> List[DocumentChunk]:
        """Split text into overlapping chunks"""
        
        max_chunk_tokens = max_chunk_tokens or DocumentProcessor.MAX_CHUNK_TOKENS
        overlap_tokens = overlap_tokens or DocumentProcessor.OVERLAP_TOKENS
        
        if not text.strip():
            return []
        
        chunks = []
        
        # Split into paragraphs first
        paragraphs = re.split(r'\n\s*\n', text)
        
        current_chunk = []
        current_tokens = 0
        chunk_id = 0
        
        for para in paragraphs:
            para = para.strip()
            if not para:
                continue
            
            para_tokens = TokenCounter.estimate(para)
            
            # If single paragraph exceeds limit, split it
            if para_tokens > max_chunk_tokens:
                # Save current chunk first
                if current_chunk:
                    chunk_text = "\n\n".join(current_chunk)
                    chunks.append(DocumentChunk(
                        chunk_id=chunk_id,
                        content=chunk_text,
                        token_count=TokenCounter.estimate(chunk_text),
                        source_document=filename
                    ))
                    chunk_id += 1
                    current_chunk = []
                    current_tokens = 0
                
                # Split large paragraph by sentences
                sentences = re.split(r'(?<=[.!?])\s+', para)
                for sent in sentences:
                    sent_tokens = TokenCounter.estimate(sent)
                    if current_tokens + sent_tokens > max_chunk_tokens:
                        if current_chunk:
                            chunk_text = " ".join(current_chunk)
                            chunks.append(DocumentChunk(
                                chunk_id=chunk_id,
                                content=chunk_text,
                                token_count=TokenCounter.estimate(chunk_text),
                                source_document=filename
                            ))
                            chunk_id += 1
                        current_chunk = [sent]
                        current_tokens = sent_tokens
                    else:
                        current_chunk.append(sent)
                        current_tokens += sent_tokens
            
            # Normal paragraph handling
            elif current_tokens + para_tokens > max_chunk_tokens:
                # Save current chunk
                if current_chunk:
                    chunk_text = "\n\n".join(current_chunk)
                    chunks.append(DocumentChunk(
                        chunk_id=chunk_id,
                        content=chunk_text,
                        token_count=TokenCounter.estimate(chunk_text),
                        source_document=filename
                    ))
                    chunk_id += 1
                
                # Start new chunk with overlap
                current_chunk = [para]
                current_tokens = para_tokens
            else:
                current_chunk.append(para)
                current_tokens += para_tokens
        
        # Save final chunk
        if current_chunk:
            chunk_text = "\n\n".join(current_chunk)
            chunks.append(DocumentChunk(
                chunk_id=chunk_id,
                content=chunk_text,
                token_count=TokenCounter.estimate(chunk_text),
                source_document=filename
            ))
        
        return chunks
    
    @staticmethod
    async def process_document(
        file_path: str,
        document_id: str,
        skip: bool = False,
        skip_reason: str = None,
        max_tokens: int = None
    ) -> ProcessedDocument:
        """Process a single document"""
        
        filename = os.path.basename(file_path)
        doc_type = DocumentProcessor.get_document_type(filename)
        
        # Handle skip flag
        if skip:
            logger.info(f"⏭️ Skipping document: {filename} (reason: {skip_reason or 'user skipped'})")
            return ProcessedDocument(
                document_id=document_id,
                filename=filename,
                document_type=doc_type,
                total_tokens=0,
                chunks=[],
                skipped=True,
                skip_reason=skip_reason or "User skipped"
            )
        
        # Check if processable
        if not DocumentProcessor.is_processable(filename):
            logger.warning(f"⚠️ Unsupported document type: {filename}")
            return ProcessedDocument(
                document_id=document_id,
                filename=filename,
                document_type=DocumentType.UNKNOWN,
                total_tokens=0,
                chunks=[],
                skipped=True,
                skip_reason=f"Unsupported file type: {doc_type.value}"
            )
        
        # Extract text
        text, page_info = await DocumentProcessor.extract_text_from_file(file_path, doc_type)
        
        if not text.strip():
            logger.warning(f"⚠️ No text extracted from: {filename}")
            return ProcessedDocument(
                document_id=document_id,
                filename=filename,
                document_type=doc_type,
                total_tokens=0,
                chunks=[],
                skipped=True,
                skip_reason="No extractable text"
            )
        
        # Create chunks
        chunks = DocumentProcessor.create_chunks(
            text=text,
            doc_id=document_id,
            filename=filename,
            max_chunk_tokens=max_tokens or DocumentProcessor.MAX_CHUNK_TOKENS
        )
        
        total_tokens = sum(chunk.token_count for chunk in chunks)
        
        logger.info(f"📄 Processed {filename}: {len(chunks)} chunks, ~{total_tokens} tokens")
        
        return ProcessedDocument(
            document_id=document_id,
            filename=filename,
            document_type=doc_type,
            total_tokens=total_tokens,
            chunks=chunks,
            skipped=False,
            metadata={
                "page_count": len(page_info),
                "char_count": len(text),
                "extraction_method": doc_type.value
            }
        )
    
    @staticmethod
    def build_context_from_documents(
        documents: List[ProcessedDocument],
        max_total_tokens: int = None
    ) -> Tuple[str, int, Dict[str, Any]]:
        """Build combined context string from processed documents"""
        
        max_total_tokens = max_total_tokens or DocumentProcessor.MAX_CONTEXT_TOKENS
        
        context_parts = []
        total_tokens = 0
        stats = {
            "documents_processed": 0,
            "documents_skipped": 0,
            "chunks_used": 0,
            "tokens_used": 0,
            "truncated": False
        }
        
        for doc in documents:
            if doc.skipped:
                stats["documents_skipped"] += 1
                continue
            
            # Add document header
            doc_header = f"\n\n### Document: {doc.filename} ###\n"
            header_tokens = TokenCounter.estimate(doc_header)
            
            if total_tokens + header_tokens > max_total_tokens:
                stats["truncated"] = True
                break
            
            context_parts.append(doc_header)
            total_tokens += header_tokens
            
            # Add chunks
            for chunk in doc.chunks:
                chunk_text = f"\n[Chunk {chunk.chunk_id + 1}]\n{chunk.content}\n"
                chunk_tokens = TokenCounter.estimate(chunk_text)
                
                if total_tokens + chunk_tokens > max_total_tokens:
                    stats["truncated"] = True
                    break
                
                context_parts.append(chunk_text)
                total_tokens += chunk_tokens
                stats["chunks_used"] += 1
            
            stats["documents_processed"] += 1
        
        stats["tokens_used"] = total_tokens
        
        context = "".join(context_parts)
        
        logger.info(f"📚 Built context: {stats['documents_processed']} docs, "
                   f"{stats['chunks_used']} chunks, ~{total_tokens} tokens")
        
        return context, total_tokens, stats


class AttachmentConfig:
    """Configuration for attachment processing"""
    
    def __init__(
        self,
        attachment_id: str,
        filename: str,
        file_path: Optional[str] = None,
        download_url: Optional[str] = None,
        skip: bool = False,
        skip_reason: Optional[str] = None,
        priority: int = 0,  # Higher = processed first
        max_tokens: Optional[int] = None
    ):
        self.attachment_id = attachment_id
        self.filename = filename
        self.file_path = file_path
        self.download_url = download_url
        self.skip = skip
        self.skip_reason = skip_reason
        self.priority = priority
        self.max_tokens = max_tokens
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "attachment_id": self.attachment_id,
            "filename": self.filename,
            "file_path": self.file_path,
            "download_url": self.download_url,
            "skip": self.skip,
            "skip_reason": self.skip_reason,
            "priority": self.priority,
            "max_tokens": self.max_tokens
        }