lib/parser/ntr_fileparser/parsers/specific_parsers/pdf_parser.py

"""
Модуль с парсером для PDF документов.
"""

import io
import logging
import os
from typing import BinaryIO

import fitz  # PyMuPDF

from ...data_classes import ParsedDocument, ParsedMeta
from ..abstract_parser import AbstractParser
from ..file_types import FileType
from .pdf.formula_parser import PDFFormulaParser
from .pdf.image_parser import PDFImageParser
from .pdf.meta_parser import PDFMetaParser
from .pdf.paragraph_parser import PDFParagraphParser
from .pdf.table_parser import PDFTableParser

logger = logging.getLogger(__name__)


class PDFParser(AbstractParser):
    """
    Парсер для PDF документов.
    
    Использует PyMuPDF (fitz) для извлечения текста, изображений, таблиц,
    формул и метаданных из документа.
    """
    
    def __init__(self):
        """
        Инициализирует PDF парсер и его компоненты.
        """
        super().__init__(FileType.PDF)
        self.meta_parser = PDFMetaParser()
        self.paragraph_parser = PDFParagraphParser()
        self.table_parser = PDFTableParser()
        self.image_parser = PDFImageParser()
        self.formula_parser = PDFFormulaParser()
    
    def parse_by_path(self, file_path: str) -> ParsedDocument:
        """
        Парсит PDF документ по пути к файлу и возвращает его структурное представление.
        
        Args:
            file_path (str): Путь к PDF файлу для парсинга.
            
        Returns:
            ParsedDocument: Структурное представление документа.
            
        Raises:
            ValueError: Если файл не существует или не может быть прочитан.
        """
        logger.debug(f"Parsing PDF file: {file_path}")
        
        if not os.path.exists(file_path):
            raise ValueError(f"File not found: {file_path}")
        
        try:
            # Открываем PDF с помощью PyMuPDF
            pdf_doc = fitz.open(file_path)
            filename = os.path.basename(file_path)
            
            return self._parse_document(pdf_doc, filename, file_path)
            
        except Exception as e:
            logger.error(f"Failed to open PDF file: {e}")
            raise ValueError(f"Cannot open PDF file: {str(e)}")
    
    def parse(self, file: BinaryIO, file_type: FileType | str | None = None) -> ParsedDocument:
        """
        Парсит PDF документ из объекта файла и возвращает его структурное представление.
        
        Args:
            file (BinaryIO): Объект файла для парсинга.
            file_type: Тип файла, если известен.
                Может быть объектом FileType или строкой с расширением (".pdf").
            
        Returns:
            ParsedDocument: Структурное представление документа.
            
        Raises:
            ValueError: Если файл не может быть прочитан или распарсен.
        """
        logger.debug("Parsing PDF from file object")
        
        # Проверяем соответствие типа файла
        if file_type and isinstance(file_type, FileType) and file_type != FileType.PDF:
            logger.warning(
                f"Provided file_type {file_type} doesn't match parser type {FileType.PDF}"
            )
        
        try:
            # Читаем содержимое файла в память
            content = file.read()
            
            # Открываем PDF из потока с помощью PyMuPDF
            pdf_stream = io.BytesIO(content)
            pdf_doc = fitz.open(stream=pdf_stream, filetype="pdf")
            
            return self._parse_document(pdf_doc, "unknown.pdf", None)
            
        except Exception as e:
            logger.error(f"Failed to parse PDF from stream: {e}")
            raise ValueError(f"Cannot parse PDF content: {str(e)}")
    
    def _parse_document(
        self,
        pdf_doc: fitz.Document,
        filename: str,
        filepath: str | None,
    ) -> ParsedDocument:
        """
        Внутренний метод для парсинга открытого PDF документа.
        
        Args:
            pdf_doc (fitz.Document): Открытый PDF документ.
            filename (str): Имя файла для документа.
            filepath (str | None): Путь к файлу (или None, если из объекта).
            
        Returns:
            ParsedDocument: Структурное представление документа.
            
        Raises:
            ValueError: Если содержимое не может быть распарсено.
        """
        # Создание базового документа
        doc = ParsedDocument(name=filename, type=FileType.PDF)
        
        try:
            # Извлечение метаданных
            meta_dict = self.meta_parser.parse(pdf_doc, filepath)
            
            # Преобразуем словарь метаданных в объект ParsedMeta
            meta = ParsedMeta()
            if 'author' in meta_dict:
                meta.owner = meta_dict['author']
            if 'creation_date' in meta_dict:
                meta.date = meta_dict['creation_date']
            if filepath:
                meta.source = filepath
            
            # Сохраняем остальные метаданные в поле note
            meta.note = meta_dict
            
            doc.meta = meta
            logger.debug("Parsed metadata")
            
            # Последовательный вызов парсеров
            try:
                # Парсим таблицы
                doc.tables.extend(self.table_parser.parse(pdf_doc))
                logger.debug(f"Parsed {len(doc.tables)} tables")
                
                # Парсим изображения
                doc.images.extend(self.image_parser.parse(pdf_doc))
                logger.debug(f"Parsed {len(doc.images)} images")
                
                # Парсим формулы
                doc.formulas.extend(self.formula_parser.parse(pdf_doc))
                logger.debug(f"Parsed {len(doc.formulas)} formulas")
                
                # Парсим текст
                doc.paragraphs.extend(self.paragraph_parser.parse(pdf_doc))
                logger.debug(f"Parsed {len(doc.paragraphs)} paragraphs")
                
                # Связываем элементы с их заголовками
                self._link_elements_with_captions(doc)
                logger.debug("Linked elements with captions")
                
            except Exception as e:
                logger.error(f"Error during parsing components: {e}")
                logger.exception(e)
                raise ValueError(f"Error parsing document components: {str(e)}")
            
            return doc
            
        finally:
            # Закрываем документ после использования
            pdf_doc.close()
            
    def _link_elements_with_captions(self, doc: ParsedDocument) -> None:
        """
        Связывает таблицы, изображения и формулы с их заголовками на основе анализа текста.
        
        Args:
            doc (ParsedDocument): Документ для обработки.
        """
        # Находим параграфы, которые могут быть заголовками
        caption_paragraphs = {}
        for i, para in enumerate(doc.paragraphs):
            text = para.text.lower()
            if any(keyword in text for keyword in ["таблица", "рисунок", "формула", "рис.", "табл."]):
                caption_paragraphs[i] = {
                    "text": text,
                    "page": para.page_number
                }
        
        # Для таблиц ищем соответствующие заголовки
        for table in doc.tables:
            table_page = table.page_number
            # Ищем заголовки на той же странице или на предыдущей
            for para_idx, caption_info in caption_paragraphs.items():
                if ("таблица" in caption_info["text"] or "табл." in caption_info["text"]) and \
                   (caption_info["page"] == table_page or caption_info["page"] == table_page - 1):
                    table.title_index_in_paragraphs = para_idx
                    break
        
        # Для изображений ищем соответствующие заголовки
        for image in doc.images:
            image_page = image.page_number
            # Ищем заголовки на той же странице
            for para_idx, caption_info in caption_paragraphs.items():
                if ("рисунок" in caption_info["text"] or "рис." in caption_info["text"]) and \
                   caption_info["page"] == image_page:
                    image.referenced_element_index = para_idx
                    break
        
        # Для формул ищем соответствующие заголовки
        for formula in doc.formulas:
            formula_page = formula.page_number
            # Ищем заголовки на той же странице
            for para_idx, caption_info in caption_paragraphs.items():
                if "формула" in caption_info["text"] and caption_info["page"] == formula_page:
                    formula.referenced_element_index = para_idx
                    break