"""
文档提取模块（用于目录提取）
支持从PDF和Word文档中提取文本内容
"""

import io
from pathlib import Path
import fitz  # PyMuPDF
from docx import Document

try:
    from ..config.config_loader import get_config
    from .toc_pattern_matcher import TOCPatternMatcher
except ImportError:
    from config.config_loader import get_config
    from toc_pattern_matcher import TOCPatternMatcher


class DocumentExtractorTOC:
    """文档提取器（用于目录提取）"""
    
    def __init__(self):
        self.config = get_config()
        self.pattern_matcher = TOCPatternMatcher()
    
    def extract_pdf_pages(self, pdf_input, max_pages=None, is_bytes=False):
        """从PDF文件的前几页提取文本"""
        if max_pages is None:
            max_pages = self.config.toc_max_pages
        try:
            if is_bytes:
                bytes_io = io.BytesIO(pdf_input)
                doc = fitz.open(stream=bytes_io)
            else:
                doc = fitz.open(pdf_input)
            
            pages_text = []
            
            for page_num in range(min(len(doc), max_pages)):
                page = doc[page_num]
                text = page.get_text()
                pages_text.append({
                    'page_num': page_num + 1,
                    'text': text
                })
            
            doc.close()
            return pages_text
        except Exception as e:
            print(f"  错误: 无法读取PDF - {str(e)}")
            return []
    
    def extract_word_pages(self, word_input, max_pages=None, is_bytes=False):
        """从Word文件的前几页提取文本"""
        if max_pages is None:
            max_pages = self.config.toc_max_pages
        
        try:
            if is_bytes:
                bytes_io = io.BytesIO(word_input)
                doc = Document(bytes_io)
            else:
                doc = Document(word_input)
            
            pages_text = []
            
            all_text = []
            for para in doc.paragraphs:
                text = para.text.strip()
                if text:
                    all_text.append(text)
            
            # 模拟分页：从配置读取每页段落数
            paragraphs_per_page = self.config.paragraphs_per_page
            for i in range(0, min(len(all_text), max_pages * paragraphs_per_page), paragraphs_per_page):
                page_text = '\n'.join(all_text[i:i+paragraphs_per_page])
                pages_text.append({
                    'page_num': i // paragraphs_per_page + 1,
                    'text': page_text
                })
            
            return pages_text
        except Exception as e:
            print(f"  错误: 无法读取Word - {str(e)}")
            return []
    
    def extract_builtin_toc(self, word_input, is_bytes=False):
        """提取Word文档的内置目录结构"""
        try:
            if is_bytes:
                bytes_io = io.BytesIO(word_input)
                doc = Document(bytes_io)
            else:
                doc = Document(word_input)
            
            toc_items = []
            
            for para in doc.paragraphs:
                style_name = para.style.name if para.style else ""
                text = para.text.strip()
                
                if not text:
                    continue
                
                # 检查是否是标题样式
                if style_name.startswith('Heading'):
                    if not self.pattern_matcher.has_numbering(text):
                        continue
                    
                    try:
                        level = int(style_name.split()[-1]) if len(style_name.split()) > 1 else 1
                    except:
                        level = 1
                    
                    toc_items.append({
                        'title': text,
                        'level': level,
                        'page': '?',
                        'original': text,
                        'source': 'heading_style'
                    })
                # 检查是否是TOC样式
                elif 'TOC' in style_name or 'toc' in style_name.lower():
                    import re
                    match = re.search(r'(\d+)\s*$', text)
                    page = match.group(1) if match else '?'
                    
                    title = re.sub(r'\s*\d+\s*$', '', text).strip()
                    
                    if not self.pattern_matcher.has_numbering(title):
                        continue
                    
                    level_match = re.search(r'TOC\s*(\d+)', style_name, re.IGNORECASE)
                    level = int(level_match.group(1)) if level_match else 1
                    
                    if title:
                        toc_items.append({
                            'title': title,
                            'level': level,
                            'page': page,
                            'original': text,
                            'source': 'toc_style'
                        })
            
            return toc_items
        except Exception as e:
            print(f"  错误: 无法读取Word内置目录 - {str(e)}")
            return []


