src/preprocessing/pdf_parser.py

"""
PDF Parser Module for Document Processing.
Extracts text, tables, and metadata from PDF documents.
"""

import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Union

import pandas as pd

from ..utils import get_logger, LoggerMixin

logger = get_logger(__name__)


@dataclass
class PageContent:
    """Content extracted from a single PDF page."""
    
    page_number: int
    text: str
    tables: List[pd.DataFrame] = field(default_factory=list)
    metadata: Dict = field(default_factory=dict)
    
    def to_dict(self) -> Dict:
        return {
            "page_number": self.page_number,
            "text": self.text,
            "tables": [t.to_dict() for t in self.tables],
            "metadata": self.metadata
        }


@dataclass
class DocumentContent:
    """Complete content extracted from a PDF document."""
    
    source_file: str
    total_pages: int
    pages: List[PageContent]
    metadata: Dict = field(default_factory=dict)
    
    @property
    def full_text(self) -> str:
        """Get concatenated text from all pages."""
        return "\n\n".join(page.text for page in self.pages)
    
    @property
    def all_tables(self) -> List[pd.DataFrame]:
        """Get all tables from all pages."""
        tables = []
        for page in self.pages:
            tables.extend(page.tables)
        return tables
    
    def to_dict(self) -> Dict:
        return {
            "source_file": self.source_file,
            "total_pages": self.total_pages,
            "pages": [p.to_dict() for p in self.pages],
            "metadata": self.metadata
        }


class TableExtractor(LoggerMixin):
    """
    Extract tables from PDF documents.
    
    Uses Tabula and Camelot for table detection and extraction.
    """
    
    def __init__(self, method: str = "tabula"):
        """
        Initialize table extractor.
        
        Args:
            method: Extraction method - "tabula" or "camelot"
        """
        self.method = method
        self._check_dependencies()
    
    def _check_dependencies(self):
        """Check if required libraries are available."""
        if self.method == "tabula":
            try:
                import tabula
                self.tabula = tabula
                self.logger.debug("Tabula initialized successfully")
            except ImportError:
                self.tabula = None
                self.logger.warning("tabula-py not installed")
        elif self.method == "camelot":
            try:
                import camelot
                self.camelot = camelot
                self.logger.debug("Camelot initialized successfully")
            except ImportError:
                self.camelot = None
                self.logger.warning("camelot-py not installed")
    
    def extract_tables(
        self,
        pdf_path: Union[str, Path],
        pages: str = "all"
    ) -> Dict[int, List[pd.DataFrame]]:
        """
        Extract tables from PDF.
        
        Args:
            pdf_path: Path to PDF file
            pages: Pages to extract - "all" or specific pages like "1,2,3"
            
        Returns:
            Dictionary mapping page numbers to list of DataFrames
        """
        pdf_path = Path(pdf_path)
        self.logger.info(f"Extracting tables from: {pdf_path.name}")
        
        if self.method == "tabula" and self.tabula:
            return self._extract_with_tabula(pdf_path, pages)
        elif self.method == "camelot" and self.camelot:
            return self._extract_with_camelot(pdf_path, pages)
        else:
            self.logger.warning("No table extraction library available")
            return {}
    
    def _extract_with_tabula(
        self,
        pdf_path: Path,
        pages: str
    ) -> Dict[int, List[pd.DataFrame]]:
        """Extract tables using Tabula."""
        try:
            # Read all tables
            tables = self.tabula.read_pdf(
                str(pdf_path),
                pages=pages,
                multiple_tables=True,
                pandas_options={'header': None}
            )
            
            # Group by page (tabula returns flat list)
            # For simplicity, assume sequential pages
            result = {}
            for i, table in enumerate(tables):
                if not table.empty:
                    page_num = i + 1
                    if page_num not in result:
                        result[page_num] = []
                    result[page_num].append(table)
            
            self.logger.debug(f"Extracted {len(tables)} tables")
            return result
            
        except Exception as e:
            self.logger.error(f"Tabula extraction failed: {e}")
            return {}
    
    def _extract_with_camelot(
        self,
        pdf_path: Path,
        pages: str
    ) -> Dict[int, List[pd.DataFrame]]:
        """Extract tables using Camelot."""
        try:
            tables = self.camelot.read_pdf(
                str(pdf_path),
                pages=pages if pages != "all" else "1-end",
                flavor='lattice'
            )
            
            result = {}
            for table in tables:
                page_num = table.page
                if page_num not in result:
                    result[page_num] = []
                result[page_num].append(table.df)
            
            self.logger.debug(f"Extracted {len(tables)} tables")
            return result
            
        except Exception as e:
            self.logger.error(f"Camelot extraction failed: {e}")
            return {}
    
    def table_to_text(self, table: pd.DataFrame) -> str:
        """
        Convert a DataFrame table to structured text.
        
        Args:
            table: Pandas DataFrame
            
        Returns:
            Formatted text representation
        """
        if table.empty:
            return ""
        
        # Clean up the table
        table = table.fillna("")
        
        # Convert to markdown-like format
        lines = []
        
        # Header (first row if it looks like headers)
        if len(table) > 1:
            header = " | ".join(str(cell) for cell in table.iloc[0])
            lines.append(header)
            lines.append("-" * len(header))
            start_row = 1
        else:
            start_row = 0
        
        # Data rows
        for _, row in table.iloc[start_row:].iterrows():
            line = " | ".join(str(cell) for cell in row)
            lines.append(line)
        
        return "\n".join(lines)


