"""
文档处理器
负责文档解析、内容提取和结构化处理
"""

import io   
from docx import Document
from typing import Dict, Any, Optional, Callable
from datetime import datetime

from foundation.logger.loggering import server_logger as logger

from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

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

    def __init__(self):
        self.supported_types = ['pdf', 'docx']

    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 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)

            if progress_callback:
                progress_callback(100, "文档处理完成")

            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内容"""
        try:
            # 保存到临时文件
            import tempfile
            with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as temp_file:
                temp_file.write(file_content)
                temp_file_path = temp_file.name

            # 使用PyPDFLoader解析
            loader = PyPDFLoader(temp_file_path)
            documents = loader.load()

            # 文本分块
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=1000,
                chunk_overlap=200,
                separators=["\n\n", "\n", " ", ""]
            )
            splits = text_splitter.split_documents(documents)
            original_count = len(splits)  # 记录原始分块数量

            # 过滤空内容切块，确保每个切块内容不为空
            valid_splits = []
            for split in splits:
                content = split.page_content.strip()
                if content:  # 确保内容不为空
                    split.page_content = content  # 更新清理后的内容
                    valid_splits.append(split)

            splits = valid_splits  # 使用过滤后的切块
            logger.info(f"PDF解析完成，过滤前分块数量: {original_count}，过滤后有效分块数量: {len(splits)}")

            return {
                'document_type': 'pdf',
                'total_pages': len(documents),
                'total_chunks': len(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 splits
                ]
            }

        except Exception as e:
            logger.error(f"PDF解析失败: {str(e)}")
            raise

    async def parse_docx_content(self, file_content: bytes) -> Dict[str, Any]:
        """解析DOCX内容"""
        try:
            # 简化实现：直接返回文本内容
            # 实际实现中可以使用python-docx库


            doc = Document(io.BytesIO(file_content))
            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"DOCX解析完成，有效分块数量: {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]:
        """结构化处理"""
        try:
            if raw_content['document_type'] == 'pdf':
                # PDF结构化
                chunks = []
                for i, chunk in enumerate(raw_content['chunks']):
                    content = chunk['content'].strip()
                    if content:  # 确保内容不为空
                        chunks.append({
                            'chunk_id': f'chunk_{i+1}',
                            'page': chunk['page'],
                            'content': content,
                            'chapter': f'第{chunk["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', {})
                        })

            return {
                'document_name': f"施工方案文档_{raw_content.get('document_type', 'unknown')}",
                'document_type': raw_content['document_type'],
                'total_chunks': len(chunks),  # 使用实际的切块数量
                'chunks': chunks,
                'metadata': raw_content.get('metadata', {})
            }

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