utils/pdf_processor.py

"""
PDF processing and text extraction with chunking.
"""
import logging
from pathlib import Path
from typing import List, Optional
import hashlib
import tiktoken
from pypdf import PdfReader

from utils.schemas import PaperChunk, Paper

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class PDFProcessor:
    """Process PDFs and extract text with intelligent chunking."""

    def __init__(
        self,
        chunk_size: int = 500,
        chunk_overlap: int = 50,
        encoding_name: str = "cl100k_base"
    ):
        """
        Initialize PDF processor.

        Args:
            chunk_size: Target chunk size in tokens
            chunk_overlap: Overlap between chunks in tokens
            encoding_name: Tiktoken encoding name
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.encoding = tiktoken.get_encoding(encoding_name)

    def extract_text(self, pdf_path: Path) -> Optional[str]:
        """
        Extract text from PDF.

        Args:
            pdf_path: Path to PDF file

        Returns:
            Extracted text or None if extraction fails
        """
        try:
            reader = PdfReader(str(pdf_path))
            text_parts = []

            for page_num, page in enumerate(reader.pages, start=1):
                try:
                    text = page.extract_text()
                    if text.strip():
                        text_parts.append(f"[Page {page_num}]\n{text}")
                except Exception as e:
                    logger.warning(f"Failed to extract text from page {page_num}: {str(e)}")
                    continue

            if not text_parts:
                logger.error(f"No text extracted from {pdf_path}")
                return None

            full_text = "\n\n".join(text_parts)
            logger.info(f"Extracted {len(full_text)} characters from {pdf_path.name}")
            return full_text

        except Exception as e:
            logger.error(f"Error reading PDF {pdf_path}: {str(e)}")
            return None

    def _generate_chunk_id(self, paper_id: str, chunk_index: int) -> str:
        """Generate unique chunk ID."""
        content = f"{paper_id}_{chunk_index}"
        return hashlib.md5(content.encode()).hexdigest()

    def chunk_text(
        self,
        text: str,
        paper: Paper
    ) -> List[PaperChunk]:
        """
        Chunk text into overlapping segments.

        Args:
            text: Full text to chunk
            paper: Paper metadata

        Returns:
            List of PaperChunk objects
        """
        chunks = []
        tokens = self.encoding.encode(text)

        # Extract page information from text
        page_markers = []
        lines = text.split('\n')
        current_char = 0
        for line in lines:
            if line.startswith('[Page ') and line.endswith(']'):
                try:
                    page_num = int(line[6:-1])
                    page_markers.append((current_char, page_num))
                except ValueError:
                    pass
            current_char += len(line) + 1

        chunk_index = 0
        start_idx = 0

        while start_idx < len(tokens):
            # Calculate end index
            end_idx = min(start_idx + self.chunk_size, len(tokens))

            # Get chunk tokens and decode
            chunk_tokens = tokens[start_idx:end_idx]
            chunk_text = self.encoding.decode(chunk_tokens)

            # Determine page number
            chunk_start_char = len(self.encoding.decode(tokens[:start_idx]))
            page_number = self._get_page_number(chunk_start_char, page_markers)

            # Extract section if possible (simple heuristic)
            section = self._extract_section(chunk_text)

            # Create metadata with type validation
            try:
                # Ensure authors is a list of strings
                authors_metadata = paper.authors
                if not isinstance(authors_metadata, list):
                    logger.warning(f"Paper {paper.arxiv_id} has invalid authors type: {type(authors_metadata)}, converting to list")
                    authors_metadata = [str(authors_metadata)] if authors_metadata else []

                # Ensure title is a string
                title_metadata = str(paper.title) if paper.title else ""

                metadata = {
                    "title": title_metadata,
                    "authors": authors_metadata,
                    "chunk_index": chunk_index,
                    "token_count": len(chunk_tokens)
                }
            except Exception as e:
                logger.warning(f"Error creating metadata for chunk {chunk_index}: {str(e)}, using fallback")
                metadata = {
                    "title": str(paper.title) if hasattr(paper, 'title') else "",
                    "authors": [],
                    "chunk_index": chunk_index,
                    "token_count": len(chunk_tokens)
                }

            # Create chunk with validated data
            try:
                chunk = PaperChunk(
                    chunk_id=self._generate_chunk_id(paper.arxiv_id, chunk_index),
                    paper_id=paper.arxiv_id,
                    content=chunk_text.strip(),
                    section=section,
                    page_number=page_number,
                    arxiv_url=str(paper.pdf_url) if paper.pdf_url else "",
                    metadata=metadata
                )
                chunks.append(chunk)
            except Exception as e:
                logger.error(f"Error creating chunk {chunk_index} for paper {paper.arxiv_id}: {str(e)}")
                # Continue processing other chunks rather than failing completely
                continue

            # Move to next chunk with overlap
            start_idx += self.chunk_size - self.chunk_overlap
            chunk_index += 1

        logger.info(f"Created {len(chunks)} chunks for paper {paper.arxiv_id}")
        return chunks

    def _get_page_number(
        self,
        char_position: int,
        page_markers: List[tuple]
    ) -> Optional[int]:
        """Determine page number for character position."""
        if not page_markers:
            return None

        for i, (marker_pos, page_num) in enumerate(page_markers):
            if char_position < marker_pos:
                return page_markers[i - 1][1] if i > 0 else None
        return page_markers[-1][1]

    def _extract_section(self, text: str) -> Optional[str]:
        """
        Extract section name from chunk (simple heuristic).

        Looks for common section headers.
        """
        section_keywords = [
            'abstract', 'introduction', 'related work', 'methodology',
            'method', 'experiments', 'results', 'discussion',
            'conclusion', 'references', 'appendix'
        ]

        lines = text.split('\n')[:5]  # Check first 5 lines
        for line in lines:
            line_lower = line.lower().strip()
            for keyword in section_keywords:
                if keyword in line_lower and len(line.split()) < 10:
                    return line.strip()
        return None

    def process_paper(
        self,
        pdf_path: Path,
        paper: Paper
    ) -> List[PaperChunk]:
        """
        Process a paper PDF into chunks.

        Args:
            pdf_path: Path to PDF file
            paper: Paper metadata

        Returns:
            List of PaperChunk objects
        """
        # Extract text
        text = self.extract_text(pdf_path)
        if not text:
            logger.error(f"Failed to extract text from {pdf_path}")
            return []

        # Chunk text
        chunks = self.chunk_text(text, paper)
        return chunks