"""
文本切分模块
实现按目录层级和字符数的智能切分逻辑
"""

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

try:
    from ..config.config_loader import get_config
    from .title_matcher import TitleMatcher
    from .text_utils import TextUtils
    from .chunk_splitter import ChunkSplitter
    from .chunk_merger import ChunkMerger
    from .chunk_metadata import ChunkMetadata
    from .hierarchy_processor import HierarchyProcessor
except ImportError:
    from config.config_loader import get_config
    from title_matcher import TitleMatcher
    from text_utils import TextUtils
    from chunk_splitter import ChunkSplitter
    from chunk_merger import ChunkMerger
    from chunk_metadata import ChunkMetadata
    from hierarchy_processor import HierarchyProcessor


class TextSplitter:
    """文本切分器，支持PDF和Word格式，支持文件路径和字节流输入"""
    
    def __init__(self):
        self.config = get_config()
        self.title_matcher = TitleMatcher()
        self.text_utils = TextUtils()
        self.chunk_splitter = ChunkSplitter()
        self.chunk_merger = ChunkMerger()
        self.chunk_metadata = ChunkMetadata()
        self.hierarchy_processor = HierarchyProcessor()
    
    def extract_full_text(self, file_input: Union[str, Path, bytes], file_type: str = None):
        """
        提取文档的全文内容
        
        参数:
            file_input: 文档路径(PDF或Word)或字节流
            file_type: 文件类型（'pdf'或'docx'），当file_input为bytes时必需
            
        返回:
            list: 每页的文本内容
        """
        # 判断输入类型
        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, is_bytes=False):
        """提取PDF的全文内容"""
        try:
            if is_bytes:
                # 从字节流打开
                bytes_io = io.BytesIO(pdf_input)
                doc = fitz.open(stream=bytes_io)
                source_file = 'bytes_stream'
            else:
                # 从文件路径打开
                doc = fitz.open(pdf_input)
                source_file = str(pdf_input)
            
            pages_content = []
            current_pos = 0
            
            for page_num in range(len(doc)):
                page = doc[page_num]
                text = page.get_text()
                
                pages_content.append({
                    'page_num': page_num + 1,
                    'text': text,
                    'start_pos': current_pos,
                    'end_pos': current_pos + len(text),
                    'source_file': source_file
                })
                
                current_pos += len(text)
            
            doc.close()
            return pages_content
        except Exception as e:
            print(f"  错误: 无法读取PDF全文 - {str(e)}")
            return []
    
    def _extract_from_word(self, word_input, is_bytes=False):
        """提取Word的全文内容（包括段落和表格）"""
        try:
            if is_bytes:
                # 从字节流打开
                bytes_io = io.BytesIO(word_input)
                doc = Document(bytes_io)
                source_file = 'bytes_stream'
            else:
                # 从文件路径打开
                doc = Document(word_input)
                source_file = str(word_input)
            
            pages_content = []
            current_pos = 0
            
            # 提取所有内容（段落和表格按文档顺序）
            all_content = []
            
            # 遍历文档的所有元素（段落和表格）
            for element in doc.element.body:
                # 检查是段落还是表格
                if element.tag.endswith('p'):  # 段落
                    for para in doc.paragraphs:
                        if para._element == element:
                            text = para.text
                            if text.strip():
                                all_content.append(text)
                            break
                elif element.tag.endswith('tbl'):  # 表格
                    for table in doc.tables:
                        if table._element == element:
                            table_text = self._extract_table_text(table)
                            all_content.append(table_text)
                            break
            
            # 模拟分页：每30个元素作为一"页"
            elements_per_page = 30
            for page_num in range(0, len(all_content), elements_per_page):
                page_elements = all_content[page_num:page_num + elements_per_page]
                page_text = '\n'.join(page_elements)
                
                pages_content.append({
                    'page_num': page_num // elements_per_page + 1,
                    'text': page_text,
                    'start_pos': current_pos,
                    'end_pos': current_pos + len(page_text),
                    'source_file': source_file
                })
                
                current_pos += len(page_text)
            
            return pages_content
        except Exception as e:
            print(f"  错误: 无法读取Word全文 - {str(e)}")
            return []
    
    def _extract_table_text(self, table):
        """提取表格内容为文本格式"""
        table_text = []
        for row in table.rows:
            row_text = []
            for cell in row.cells:
                cell_text = cell.text.strip().replace('\n', ' ')
                row_text.append(cell_text)
            table_text.append('\t'.join(row_text))
        
        return '\n[表格开始]\n' + '\n'.join(table_text) + '\n[表格结束]\n'
    
    def split_by_hierarchy(self, classified_items, pages_content, toc_info, 
                          target_level=2, max_chunk_size=1000, min_chunk_size=500, 
                          use_concurrent=True, max_workers=None):
        """
        按目录层级和字符数智能切分文本
        
        新的分块逻辑：
        1. 按目录项定位到指定层级的正文标题
        2. 在指定层级正文标题所属的正文块中，先按目录项的最低层级子标题进行分块
        3. 然后逐个判断字符数：
           - 超过max_chunk_size的进行分割(句子级，保持语义完整，分割的块不参与合并)
           - 不足min_chunk_size的块进行合并(合并后不能超过max_chunk_size，否则不合并)
        
        参数:
            classified_items: 已分类的目录项列表
            pages_content: 文档全文内容（按页）
            toc_info: 目录信息
            target_level: 目标层级
            max_chunk_size: 最大分块字符数
            min_chunk_size: 最小分块字符数
            use_concurrent: 是否使用并发处理（默认True）
            max_workers: 最大并发线程数（默认None，使用系统默认值）
            
        返回:
            list: 带分类信息的文本块列表
        """
        full_text = ''.join([page['text'] for page in pages_content])
        
        print(f"  正在定位{len(classified_items)}个已分类的标题...")
        print(f"  目录所在页: {toc_info['toc_pages']}")
        
        # 步骤1: 在正文中定位已分类的标题（跳过目录页）
        located_titles = self.title_matcher.find_title_positions(
            classified_items, 
            full_text, 
            pages_content, 
            toc_info['toc_pages']
        )
        
        # 只保留成功定位的标题
        found_titles = [t for t in located_titles if t['found']]
        
        if not found_titles:
            print(f"  错误: 未能在正文中定位任何标题")
            return []
        
        print(f"  成功定位 {len(found_titles)}/{len(classified_items)} 个标题")
        
        # 按位置排序
        found_titles.sort(key=lambda x: x['position'])
        
        # 步骤2: 提取所有层级的目录项，用于在正文块中查找子标题
        all_toc_items = toc_info['toc_items']
        
        # 步骤2.5: 为每个找到的标题构建完整的层级路径
        for title_info in found_titles:
            hierarchy_path = self.chunk_metadata.build_hierarchy_path(title_info['title'], all_toc_items, target_level)
            title_info['hierarchy_path'] = hierarchy_path
        
        # 步骤3: 按一级目录分组并并发处理
        if use_concurrent:
            all_chunks = self.hierarchy_processor.split_by_hierarchy_concurrent(
                found_titles, full_text, pages_content, all_toc_items,
                target_level, max_chunk_size, min_chunk_size, max_workers
            )
        else:
            all_chunks = self.hierarchy_processor.split_by_hierarchy_sequential(
                found_titles, full_text, pages_content, all_toc_items,
                target_level, max_chunk_size, min_chunk_size
            )
        
        # 步骤4: 对小块进行合并
        merged_chunks = self.chunk_merger.merge_small_chunks(all_chunks, max_chunk_size, min_chunk_size, target_level)
        
        # 步骤5: 生成最终的chunk_id和serial_number
        final_chunks = self.chunk_metadata.finalize_chunk_ids(merged_chunks)
        
        print(f"  初始切分: {len(all_chunks)} 个块")
        print(f"  合并后: {len(merged_chunks)} 个块")
        
        return final_chunks
