"""
核心处理模块
提供统一的文档处理接口
"""

from pathlib import Path
from collections import Counter
import time

try:
    from .toc.toc_extractor import TOCExtractor
    from .classification.hierarchy_classifier import HierarchyClassifier
    from .chunking.text_splitter import TextSplitter
    from .output.result_saver import ResultSaver
    from .config.config_loader import get_config
except ImportError:
    from toc.toc_extractor import TOCExtractor
    from classification.hierarchy_classifier import HierarchyClassifier
    from chunking.text_splitter import TextSplitter
    from output.result_saver import ResultSaver
    from config.config_loader import get_config


class DocumentClassifier:
    """
    文档分类切分器
    
    支持PDF和Word文档的目录提取、分类和文本切分
    """
    
    def __init__(self):
        """
        初始化文档分类器
        """
        self.config = get_config()
        self.toc_extractor = TOCExtractor()
        self.hierarchy_classifier = HierarchyClassifier()
        self.text_splitter = TextSplitter()
        self.result_saver = ResultSaver()
    
    def process_document(self, file_path, target_level=None, output_dir=None, 
                        max_chunk_size=None, min_chunk_size=None, save_results=None):
        """
        处理文档：提取目录、分类、切分文本块
        
        参数:
            file_path: 文档文件路径(PDF或Word)
            target_level: 要分类的目标层级(可选，默认从配置文件读取)
            output_dir: 输出目录(可选，仅在save_results=True时使用)
            max_chunk_size: 最大分块字符数(可选，默认从配置文件读取)
            min_chunk_size: 最小分块字符数(可选，默认从配置文件读取)
            save_results: 是否保存结果到文件(可选，默认从配置文件读取)
            
        返回:
            dict: 处理结果，包含目录、分类和文本块信息
        """
        # 从配置文件读取默认值
        if target_level is None:
            target_level = self.config.target_level
        if max_chunk_size is None:
            max_chunk_size = self.config.max_chunk_size
        if min_chunk_size is None:
            min_chunk_size = self.config.min_chunk_size
        if save_results is None:
            save_results = self.config.save_results_default
        file_path = Path(file_path)
        
        # 检查文件是否存在
        if not file_path.exists():
            raise FileNotFoundError(f"文件不存在: {file_path}")
        
        # 检查文件格式
        file_ext = file_path.suffix.lower()
        if file_ext not in ['.pdf', '.docx', '.doc']:
            raise ValueError(f"不支持的文件格式: {file_ext}")
        
        print("=" * 100)
        print("文档分类切分工具 v2.0")
        print("=" * 100)
        print(f"\n文件: {file_path}")
        print(f"格式: {file_ext.upper()}")
        print(f"目标层级: {target_level}级")
        print(f"分块大小: {min_chunk_size}-{max_chunk_size}字符")
        
        # 初始化时间记录
        step_times = {}
        total_start_time = time.time()
        
        # 设置输出目录
        if output_dir is None:
            output_dir = file_path.parent / self.config.default_output_dir
        else:
            output_dir = Path(output_dir)
        
        # ========== 步骤1: 提取目录 ==========
        print("\n" + "=" * 100)
        print("步骤1: 提取文档目录")
        print("=" * 100)
        
        step1_start = time.time()
        toc_info = self.toc_extractor.extract_toc(file_path)
        step1_end = time.time()
        step_times['步骤1_提取目录'] = step1_end - step1_start
        
        if toc_info['toc_count'] == 0:
            raise ValueError("未在文档中检测到目录，无法继续处理")
        
        print(f"\n成功提取 {toc_info['toc_count']} 个目录项")
        print(f"目录所在页: {', '.join(map(str, toc_info['toc_pages']))}")
        print(f"[TIME] 耗时: {step_times['步骤1_提取目录']:.2f}秒")
        
        # ========== 步骤2: 目录层级校对 ==========
        print("\n" + "=" * 100)
        print("步骤2: 目录层级校对")
        print("=" * 100)
        
        step2_start = time.time()
        # 注意：toc_extractor.extract_toc 已经包含了层级识别
        # 这里只是显示层级统计信息
        level_counts = Counter([item['level'] for item in toc_info['toc_items']])
        print("\n目录层级分布:")
        for level in sorted(level_counts.keys()):
            print(f"  {level}级: {level_counts[level]} 项")
        
        # 显示前几个目录项的层级信息
        print("\n目录层级示例（前5项）:")
        for i, item in enumerate(toc_info['toc_items'][:5], 1):
            print(f"  [{i}] 第{item['level']}级: {item['title']}")
        if len(toc_info['toc_items']) > 5:
            print(f"  ... 还有 {len(toc_info['toc_items']) - 5} 个目录项")
        
        step2_end = time.time()
        step_times['步骤2_层级校对'] = step2_end - step2_start
        print(f"[TIME] 耗时: {step_times['步骤2_层级校对']:.2f}秒")
        
        # ========== 步骤3: 目录分类（基于二级目录关键词匹配） ==========
        print("\n" + "=" * 100)
        print("步骤3: 目录分类（基于二级目录关键词匹配）")
        print("=" * 100)
        
        step3_start = time.time()
        classification_result = self.hierarchy_classifier.classify(
            toc_info['toc_items'],
            target_level=target_level
        )
        step3_end = time.time()
        step_times['步骤3_目录分类'] = step3_end - step3_start
        
        if classification_result is None:
            raise ValueError("分类失败，无法继续处理")
        
        # 显示分类统计
        category_counts = Counter([item['category'] for item in classification_result['items']])
        print(f"\n分类统计:")
        for category, count in sorted(category_counts.items(), key=lambda x: x[1], reverse=True):
            print(f"  {category}: {count} 项")
        
        # 显示分类详情（前几项）
        print("\n分类详情示例（前3项）:")
        for i, item in enumerate(classification_result['items'][:3], 1):
            print(f"  [{i}] {item['title']}")
            print(f"      分类: {item['category']}")
            print(f"      二级目录数: {item['level2_count']}")
            if item['level2_titles']:
                print(f"      二级目录: {', '.join(item['level2_titles'][:3])}")
                if len(item['level2_titles']) > 3:
                    print(f"                ... 还有 {len(item['level2_titles']) - 3} 个")
        if len(classification_result['items']) > 3:
            print(f"  ... 还有 {len(classification_result['items']) - 3} 个一级目录")
        print(f"[TIME] 耗时: {step_times['步骤3_目录分类']:.2f}秒")
        
        # ========== 步骤4: 提取文档全文 ==========
        print("\n" + "=" * 100)
        print("步骤4: 提取文档全文")
        print("=" * 100)
        
        step4_start = time.time()
        pages_content = self.text_splitter.extract_full_text(file_path)
        step4_end = time.time()
        step_times['步骤4_提取全文'] = step4_end - step4_start
        
        if not pages_content:
            raise ValueError("无法提取文档全文")
        
        total_chars = sum(len(page['text']) for page in pages_content)
        print(f"\n提取完成，共 {len(pages_content)} 页，{total_chars} 个字符")
        print(f"[TIME] 耗时: {step_times['步骤4_提取全文']:.2f}秒")
        
        # ========== 步骤5: 按分类标题切分文本 ==========
        print("\n" + "=" * 100)
        print("步骤5: 按分类标题智能切分文本")
        print("=" * 100)
        
        step5_start = time.time()
        chunks = self.text_splitter.split_by_hierarchy(
            classification_result['items'],
            pages_content,
            toc_info,
            target_level=target_level,
            max_chunk_size=max_chunk_size,
            min_chunk_size=min_chunk_size
        )
        step5_end = time.time()
        step_times['步骤5_切分文本'] = step5_end - step5_start
        
        if not chunks:
            raise ValueError("未能生成任何文本块")
        
        print(f"\n切分完成，共生成 {len(chunks)} 个文本块")
        
        # 显示前5个文本块的信息
        print("\n文本块预览:")
        for i, chunk in enumerate(chunks[:5], 1):
            print(f"  [{i}] {chunk['section_label']} ({len(chunk['review_chunk_content'])} 字符)")
        if len(chunks) > 5:
            print(f"  ... 还有 {len(chunks) - 5} 个文本块")
        print(f"[TIME] 耗时: {step_times['步骤5_切分文本']:.2f}秒")
        
        # ========== 步骤6: 保存结果（可选） ==========
        saved_files = None
        if save_results:
            print("\n" + "=" * 100)
            print("步骤6: 保存结果")
            print("=" * 100)
            
            step6_start = time.time()
            # 保存结果
            saved_files = self.result_saver.save_all(
                file_path, 
                toc_info, 
                classification_result, 
                chunks, 
                output_dir
            )
            step6_end = time.time()
            step_times['步骤6_保存结果'] = step6_end - step6_start
            print(f"[TIME] 耗时: {step_times['步骤6_保存结果']:.2f}秒")
        
        # ========== 完成 ==========
        total_end_time = time.time()
        total_time = total_end_time - total_start_time
        
        print("\n" + "=" * 100)
        print("处理完成！")
        print("=" * 100)
        
        if save_results:
            print(f"\n结果已保存到: {output_dir}")
        print(f"文本块总数: {len(chunks)}")
        print(f"类别数量: {len(category_counts)}")
        
        # 显示时间统计
        print("\n" + "=" * 100)
        print("[TIME] 时间统计")
        print("=" * 100)
        print(f"\n总耗时: {total_time:.2f}秒")
        print("\n各步骤耗时:")
        for step_name, step_time in step_times.items():
            percentage = (step_time / total_time * 100) if total_time > 0 else 0
            print(f"  {step_name}: {step_time:.2f}秒 ({percentage:.1f}%)")
        
        # 找出最耗时的步骤
        if step_times:
            slowest_step = max(step_times.items(), key=lambda x: x[1])
            print(f"\n[WARN] 最耗时步骤: {slowest_step[0]} ({slowest_step[1]:.2f}秒)")
        
        return {
            'toc_info': toc_info,
            'classification': classification_result,
            'chunks': chunks,
            'saved_files': saved_files,
            'output_dir': str(output_dir) if output_dir else None
        }

