services/pdf_processor.py

"""
PDF 파일 처리 및 텍스트 추출
"""
import os
from pathlib import Path
from typing import List, Dict, Optional
import PyPDF2
import pdfplumber
from loguru import logger
from tqdm import tqdm


class PDFProcessor:
    """PDF 파일에서 텍스트와 메타데이터를 추출하는 클래스"""

    def __init__(self, pdf_directory: str):
        """
        Args:
            pdf_directory: PDF 파일들이 있는 디렉토리 경로
        """
        self.pdf_directory = Path(pdf_directory)
        self.processed_docs = []

    def get_pdf_files(self) -> List[Path]:
        """디렉토리에서 모든 PDF 파일 찾기"""
        if not self.pdf_directory.exists():
            raise FileNotFoundError(f"Directory not found: {self.pdf_directory}")

        pdf_files = list(self.pdf_directory.glob("**/*.pdf"))
        logger.info(f"Found {len(pdf_files)} PDF files in {self.pdf_directory}")
        return pdf_files

    def extract_text_from_pdf(self, pdf_path: Path) -> Optional[Dict[str, any]]:
        """
        단일 PDF 파일에서 텍스트 추출

        Args:
            pdf_path: PDF 파일 경로

        Returns:
            Dict with 'text', 'metadata', 'filename', 'page_count'
        """
        try:
            # pdfplumber를 사용한 텍스트 추출 (더 정확함)
            with pdfplumber.open(pdf_path) as pdf:
                text = ""
                for page in pdf.pages:
                    page_text = page.extract_text()
                    if page_text:
                        text += page_text + "\n\n"

                # PyPDF2로 메타데이터 추출
                with open(pdf_path, 'rb') as f:
                    pdf_reader = PyPDF2.PdfReader(f)
                    metadata = pdf_reader.metadata if pdf_reader.metadata else {}
                    page_count = len(pdf_reader.pages)

                return {
                    'text': text.strip(),
                    'metadata': {
                        'title': metadata.get('/Title', ''),
                        'author': metadata.get('/Author', ''),
                        'subject': metadata.get('/Subject', ''),
                        'creator': metadata.get('/Creator', ''),
                    },
                    'filename': pdf_path.name,
                    'filepath': str(pdf_path),
                    'page_count': page_count
                }

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

    def process_all_pdfs(self) -> List[Dict[str, any]]:
        """
        모든 PDF 파일 처리

        Returns:
            List of dictionaries containing extracted text and metadata
        """
        pdf_files = self.get_pdf_files()
        self.processed_docs = []

        logger.info(f"Processing {len(pdf_files)} PDF files...")

        for pdf_path in tqdm(pdf_files, desc="Processing PDFs"):
            doc_data = self.extract_text_from_pdf(pdf_path)
            if doc_data and doc_data['text']:  # 텍스트가 있는 경우만 추가
                self.processed_docs.append(doc_data)
            else:
                logger.warning(f"No text extracted from {pdf_path.name}")

        logger.info(f"Successfully processed {len(self.processed_docs)} PDFs")
        return self.processed_docs

    def get_statistics(self) -> Dict[str, any]:
        """처리된 문서들의 통계 정보"""
        if not self.processed_docs:
            return {}

        total_pages = sum(doc['page_count'] for doc in self.processed_docs)
        total_chars = sum(len(doc['text']) for doc in self.processed_docs)

        return {
            'total_documents': len(self.processed_docs),
            'total_pages': total_pages,
            'total_characters': total_chars,
            'avg_pages_per_doc': total_pages / len(self.processed_docs),
            'avg_chars_per_doc': total_chars / len(self.processed_docs),
        }