"""
目录提取模块
支持从PDF和Word文档中提取目录结构
"""

from pathlib import Path
from typing import Union

try:
    from ..config.config_loader import get_config
    from .document_extractor_toc import DocumentExtractorTOC
    from .toc_pattern_matcher import TOCPatternMatcher
    from .toc_level_identifier import TOCLevelIdentifier
except ImportError:
    from config.config_loader import get_config
    from document_extractor_toc import DocumentExtractorTOC
    from toc_pattern_matcher import TOCPatternMatcher
    from toc_level_identifier import TOCLevelIdentifier


class TOCExtractor:
    """目录提取器，支持PDF和Word格式，支持文件路径和字节流输入"""
    
    def __init__(self):
        self.config = get_config()
        self.document_extractor = DocumentExtractorTOC()
        self.pattern_matcher = TOCPatternMatcher()
        self.level_identifier = TOCLevelIdentifier()
    
    def extract_toc(self, file_input: Union[str, Path, bytes], file_type: str = None):
        """
        提取文档目录
        
        参数:
            file_input: 文档路径(PDF或Word)或字节流
            file_type: 文件类型（'pdf'或'docx'），当file_input为bytes时必需
            
        返回:
            dict: 包含目录项和统计信息的字典
        """
        # 判断输入类型
        if isinstance(file_input, bytes):
            if not file_type:
                raise ValueError("当输入为字节流时，必须指定file_type参数（'pdf'或'docx'）")
            file_ext = f'.{file_type.lower()}'
            if file_ext == '.pdf':
                return self._extract_from_pdf(file_input, is_bytes=True)
            elif file_ext in ['.docx', '.doc']:
                return self._extract_from_word(file_input, is_bytes=True)
            else:
                raise ValueError(f"不支持的文件格式: {file_ext}")
        else:
            # 文件路径输入（保持向后兼容）
            file_path = Path(file_input)
            file_ext = file_path.suffix.lower()
            
            if file_ext == '.pdf':
                return self._extract_from_pdf(file_path, is_bytes=False)
            elif file_ext in ['.docx', '.doc']:
                return self._extract_from_word(file_path, is_bytes=False)
            else:
                raise ValueError(f"不支持的文件格式: {file_ext}")
    
    def _extract_from_pdf(self, pdf_input, max_pages=None, is_bytes=False):
        """从PDF中提取目录"""
        if max_pages is None:
            max_pages = self.config.toc_max_pages
        pages_text = self.document_extractor.extract_pdf_pages(pdf_input, max_pages, is_bytes=is_bytes)
        
        all_toc_items = []
        toc_page_nums = []
        
        for page_info in pages_text:
            toc_items = self.pattern_matcher.detect_toc_patterns(page_info['text'])
            
            if toc_items:
                all_toc_items.extend(toc_items)
                toc_page_nums.append(page_info['page_num'])
        
        # 去重
        unique_toc = []
        seen = set()
        for item in all_toc_items:
            key = (item['title'], item['page'])
            if key not in seen:
                seen.add(key)
                unique_toc.append(item)
        
        # 使用递归层级识别方法重新识别层级
        unique_toc = self.level_identifier.identify_levels(unique_toc)
        
        return {
            'toc_items': unique_toc,
            'toc_count': len(unique_toc),
            'toc_pages': toc_page_nums
        }
    
    def _extract_from_word(self, word_input, max_pages=None, is_bytes=False):
        """从Word中提取目录"""
        if max_pages is None:
            max_pages = self.config.toc_max_pages
        
        # 方法1: 尝试提取内置目录结构
        builtin_toc = self.document_extractor.extract_builtin_toc(word_input, is_bytes=is_bytes)
        
        # 方法2: 文本模式匹配（作为补充）
        pages_text = self.document_extractor.extract_word_pages(word_input, max_pages, is_bytes=is_bytes)
        
        pattern_toc_items = []
        toc_page_nums = []
        
        for page_info in pages_text:
            toc_items = self.pattern_matcher.detect_toc_patterns(page_info['text'])
            
            if toc_items:
                pattern_toc_items.extend(toc_items)
                toc_page_nums.append(page_info['page_num'])
        
        # 合并两种方法的结果
        all_toc_items = []
        
        # 优先使用内置目录
        if builtin_toc:
            all_toc_items.extend(builtin_toc)
        
        # 如果内置目录为空或数量较少，使用模式匹配的结果
        if len(builtin_toc) < 3:
            all_toc_items.extend(pattern_toc_items)
        
        # 去重
        unique_toc = []
        seen = set()
        for item in all_toc_items:
            key = (item['title'], item.get('page', '?'))
            if key not in seen:
                seen.add(key)
                unique_toc.append(item)
        
        # 使用递归层级识别方法重新识别层级
        unique_toc = self.level_identifier.identify_levels(unique_toc)
        
        return {
            'toc_items': unique_toc,
            'toc_count': len(unique_toc),
            'toc_pages': toc_page_nums if toc_page_nums else [1]
        }
    
    def _detect_level(self, title):
        """
        检测目录项的层级（已废弃）
        
        注意：此方法已不再使用，现在使用递归层级识别（_identify_levels）代替。
        保留此方法仅用于向后兼容和测试。
        """
        import re
        if re.match(r'^【\d+】', title):
            return 1
        
        # 检查数字编号层级（如 1.1, 1.1.1, 1.1.1.1）
        number_match = re.match(r'^(\d+(?:\.\d+)*)\s', title)
        if number_match:
            number_part = number_match.group(1)
            dot_count = number_part.count('.')
            return dot_count + 1
        
        # 检查〖〗格式的编号
        bracket_match = re.match(r'^〖(\d+(?:\.\d+)*)〗', title)
        if bracket_match:
            number_part = bracket_match.group(1)
            dot_count = number_part.count('.')
            return dot_count + 1
        
        # 使用硬编码的简单规则（不再从配置读取，因为配置已删除）
        # 这些规则仅用于向后兼容
        if re.match(r'^第[一二三四五六七八九十\d]+章', title):
            return 1
        if re.match(r'^第[一二三四五六七八九十\d]+节', title):
            return 2
        if re.match(r'^\([一二三四五六七八九十]+\)', title):
            return 3
        
        return 1

