"""
文档处理器
负责文档解析、内容提取和结构化处理
集成doc_worker模块的智能处理能力
"""

import io
import os
import tempfile
from pathlib import Path
from typing import Dict, Any, Optional, Callable
from datetime import datetime

from foundation.logger.loggering import server_logger as logger

# 引入doc_worker核心组件
try:
    from .doc_worker import TOCExtractor, TextSplitter, LLMClassifier
    from .doc_worker.config.config_loader import get_config
except ImportError:
    from core.construction_review.component.doc_worker import TOCExtractor, TextSplitter, LLMClassifier
    from core.construction_review.component.doc_worker.config.config_loader import get_config

class DocumentProcessor:
    """文档处理器"""

    def __init__(self):
        self.supported_types = ['pdf', 'docx']
        # 初始化doc_worker组件
        self.toc_extractor = TOCExtractor()
        self.text_splitter = TextSplitter()
        self.config = get_config()
        # LLM分类器可选，如果配置了模型URL则初始化
        self.llm_classifier = None
        try:
            model_url = self.config.llm_model_url
            if model_url:
                self.llm_classifier = LLMClassifier(model_url)
        except Exception as e:
            logger.warning(f"LLM分类器初始化失败，将使用基础处理模式: {str(e)}")

    async def process_document(self, file_content: bytes, file_type: str,
                            #  progress_callback: Optional[Callable[[int, str], None]] = None
                             ) -> Dict[str, Any]:
        """
        处理文档

        Args:
            file_content: 文件内容
            file_type: 文件类型
            progress_callback: 进度回调函数

        Returns:
            Dict: 解析结果
        """
        try:
            logger.info(f"开始处理文档，类型: {file_type}")
            # if progress_callback:
                #  progress_callback(20, "开始文档处理")
            # 简化处理：直接解析
            if file_type.lower() == 'pdf':
                result = await self.parse_pdf_content(file_content)
            elif file_type.lower() == 'docx':
                result = await self.parse_docx_content(file_content)
            else:
                raise ValueError(f"不支持的文件类型: {file_type}")

            # 结构化内容
            structured_result = self.structure_content(result)



            return structured_result

        except Exception as e:
            logger.error(f"文档处理失败: {str(e)}")
            raise

    async def parse_pdf_content(self, file_content: bytes) -> Dict[str, Any]:
        """解析PDF内容，使用doc_worker的智能处理能力"""
        temp_file_path = None
        try:
            # 保存到临时文件
            with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as temp_file:
                temp_file.write(file_content)
                temp_file_path = temp_file.name

            logger.info(f"开始使用doc_worker处理PDF文档: {temp_file_path}")

            # 步骤1: 提取目录
            logger.info("步骤1: 提取文档目录")
            toc_info = self.toc_extractor.extract_toc(temp_file_path)
            
            if toc_info['toc_count'] == 0:
                logger.warning("未检测到目录，使用基础处理模式")
                return await self._fallback_pdf_processing(temp_file_path)

            logger.info(f"成功提取 {toc_info['toc_count']} 个目录项")

            # 步骤2: 使用LLM进行分类（如果可用）
            classified_items = None
            target_level = self.config.target_level
            
            # if self.llm_classifier:
            #     try:
            #         logger.info(f"步骤2: 使用LLM对{target_level}级目录进行分类")
            #         classification_result = self.llm_classifier.classify(
            #             toc_info['toc_items'],
            #             target_level=target_level
            #         )
            #         if classification_result:
            #             classified_items = classification_result['items']
            #             logger.info(f"分类完成，共分类 {len(classified_items)} 个目录项")
            #     except Exception as e:
            #         logger.warning(f"LLM分类失败，使用目录项直接处理: {str(e)}")
            
            # 如果没有分类结果，使用原始目录项（筛选目标层级）
            if not classified_items:
                classified_items = [
                    item for item in toc_info['toc_items'] 
                    if item['level'] == target_level
                ]
                # 为每个目录项添加默认分类信息
                for item in classified_items:
                    item['category'] = '未分类'
                    item['category_code'] = 'other'

            # 步骤3: 提取文档全文
            logger.info("步骤3: 提取文档全文")
            pages_content = self.text_splitter.extract_full_text(temp_file_path)
            
            if not pages_content:
                logger.warning("无法提取文档全文，使用基础处理模式")
                return await self._fallback_pdf_processing(temp_file_path)

            total_chars = sum(len(page['text']) for page in pages_content)
            logger.info(f"提取完成，共 {len(pages_content)} 页，{total_chars} 个字符")

            # 步骤4: 按分类标题智能切分文本
            logger.info("步骤4: 按分类标题智能切分文本")
            max_chunk_size = self.config.max_chunk_size
            min_chunk_size = self.config.min_chunk_size
            
            chunks = self.text_splitter.split_by_hierarchy(
                classified_items,
                pages_content,
                toc_info,
                target_level=target_level,
                max_chunk_size=max_chunk_size,
                min_chunk_size=min_chunk_size
            )

            if not chunks:
                logger.warning("未能生成任何文本块，使用基础处理模式")
                return await self._fallback_pdf_processing(temp_file_path)

            logger.info(f"切分完成，共生成 {len(chunks)} 个文本块")

            # 适配返回格式
            return {
                'document_type': 'pdf',
                'total_pages': len(pages_content),
                'total_chunks': len(chunks),
                'chunks': [
                    {
                        'page': chunk.get('element_tag', {}).get('page', 0),
                        'content': chunk.get('review_chunk_content', ''),
                        'metadata': {
                            'chunk_id': chunk.get('chunk_id', ''),
                            'section_label': chunk.get('section_label', ''),
                            'project_plan_type': chunk.get('project_plan_type', ''),
                            'element_tag': chunk.get('element_tag', {})
                        }
                    }
                    for chunk in chunks
                ],
                'splits': [
                    {
                        'content': chunk.get('review_chunk_content', ''),
                        'metadata': {
                            'chunk_id': chunk.get('chunk_id', ''),
                            'section_label': chunk.get('section_label', ''),
                            'page': chunk.get('element_tag', {}).get('page', 0)
                        }
                    }
                    for chunk in chunks
                ],
                'toc_info': toc_info,
                'classification': {
                    'items': classified_items,
                    'target_level': target_level
                } if classified_items else None
            }

        except Exception as e:
            logger.error(f"PDF解析失败: {str(e)}")
            # 如果智能处理失败，尝试基础处理
            if temp_file_path and os.path.exists(temp_file_path):
                try:
                    logger.info("尝试使用基础处理模式")
                    return await self._fallback_pdf_processing(temp_file_path)
                except Exception as fallback_error:
                    logger.error(f"基础处理模式也失败: {str(fallback_error)}")
            raise
        finally:
            # 清理临时文件
            if temp_file_path and os.path.exists(temp_file_path):
                try:
                    os.unlink(temp_file_path)
                except Exception as e:
                    logger.warning(f"清理临时文件失败: {str(e)}")

    async def parse_docx_content(self, file_content: bytes) -> Dict[str, Any]:
        """解析DOCX内容，使用doc_worker的智能处理能力"""
        temp_file_path = None
        try:
            # 保存到临时文件
            with tempfile.NamedTemporaryFile(delete=False, suffix='.docx') as temp_file:
                temp_file.write(file_content)
                temp_file_path = temp_file.name

            logger.info(f"开始使用doc_worker处理DOCX文档: {temp_file_path}")

            # 步骤1: 提取目录
            logger.info("步骤1: 提取文档目录")
            toc_info = self.toc_extractor.extract_toc(temp_file_path)
            
            if toc_info['toc_count'] == 0:
                logger.warning("未检测到目录，使用基础处理模式")
                return await self._fallback_docx_processing(temp_file_path)

            logger.info(f"成功提取 {toc_info['toc_count']} 个目录项")

            # 步骤2: 使用LLM进行分类（如果可用）
            classified_items = None
            target_level = self.config.target_level
            
            if self.llm_classifier:
                try:
                    logger.info(f"步骤2: 使用LLM对{target_level}级目录进行分类")
                    classification_result = self.llm_classifier.classify(
                        toc_info['toc_items'],
                        target_level=target_level
                    )
                    if classification_result:
                        classified_items = classification_result['items']
                        logger.info(f"分类完成，共分类 {len(classified_items)} 个目录项")
                except Exception as e:
                    logger.warning(f"LLM分类失败，使用目录项直接处理: {str(e)}")
            
            # 如果没有分类结果，使用原始目录项（筛选目标层级）
            if not classified_items:
                classified_items = [
                    item for item in toc_info['toc_items'] 
                    if item['level'] == target_level
                ]
                # 为每个目录项添加默认分类信息
                for item in classified_items:
                    item['category'] = '未分类'
                    item['category_code'] = 'other'

            # 步骤3: 提取文档全文
            logger.info("步骤3: 提取文档全文")
            pages_content = self.text_splitter.extract_full_text(temp_file_path)
            
            if not pages_content:
                logger.warning("无法提取文档全文，使用基础处理模式")
                return await self._fallback_docx_processing(temp_file_path)

            total_chars = sum(len(page['text']) for page in pages_content)
            logger.info(f"提取完成，共 {len(pages_content)} 页，{total_chars} 个字符")

            # 步骤4: 按分类标题智能切分文本
            logger.info("步骤4: 按分类标题智能切分文本")
            max_chunk_size = self.config.max_chunk_size
            min_chunk_size = self.config.min_chunk_size
            
            chunks = self.text_splitter.split_by_hierarchy(
                classified_items,
                pages_content,
                toc_info,
                target_level=target_level,
                max_chunk_size=max_chunk_size,
                min_chunk_size=min_chunk_size
            )

            if not chunks:
                logger.warning("未能生成任何文本块，使用基础处理模式")
                return await self._fallback_docx_processing(temp_file_path)

            logger.info(f"切分完成，共生成 {len(chunks)} 个文本块")

            # 适配返回格式
            return {
                'document_type': 'docx',
                'total_pages': len(pages_content),
                'total_chunks': len(chunks),
                'chunks': [
                    {
                        'page': chunk.get('element_tag', {}).get('page', 0),
                        'content': chunk.get('review_chunk_content', ''),
                        'metadata': {
                            'chunk_id': chunk.get('chunk_id', ''),
                            'section_label': chunk.get('section_label', ''),
                            'project_plan_type': chunk.get('project_plan_type', ''),
                            'element_tag': chunk.get('element_tag', {})
                        }
                    }
                    for chunk in chunks
                ],
                'splits': [
                    {
                        'content': chunk.get('review_chunk_content', ''),
                        'metadata': {
                            'chunk_id': chunk.get('chunk_id', ''),
                            'section_label': chunk.get('section_label', ''),
                            'page': chunk.get('element_tag', {}).get('page', 0)
                        }
                    }
                    for chunk in chunks
                ],
                'full_text': ''.join([page['text'] for page in pages_content]),
                'toc_info': toc_info,
                'classification': {
                    'items': classified_items,
                    'target_level': target_level
                } if classified_items else None,
                'metadata': {
                    'total_pages': len(pages_content),
                    'total_chars': total_chars
                }
            }

        except Exception as e:
            logger.error(f"DOCX解析失败: {str(e)}")
            # 如果智能处理失败，尝试基础处理
            if temp_file_path and os.path.exists(temp_file_path):
                try:
                    logger.info("尝试使用基础处理模式")
                    return await self._fallback_docx_processing(temp_file_path)
                except Exception as fallback_error:
                    logger.error(f"基础处理模式也失败: {str(fallback_error)}")
            raise
        finally:
            # 清理临时文件
            if temp_file_path and os.path.exists(temp_file_path):
                try:
                    os.unlink(temp_file_path)
                except Exception as e:
                    logger.warning(f"清理临时文件失败: {str(e)}")

    async def _fallback_pdf_processing(self, file_path: str) -> Dict[str, Any]:
        """PDF基础处理模式（当智能处理失败时使用）"""
        try:
            from langchain_community.document_loaders import PyPDFLoader
            from langchain.text_splitter import RecursiveCharacterTextSplitter
            
            logger.info("使用基础PDF处理模式")
            loader = PyPDFLoader(file_path)
            documents = loader.load()

            # 文本分块
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=1000,
                chunk_overlap=200,
                separators=["\n\n", "\n", " ", ""]
            )
            splits = text_splitter.split_documents(documents)

            # 过滤空内容切块
            valid_splits = []
            for split in splits:
                content = split.page_content.strip()
                if content:
                    split.page_content = content
                    valid_splits.append(split)

            logger.info(f"基础处理完成，有效分块数量: {len(valid_splits)}")

            return {
                'document_type': 'pdf',
                'total_pages': len(documents),
                'total_chunks': len(valid_splits),
                'chunks': [
                    {
                        'page': doc.metadata.get('page', 0),
                        'content': doc.page_content,
                        'metadata': doc.metadata
                    }
                    for doc in documents
                ],
                'splits': [
                    {
                        'content': split.page_content,
                        'metadata': split.metadata
                    }
                    for split in valid_splits
                ]
            }
        except Exception as e:
            logger.error(f"基础PDF处理失败: {str(e)}")
            raise

    async def _fallback_docx_processing(self, file_path: str) -> Dict[str, Any]:
        """DOCX基础处理模式（当智能处理失败时使用）"""
        try:
            from docx import Document
            
            logger.info("使用基础DOCX处理模式")
            doc = Document(file_path)
            full_text = '\n'.join([paragraph.text for paragraph in doc.paragraphs])

            # 简单分块，并过滤空内容
            chunks = []
            chunk_size = 1000
            chunk_index = 1
            for i in range(0, len(full_text), chunk_size):
                chunk_text = full_text[i:i+chunk_size].strip()
                if chunk_text:
                    chunks.append({
                        'chunk_id': f'chunk_{chunk_index}',
                        'content': chunk_text,
                        'metadata': {'chunk_index': chunk_index}
                    })
                    chunk_index += 1

            logger.info(f"基础处理完成，有效分块数量: {len(chunks)}")

            return {
                'document_type': 'docx',
                'total_chunks': len(chunks),
                'full_text': full_text,
                'chunks': chunks,
                'metadata': {
                    'paragraphs_count': len(doc.paragraphs),
                    'word_count': len(full_text.split())
                }
            }
        except Exception as e:
            logger.error(f"基础DOCX处理失败: {str(e)}")
            raise

    def structure_content(self, raw_content: Dict[str, Any]) -> Dict[str, Any]:
        """结构化处理，适配doc_worker返回的格式"""
        try:
            document_type = raw_content.get('document_type', 'unknown')
            
            # 检查是否使用了doc_worker的智能处理（有toc_info或classification字段）
            is_smart_processing = 'toc_info' in raw_content or 'classification' in raw_content
            
            if is_smart_processing:
                # 使用doc_worker智能处理的结果
                chunks = []
                for chunk in raw_content.get('chunks', []):
                    content = chunk.get('content', '').strip()
                    if content:
                        metadata = chunk.get('metadata', {})
                        element_tag = metadata.get('element_tag', {})
                        
                        chunks.append({
                            'chunk_id': metadata.get('chunk_id', ''),
                            'page': chunk.get('page', 0),
                            'content': content,
                            'section_label': metadata.get('section_label', ''),
                            'project_plan_type': metadata.get('project_plan_type', ''),
                            'element_tag': element_tag,
                            'chapter': metadata.get('section_label', f'第{chunk.get("page", 0)}页'),
                            'title': metadata.get('section_label', ''),
                            'original_content': content[:100] + '...' if len(content) > 100 else content
                        })
            else:
                # 使用基础处理的结果
                if document_type == 'pdf':
                    chunks = []
                    for i, chunk in enumerate(raw_content.get('chunks', [])):
                        content = chunk.get('content', '').strip() if isinstance(chunk, dict) else str(chunk).strip()
                        if content:
                            page = chunk.get('page', 0) if isinstance(chunk, dict) else 0
                            chunks.append({
                                'chunk_id': f'chunk_{i+1}',
                                'page': page,
                                'content': content,
                                'chapter': f'第{page}页',
                                'title': f'内容块{i+1}',
                                'original_content': content[:100] + '...' if len(content) > 100 else content
                            })
                else:
                    # DOCX基础处理
                    all_chunks = raw_content.get('chunks', [])
                    chunks = []
                    for chunk in all_chunks:
                        content = chunk.get('content', '').strip()
                        if content:
                            chunks.append({
                                'chunk_id': chunk.get('chunk_id', f'chunk_{len(chunks)+1}'),
                                'content': content,
                                'metadata': chunk.get('metadata', {})
                            })

            # 构建返回结果
            result = {
                'document_name': f"施工方案文档_{document_type}",
                'document_type': document_type,
                'total_chunks': len(chunks),
                'chunks': chunks,
                'metadata': raw_content.get('metadata', {})
            }
            
            # 如果使用了智能处理，保留额外信息
            if is_smart_processing:
                if 'toc_info' in raw_content:
                    result['toc_info'] = raw_content['toc_info']
                if 'classification' in raw_content:
                    result['classification'] = raw_content['classification']

            return result

        except Exception as e:
            logger.error(f"内容结构化失败: {str(e)}")
            raise