class ImageExtractor(LoggerMixin):
    """
    Extract images from PDF documents with optional captioning.
    
    Uses PyMuPDF for extraction and BLIP/CLIP for image understanding.
    """
    
    def __init__(
        self,
        output_dir: Optional[Union[str, Path]] = None,
        min_size: int = 100,
        generate_captions: bool = True
    ):
        """
        Initialize image extractor.
        
        Args:
            output_dir: Directory to save extracted images
            min_size: Minimum image dimension to extract
            generate_captions: Whether to generate captions using vision model
        """
        self.output_dir = Path(output_dir) if output_dir else Path("extracted_images")
        self.min_size = min_size
        self.generate_captions = generate_captions
        self.caption_model = None
        self.caption_processor = None
        
        self.output_dir.mkdir(parents=True, exist_ok=True)
    
    def _load_caption_model(self):
        """Lazy load BLIP captioning model."""
        if self.caption_model is None and self.generate_captions:
            try:
                from transformers import BlipProcessor, BlipForConditionalGeneration
                import torch
                
                self.logger.info("Loading BLIP captioning model...")
                self.caption_processor = BlipProcessor.from_pretrained(
                    "Salesforce/blip-image-captioning-base"
                )
                self.caption_model = BlipForConditionalGeneration.from_pretrained(
                    "Salesforce/blip-image-captioning-base"
                )
                
                if torch.cuda.is_available():
                    self.caption_model = self.caption_model.cuda()
                    
                self.logger.info("BLIP model loaded successfully")
            except Exception as e:
                self.logger.warning(f"BLIP not available: {e}")
                self.generate_captions = False
    
    def extract_images(
        self,
        pdf_path: Union[str, Path],
        pages: Optional[List[int]] = None
    ) -> List[Dict]:
        """
        Extract images from PDF.
        
        Args:
            pdf_path: Path to PDF file
            pages: Specific pages to extract from (1-indexed), None for all
            
        Returns:
            List of dicts with image info: path, page, caption, etc.
        """
        try:
            import fitz
        except ImportError:
            self.logger.error("PyMuPDF required for image extraction")
            return []
        
        pdf_path = Path(pdf_path)
        self.logger.info(f"Extracting images from: {pdf_path.name}")
        
        extracted = []
        doc = fitz.open(str(pdf_path))
        
        for page_num in range(len(doc)):
            if pages and (page_num + 1) not in pages:
                continue
                
            page = doc[page_num]
            image_list = page.get_images(full=True)
            
            for img_index, img in enumerate(image_list):
                try:
                    xref = img[0]
                    base_image = doc.extract_image(xref)
                    
                    if base_image:
                        image_bytes = base_image["image"]
                        image_ext = base_image["ext"]
                        width = base_image.get("width", 0)
                        height = base_image.get("height", 0)
                        
                        # Skip small images (likely icons/bullets)
                        if width < self.min_size or height < self.min_size:
                            continue
                        
                        # Save image
                        image_name = f"{pdf_path.stem}_p{page_num+1}_img{img_index+1}.{image_ext}"
                        image_path = self.output_dir / image_name
                        
                        with open(image_path, "wb") as f:
                            f.write(image_bytes)
                        
                        # Generate caption if enabled
                        caption = ""
                        if self.generate_captions:
                            caption = self._generate_caption(image_path)
                        
                        extracted.append({
                            "path": str(image_path),
                            "page": page_num + 1,
                            "index": img_index + 1,
                            "width": width,
                            "height": height,
                            "caption": caption,
                            "format": image_ext
                        })
                        
                except Exception as e:
                    self.logger.debug(f"Failed to extract image: {e}")
                    continue
        
        doc.close()
        self.logger.info(f"Extracted {len(extracted)} images")
        return extracted
    
    def _generate_caption(self, image_path: Path) -> str:
        """Generate caption for an image using BLIP."""
        self._load_caption_model()
        
        if not self.caption_model:
            return ""
        
        try:
            from PIL import Image
            import torch
            
            image = Image.open(image_path).convert("RGB")
            inputs = self.caption_processor(images=image, return_tensors="pt")
            
            if torch.cuda.is_available():
                inputs = {k: v.cuda() for k, v in inputs.items()}
            
            with torch.no_grad():
                output = self.caption_model.generate(**inputs, max_new_tokens=50)
            
            caption = self.caption_processor.decode(output[0], skip_special_tokens=True)
            return caption
            
        except Exception as e:
            self.logger.debug(f"Caption generation failed: {e}")
            return ""
    
    def image_to_text(self, image_info: Dict) -> str:
        """Convert image info to text for indexing."""
        parts = [f"[Image on page {image_info['page']}]"]
        if image_info.get("caption"):
            parts.append(f"Caption: {image_info['caption']}")
        parts.append(f"Size: {image_info.get('width', '?')}x{image_info.get('height', '?')}")
        return " ".join(parts)


