"""
层级处理模块
用于按层级分组和处理文本块
"""

from concurrent.futures import ThreadPoolExecutor, as_completed

try:
    from .chunk_splitter import ChunkSplitter
    from .chunk_metadata import ChunkMetadata
except ImportError:
    from chunk_splitter import ChunkSplitter
    from chunk_metadata import ChunkMetadata


class HierarchyProcessor:
    """层级处理器"""
    
    def __init__(self):
        self.chunk_splitter = ChunkSplitter()
        self.chunk_metadata = ChunkMetadata()
    
    def split_by_hierarchy_sequential(self, found_titles, full_text, pages_content, 
                                       all_toc_items, target_level, max_chunk_size, min_chunk_size):
        """
        顺序处理方式（原有逻辑）
        """
        all_chunks = []
        
        for i, title_info in enumerate(found_titles):
            start_pos = title_info['position']
            
            # 确定正文块的结束位置（下一个同级标题的位置）
            if i + 1 < len(found_titles):
                end_pos = found_titles[i + 1]['position']
            else:
                end_pos = len(full_text)
            
            # 提取正文块
            content_block = full_text[start_pos:end_pos]
            
            # 在正文块中查找子标题（直接子标题，level = target_level + 1）
            sub_chunks = self.chunk_splitter.split_by_sub_titles(
                content_block,
                all_toc_items,
                title_info,
                target_level,
                max_chunk_size,
                min_chunk_size,
                full_text,
                start_pos,
                title_info.get('hierarchy_path', [title_info['title']])
            )
            
            # 为每个子块添加元数据
            for j, sub_chunk in enumerate(sub_chunks, 1):
                chunk_data = self.chunk_metadata.build_chunk_metadata(
                    sub_chunk, title_info, start_pos, pages_content, i, j
                )
                all_chunks.append(chunk_data)
        
        return all_chunks
    
    def split_by_hierarchy_concurrent(self, found_titles, full_text, pages_content, 
                                      all_toc_items, target_level, max_chunk_size, 
                                      min_chunk_size, max_workers=None):
        """
        并发处理方式：按一级目录分组，每个线程处理一个一级目录及其子目录
        """
        # 按一级目录分组
        level1_groups = self.group_by_level1(found_titles, all_toc_items)
        
        if not level1_groups:
            # 如果没有一级目录，回退到顺序处理
            print("  未找到一级目录，使用顺序处理")
            return self.split_by_hierarchy_sequential(
                found_titles, full_text, pages_content, all_toc_items,
                target_level, max_chunk_size, min_chunk_size
            )
        
        print(f"  按一级目录分组: {len(level1_groups)} 个一级目录，使用并发处理")
        
        all_chunks = []
        
        # 使用线程池并发处理
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # 提交所有任务
            future_to_group = {
                executor.submit(
                    self.process_level1_group,
                    group, full_text, pages_content, all_toc_items,
                    target_level, max_chunk_size, min_chunk_size
                ): group for group in level1_groups
            }
            
            # 收集结果
            for future in as_completed(future_to_group):
                group = future_to_group[future]
                try:
                    chunks = future.result()
                    all_chunks.extend(chunks)
                    print(f"  完成一级目录处理: {group['level1_title']} ({len(chunks)} 个块)")
                except Exception as e:
                    print(f"  处理一级目录 {group['level1_title']} 时出错: {str(e)}")
                    # 出错时回退到顺序处理该组
                    chunks = self.process_level1_group_sequential(
                        group, full_text, pages_content, all_toc_items,
                        target_level, max_chunk_size, min_chunk_size
                    )
                    all_chunks.extend(chunks)
        
        # 按位置排序所有块（因为并发处理可能打乱顺序）
        all_chunks.sort(key=lambda x: x.get('_sort_key', 0))
        
        return all_chunks
    
    def group_by_level1(self, found_titles, all_toc_items):
        """
        按一级目录分组
        
        参数:
            found_titles: 已定位的标题列表
            all_toc_items: 所有目录项
            
        返回:
            list: 分组列表，每个元素包含一级目录信息和其子目录的标题
        """
        groups = []
        
        # 找到所有一级目录在目录项中的位置，并记录它们在全文中的位置
        level1_indices = []
        level1_positions = {}  # 一级目录标题 -> 在全文中的位置
        
        for idx, toc_item in enumerate(all_toc_items):
            if toc_item.get('level', 1) == 1:
                level1_title = toc_item['title']
                # 查找该一级目录在found_titles中的位置
                for title_info in found_titles:
                    if title_info['title'] == level1_title:
                        level1_positions[level1_title] = title_info['position']
                        break
                level1_indices.append((idx, toc_item))
        
        if not level1_indices:
            return []
        
        # 为每个一级目录创建组
        for i, (level1_idx, level1_item) in enumerate(level1_indices):
            level1_title = level1_item['title']
            
            # 确定该一级目录的范围（到下一个一级目录之前）
            if i + 1 < len(level1_indices):
                next_level1_idx = level1_indices[i + 1][0]
                next_level1_title = level1_indices[i + 1][1]['title']
                next_level1_position = level1_positions.get(next_level1_title, None)
            else:
                next_level1_idx = len(all_toc_items)
                next_level1_position = None
            
            # 找到属于该一级目录的所有标题（在found_titles中）
            group_titles = []
            for title_info in found_titles:
                title = title_info['title']
                # 检查该标题是否在当前一级目录的范围内
                for idx in range(level1_idx, next_level1_idx):
                    if idx < len(all_toc_items) and all_toc_items[idx]['title'] == title:
                        group_titles.append(title_info)
                        break
            
            if group_titles:
                groups.append({
                    'level1_title': level1_title,
                    'level1_index': level1_idx,
                    'level1_end_index': next_level1_idx,
                    'next_level1_position': next_level1_position,  # 下一个一级目录在全文中的位置
                    'titles': group_titles
                })
        
        return groups
    
    def process_level1_group(self, group, full_text, pages_content, all_toc_items,
                             target_level, max_chunk_size, min_chunk_size):
        """
        处理单个一级目录及其子目录
        
        参数:
            group: 一级目录组信息
            full_text: 全文内容
            pages_content: 页面内容列表
            all_toc_items: 所有目录项
            target_level: 目标层级
            max_chunk_size: 最大分块字符数
            min_chunk_size: 最小分块字符数
            
        返回:
            list: 该一级目录下的所有文本块
        """
        group_titles = group['titles']
        all_chunks = []
        
        # 按位置排序
        group_titles.sort(key=lambda x: x['position'])
        
        for i, title_info in enumerate(group_titles):
            start_pos = title_info['position']
            
            # 确定正文块的结束位置（下一个同级标题的位置，或下一个一级目录的开始位置）
            if i + 1 < len(group_titles):
                end_pos = group_titles[i + 1]['position']
            else:
                # 检查是否有下一个一级目录的位置
                next_level1_pos = group.get('next_level1_position')
                if next_level1_pos is not None:
                    end_pos = next_level1_pos
                else:
                    end_pos = len(full_text)
            
            # 提取正文块
            content_block = full_text[start_pos:end_pos]
            
            # 在正文块中查找子标题
            sub_chunks = self.chunk_splitter.split_by_sub_titles(
                content_block,
                all_toc_items,
                title_info,
                target_level,
                max_chunk_size,
                min_chunk_size,
                full_text,
                start_pos,
                title_info.get('hierarchy_path', [title_info['title']])
            )
            
            # 为每个子块添加元数据
            for j, sub_chunk in enumerate(sub_chunks, 1):
                chunk_data = self.chunk_metadata.build_chunk_metadata(
                    sub_chunk, title_info, start_pos, pages_content, i, j
                )
                # 添加排序键（用于后续排序）
                chunk_data['_sort_key'] = start_pos + sub_chunk['relative_start']
                all_chunks.append(chunk_data)
        
        return all_chunks
    
    def process_level1_group_sequential(self, group, full_text, pages_content, all_toc_items,
                                         target_level, max_chunk_size, min_chunk_size):
        """
        顺序处理单个一级目录组（用于错误回退）
        """
        return self.process_level1_group(
            group, full_text, pages_content, all_toc_items,
            target_level, max_chunk_size, min_chunk_size
        )


