# pdf_processor.py
import PyPDF2
import re
from typing import List, Dict, Tuple
import os
import json
from dataclasses import dataclass
from enum import Enum


class SectionType(Enum):
    ABSTRACT = "abstract"
    INTRODUCTION = "introduction"
    METHODS = "methods"
    RESULTS = "results"
    DISCUSSION = "discussion"
    CONCLUSION = "conclusion"
    REFERENCES = "references"
    OTHER = "other"


@dataclass
class TextSegment:
    section_type: SectionType
    section_title: str
    text: str
    page_number: int
    start_pos: int
    end_pos: int
    confidence: float = 1.0


class PDFSectionExtractor:
    """提取PDF并按章节分割"""

    # 常见章节标题模式（可扩展）
    SECTION_PATTERNS = {
        r'^\s*abstract\s*$': SectionType.ABSTRACT,
        r'^\s*summary\s*$': SectionType.ABSTRACT,
        r'^\s*1\.?\s*introduction': SectionType.INTRODUCTION,
        r'^\s*introduction\s*$': SectionType.INTRODUCTION,
        r'^\s*background\s*$': SectionType.INTRODUCTION,
        r'^\s*2\.?\s*methods': SectionType.METHODS,
        r'^\s*methods\s*$': SectionType.METHODS,
        r'^\s*materials\s*and\s*methods': SectionType.METHODS,
        r'^\s*3\.?\s*results': SectionType.RESULTS,
        r'^\s*results\s*$': SectionType.RESULTS,
        r'^\s*4\.?\s*discussion': SectionType.DISCUSSION,
        r'^\s*discussion\s*$': SectionType.DISCUSSION,
        r'^\s*5\.?\s*conclusion': SectionType.CONCLUSION,
        r'^\s*conclusion\s*$': SectionType.CONCLUSION,
        r'^\s*conclusions\s*$': SectionType.CONCLUSION,
        r'^\s*references\s*$': SectionType.REFERENCES,
        r'^\s*bibliography\s*$': SectionType.REFERENCES,
    }

    def __init__(self, custom_patterns: Dict[str, SectionType] = None):
        if custom_patterns:
            self.SECTION_PATTERNS.update(custom_patterns)

    def extract_text_from_pdf(self, pdf_path: str) -> str:
        """提取PDF所有文本"""
        text = ""
        try:
            with open(pdf_path, 'rb') as file:
                pdf_reader = PyPDF2.PdfReader(file)

                for page_num in range(len(pdf_reader.pages)):
                    page = pdf_reader.pages[page_num]
                    page_text = page.extract_text()

                    if page_text:
                        # 添加页码标记
                        text += f"\n--- PAGE {page_num + 1} ---\n{page_text}"

        except Exception as e:
            print(f"读取PDF错误 {pdf_path}: {e}")

        return text

    def identify_sections(self, text: str) -> List[TextSegment]:
        """识别文本中的章节"""
        lines = text.split('\n')
        segments = []
        current_section = SectionType.OTHER
        current_title = "开始"
        current_text = []
        page_num = 1

        for i, line in enumerate(lines):
            # 检测页码标记
            if line.startswith("--- PAGE "):
                if current_text:
                    segments.append(TextSegment(
                        section_type=current_section,
                        section_title=current_title,
                        text='\n'.join(current_text),
                        page_number=page_num,
                        start_pos=i - len(current_text),
                        end_pos=i
                    ))
                current_text = []
                try:
                    page_num = int(line.split()[2])
                except:
                    pass
                continue

            # 检测章节标题
            line_lower = line.strip().lower()
            section_found = False

            for pattern, section_type in self.SECTION_PATTERNS.items():
                if re.match(pattern, line_lower, re.IGNORECASE):
                    if current_text:
                        segments.append(TextSegment(
                            section_type=current_section,
                            section_title=current_title,
                            text='\n'.join(current_text),
                            page_number=page_num,
                            start_pos=i - len(current_text),
                            end_pos=i
                        ))

                    current_section = section_type
                    current_title = line.strip()
                    current_text = [line]
                    section_found = True
                    break

            if not section_found:
                current_text.append(line)

        # 添加最后一段
        if current_text:
            segments.append(TextSegment(
                section_type=current_section,
                section_title=current_title,
                text='\n'.join(current_text),
                page_number=page_num,
                start_pos=len(lines) - len(current_text),
                end_pos=len(lines)
            ))

        return segments

    def process_pdf(self, pdf_path: str) -> Dict:
        """处理单个PDF文件"""
        print(f"处理PDF: {pdf_path}")

        # 提取文本
        raw_text = self.extract_text_from_pdf(pdf_path)
        if not raw_text:
            return None

        # 分割章节
        segments = self.identify_sections(raw_text)

        # 整理结果
        result = {
            'pdf_path': pdf_path,
            'total_pages': len(raw_text.split('--- PAGE')) - 1,
            'segments': [],
            'section_summary': {}
        }

        for seg in segments:
            segment_data = {
                'section_type': seg.section_type.value,
                'section_title': seg.section_title,
                'text': seg.text,
                'page_number': seg.page_number,
                'char_length': len(seg.text),
                'word_count': len(seg.text.split())
            }
            result['segments'].append(segment_data)

            # 统计章节信息
            if seg.section_type.value not in result['section_summary']:
                result['section_summary'][seg.section_type.value] = 0
            result['section_summary'][seg.section_type.value] += len(seg.text)

        return result