class OCRProcessor(LoggerMixin):
    """
    OCR processor for scanned PDFs.
    
    Uses Tesseract for text extraction from images.
    """
    
    def __init__(self, language: str = "eng"):
        """
        Initialize OCR processor.
        
        Args:
            language: Tesseract language code
        """
        self.language = language
        self.tesseract_available = self._check_tesseract()
    
    def _check_tesseract(self) -> bool:
        """Check if Tesseract is available."""
        try:
            import pytesseract
            pytesseract.get_tesseract_version()
            self.logger.info("Tesseract OCR available")
            return True
        except Exception:
            self.logger.warning("Tesseract not available. Install with: pip install pytesseract")
            return False
    
    def ocr_pdf(
        self,
        pdf_path: Union[str, Path],
        dpi: int = 300
    ) -> List[Dict]:
        """
        Perform OCR on a scanned PDF.
        
        Args:
            pdf_path: Path to PDF file
            dpi: Resolution for rendering pages
            
        Returns:
            List of dicts with page number and extracted text
        """
        if not self.tesseract_available:
            return []
        
        try:
            import fitz
            import pytesseract
            from PIL import Image
            import io
        except ImportError as e:
            self.logger.error(f"Missing dependency: {e}")
            return []
        
        pdf_path = Path(pdf_path)
        self.logger.info(f"OCR processing: {pdf_path.name}")
        
        results = []
        doc = fitz.open(str(pdf_path))
        
        for page_num in range(len(doc)):
            page = doc[page_num]
            
            # Render page to image
            mat = fitz.Matrix(dpi / 72, dpi / 72)
            pix = page.get_pixmap(matrix=mat)
            
            # Convert to PIL Image
            img_data = pix.tobytes("png")
            image = Image.open(io.BytesIO(img_data))
            
            # Perform OCR
            text = pytesseract.image_to_string(image, lang=self.language)
            
            results.append({
                "page": page_num + 1,
                "text": text.strip(),
                "confidence": self._get_confidence(image)
            })
        
        doc.close()
        self.logger.info(f"OCR completed for {len(results)} pages")
        return results
    
    def _get_confidence(self, image) -> float:
        """Get OCR confidence score."""
        try:
            import pytesseract
            data = pytesseract.image_to_data(image, output_type=pytesseract.Output.DICT)
            confidences = [int(c) for c in data.get("conf", []) if c != "-1" and int(c) > 0]
            return sum(confidences) / len(confidences) if confidences else 0.0
        except:
            return 0.0



class PDFParser(LoggerMixin):
    """
    PDF Parser for extracting text and structure from documents.
    
    Uses PyMuPDF (fitz) for fast and accurate text extraction
    with layout preservation.
    """
    
    def __init__(
        self,
        extract_tables: bool = True,
        table_extractor: Optional[TableExtractor] = None
    ):
        """
        Initialize PDF parser.
        
        Args:
            extract_tables: Whether to extract tables
            table_extractor: Custom table extractor instance
        """
        self.extract_tables = extract_tables
        self.table_extractor = table_extractor or TableExtractor()
        
        # Import fitz (PyMuPDF)
        try:
            import fitz
            self.fitz = fitz
            self.logger.info("PyMuPDF initialized successfully")
        except ImportError:
            self.fitz = None
            self.logger.error("PyMuPDF not installed. Install with: pip install PyMuPDF")
    
    def parse(
        self,
        pdf_path: Union[str, Path],
        start_page: int = 0,
        end_page: Optional[int] = None
    ) -> DocumentContent:
        """
        Parse a PDF document and extract all content.
        
        Args:
            pdf_path: Path to PDF file
            start_page: First page to process (0-indexed)
            end_page: Last page to process (exclusive, None for all)
            
        Returns:
            DocumentContent with extracted text and tables
        """
        pdf_path = Path(pdf_path)
        self.logger.info(f"Parsing PDF: {pdf_path.name}")
        
        if self.fitz is None:
            self.logger.error("PyMuPDF not available")
            return DocumentContent(
                source_file=str(pdf_path),
                total_pages=0,
                pages=[],
                metadata={}
            )
        
        try:
            doc = self.fitz.open(str(pdf_path))
            total_pages = len(doc)
            
            # Determine page range
            if end_page is None:
                end_page = total_pages
            end_page = min(end_page, total_pages)
            
            # Extract document metadata
            doc_metadata = self._extract_metadata(doc)
            
            # Extract tables if enabled
            tables_by_page = {}
            if self.extract_tables:
                page_spec = f"{start_page + 1}-{end_page}"
                tables_by_page = self.table_extractor.extract_tables(pdf_path, page_spec)
            
            # Process each page
            pages = []
            for page_num in range(start_page, end_page):
                page = doc[page_num]
                
                # Extract text with layout preservation
                text = self._extract_page_text(page)
                
                # Get tables for this page
                page_tables = tables_by_page.get(page_num + 1, [])
                
                # Page metadata
                page_metadata = {
                    "width": page.rect.width,
                    "height": page.rect.height,
                    "rotation": page.rotation
                }
                
                page_content = PageContent(
                    page_number=page_num + 1,
                    text=text,
                    tables=page_tables,
                    metadata=page_metadata
                )
                pages.append(page_content)
            
            doc.close()
            
            self.logger.info(f"Extracted {len(pages)} pages from PDF")
            
            return DocumentContent(
                source_file=str(pdf_path),
                total_pages=total_pages,
                pages=pages,
                metadata=doc_metadata
            )
            
        except Exception as e:
            self.logger.error(f"Error parsing PDF: {e}")
            return DocumentContent(
                source_file=str(pdf_path),
                total_pages=0,
                pages=[],
                metadata={"error": str(e)}
            )
    
    def _extract_metadata(self, doc) -> Dict:
        """Extract document metadata."""
        metadata = doc.metadata
        return {
            "title": metadata.get("title", ""),
            "author": metadata.get("author", ""),
            "subject": metadata.get("subject", ""),
            "creator": metadata.get("creator", ""),
            "creation_date": metadata.get("creationDate", ""),
            "modification_date": metadata.get("modDate", ""),
            "keywords": metadata.get("keywords", "")
        }
    
    def _extract_page_text(self, page) -> str:
        """
        Extract text from a page with layout preservation.
        
        Args:
            page: PyMuPDF page object
            
        Returns:
            Extracted text
        """
        # Get text blocks for layout-aware extraction
        blocks = page.get_text("blocks")
        
        # Sort blocks by position (top to bottom, left to right)
        blocks = sorted(blocks, key=lambda b: (b[1], b[0]))
        
        # Extract text from blocks
        text_parts = []
        for block in blocks:
            if block[6] == 0:  # Text block (not image)
                text = block[4].strip()
                if text:
                    text_parts.append(text)
        
        text = "\n\n".join(text_parts)
        
        # Clean up text
        text = self._clean_text(text)
        
        return text
    
    def _clean_text(self, text: str) -> str:
        """Clean and normalize extracted text."""
        # Fix hyphenation at line breaks
        text = re.sub(r'-\s*\n\s*', '', text)
        
        # Normalize whitespace
        text = re.sub(r'[ \t]+', ' ', text)
        text = re.sub(r'\n{3,}', '\n\n', text)
        
        # Remove page numbers (common patterns)
        text = re.sub(r'\n\s*\d+\s*\n', '\n', text)
        
        return text.strip()
    
    def extract_sections(self, doc_content: DocumentContent) -> List[Dict]:
        """
        Extract sections/headings from document.
        
        Args:
            doc_content: Parsed document content
            
        Returns:
            List of sections with titles and content
        """
        sections = []
        current_section = {"title": "Introduction", "content": [], "page": 1}
        
        # Simple heuristic: lines that are short and possibly uppercase are headings
        for page in doc_content.pages:
            lines = page.text.split('\n')
            
            for line in lines:
                line = line.strip()
                if not line:
                    continue
                
                # Check if line looks like a heading
                is_heading = (
                    len(line) < 100 and
                    (line.isupper() or 
                     line.istitle() or
                     re.match(r'^\d+\.?\s+[A-Z]', line))
                )
                
                if is_heading:
                    # Save current section
                    if current_section["content"]:
                        current_section["content"] = "\n".join(current_section["content"])
                        sections.append(current_section)
                    
                    # Start new section
                    current_section = {
                        "title": line,
                        "content": [],
                        "page": page.page_number
                    }
                else:
                    current_section["content"].append(line)
        
        # Don't forget last section
        if current_section["content"]:
            current_section["content"] = "\n".join(current_section["content"])
            sections.append(current_section)
        
        return sections


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="PDF Parser Test")
    parser.add_argument("--test", action="store_true", help="Run test mode")
    parser.add_argument("--pdf", type=str, help="PDF path to process")
    args = parser.parse_args()
    
    if args.test:
        print("PDF Parser initialized successfully!")
        pdf_parser = PDFParser()
        print(f"PyMuPDF available: {pdf_parser.fitz is not None}")
        print(f"Table extraction enabled: {pdf_parser.extract_tables}")
    
    if args.pdf:
        pdf_parser = PDFParser()
        result = pdf_parser.parse(args.pdf)
        print(f"\nDocument: {result.source_file}")
        print(f"Total pages: {result.total_pages}")
        print(f"Metadata: {result.metadata}")
        print(f"\nFirst page text (first 500 chars):\n{result.pages[0].text[:500]}...")