# requirements.txt 示例
# langchain==0.1.0
# unstructured[pdf,docx,html]==0.10.30
# pypdf==3.17.0
# chromadb==0.4.22
# qdrant-client==1.7.0
# tenacity==8.2.3  # 用于重试

import os
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
from enum import Enum
import hashlib
import logging
from tenacity import retry, stop_after_attempt, wait_exponential

from langchain.text_splitter import RecursiveCharacterTextSplitter, MarkdownHeaderTextSplitter
from langchain.document_loaders import (
    PyPDFLoader,
    UnstructuredWordDocumentLoader,
    UnstructuredHTMLLoader,
)
from unstructured.partition.auto import partition
from unstructured.cleaners.core import clean_extra_whitespace

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class DocumentType(Enum):
    PDF = "pdf"
    DOCX = "docx"
    HTML = "html"
    TXT = "txt"

@dataclass
class DocumentChunk:
    id: str  # 全局唯一ID，可用于去重和引用
    text: str
    metadata: Dict[str, Any]  # 包含：source_file, chunk_index, tenant_id, doc_id, page_num等
    embedding: Optional[List[float]] = None

class EnterpriseDocumentParser:
    """企业级文档解析器，支持多格式并提取元数据"""
    
    def __init__(self, tenant_id: str, doc_id: str):
        self.tenant_id = tenant_id
        self.doc_id = doc_id
        self._cache_dir = f"./cache/{tenant_id}"  # 生产环境应使用分布式缓存或对象存储
        os.makedirs(self._cache_dir, exist_ok=True)
    
    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
    def parse(self, file_path: str, file_type: DocumentType) -> List[Dict[str, Any]]:
        """解析文档，返回包含文本和丰富元数据的字典列表"""
        raw_elements = []
        try:
            # 使用unstructured库进行鲁棒的解析，它能更好地处理复杂布局
            if file_type == DocumentType.PDF:
                raw_elements = partition(filename=file_path, strategy="hi_res")  # 高分辨率策略处理扫描件
            elif file_type == DocumentType.DOCX:
                raw_elements = partition(filename=file_path, content_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document")
            elif file_type == DocumentType.HTML:
                raw_elements = partition(filename=file_path, content_type="text/html")
            else:
                # 回退到langchain loader
                with open(file_path, 'r', encoding='utf-8') as f:
                    raw_elements = [{"text": f.read(), "metadata": {}}]
        except Exception as e:
            logger.error(f"Failed to parse {file_path} with unstructured: {e}. Fallback to langchain.")
            raw_elements = self._fallback_parse(file_path, file_type)
        
        documents = []
        for idx, elem in enumerate(raw_elements):
            text = getattr(elem, "text", str(elem))
            text = clean_extra_whitespace(text)  # 清理多余空白
            
            # 提取结构化元数据
            metadata = {
                "tenant_id": self.tenant_id,
                "doc_id": self.doc_id,
                "source_file": os.path.basename(file_path),
                "file_type": file_type.value,
                "element_index": idx,
                "element_type": getattr(elem, "category", "unknown"),
                "page_number": getattr(elem, "metadata", {}).get("page_number", 1),
                # 提取更多业务元数据，如作者、标题（如果解析器支持）
                **getattr(elem, "metadata", {})
            }
            # 处理表格：如果元素是表格，可以特殊标记
            if getattr(elem, "category", None) == "Table":
                metadata["is_table"] = True
                # 这里可以调用专门的表格提取逻辑，将表格转为Markdown或结构化JSON
                # text = self._convert_table_to_markdown(elem)
            
            if text.strip():  # 忽略空文本
                documents.append({"text": text, "metadata": metadata})
        
        logger.info(f"Parsed {file_path} into {len(documents)} elements.")
        return documents
    
    def _fallback_parse(self, file_path: str, file_type: DocumentType) -> List[Any]:
        """备选解析方案"""
        loaders = {
            DocumentType.PDF: PyPDFLoader,
            DocumentType.DOCX: UnstructuredWordDocumentLoader,
            DocumentType.HTML: UnstructuredHTMLLoader,
        }
        if file_type in loaders:
            loader = loaders[file_type](file_path)
            return loader.load()  # 返回langchain Document对象列表
        else:
            with open(file_path, 'r', encoding='utf-8') as f:
                return [f.read()]
    
    def _convert_table_to_markdown(self, table_element) -> str:
        """将表格元素转换为Markdown格式（简化示例）"""
        # 实际应使用unstructured的table extraction功能或tabula（对于PDF）
        return "| ... | ... |\n| --- | --- |\n| ... | ... |"

class IntelligentTextSplitter:
    """智能分块器，考虑语义边界和特殊内容"""
    
    def __init__(self, chunk_size: int = 1000, chunk_overlap: int = 200):
        self.base_splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
            separators=["\n\n", "\n", "。", "！", "？", "；", "，", " ", ""]  # 中文友好分隔符
        )
        
    def split_documents(self, documents: List[Dict[str, Any]]) -> List[DocumentChunk]:
        """对解析后的文档进行智能分块"""
        all_chunks = []
        
        for doc in documents:
            text = doc["text"]
            metadata = doc["metadata"]
            
            # 策略1：如果是Markdown或HTML，按标题分割
            if metadata["file_type"] in ["html", "md"] and "# " in text:
                try:
                    headers_to_split_on = [("#", "h1"), ("##", "h2"), ("###", "h3")]
                    md_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
                    md_chunks = md_splitter.split_text(text)
                    for chunk in md_chunks:
                        chunk_metadata = {**metadata, **chunk.metadata}
                        all_chunks.append(self._create_chunk(chunk.page_content, chunk_metadata, len(all_chunks)))
                except Exception as e:
                    logger.warning(f"Markdown split failed: {e}. Fallback to recursive split.")
                    # 降级到基础分块
                    base_chunks = self.base_splitter.split_text(text)
                    for chunk in base_chunks:
                        all_chunks.append(self._create_chunk(chunk, metadata, len(all_chunks)))
            
            # 策略2：如果是表格，尽量保持完整（假设表格不大）
            elif metadata.get("is_table"):
                all_chunks.append(self._create_chunk(text, metadata, len(all_chunks)))
            
            # 策略3：普通文本，使用递归分块
            else:
                base_chunks = self.base_splitter.split_text(text)
                for chunk in base_chunks:
                    all_chunks.append(self._create_chunk(chunk, metadata, len(all_chunks)))
        
        logger.info(f"Split into {len(all_chunks)} total chunks.")
        return all_chunks
    
    def _create_chunk(self, text: str, metadata: Dict, index: int) -> DocumentChunk:
        """创建文档块对象，生成唯一ID"""
        chunk_id = hashlib.md5(
            f"{metadata['tenant_id']}_{metadata['doc_id']}_{index}_{text[:50]}".encode()
        ).hexdigest()
        return DocumentChunk(
            id=chunk_id,
            text=text,
            metadata={**metadata, "chunk_index": index}
        )

# 使用示例
if __name__ == "__main__":
    # 模拟一个租户文档
    parser = EnterpriseDocumentParser(tenant_id="tenant_001", doc_id="doc_abc123")
    # 假设我们有一个sample.pdf文件
    parsed_docs = parser.parse("./sample.pdf", DocumentType.PDF)
    
    splitter = IntelligentTextSplitter(chunk_size=800, chunk_overlap=150)
    chunks = splitter.split_documents(parsed_docs)
    
    print(f"生成 {len(chunks)} 个块。")
    print(f"第一个块元数据: {chunks[0].metadata}")
    print(f"第一个块预览: {chunks[0].text[:200]}...")

# 伪代码：检索服务中的租户隔离
from qdrant_client import QdrantClient
from qdrant_client.http import models

class SecureRetrievalService:
    def __init__(self, qdrant_client: QdrantClient):
        self.client = qdrant_client
    
    def search(self, query_vector: List[float], tenant_id: str, top_k: int = 5):
        # 关键：在搜索请求中强制加入租户过滤器
        search_filter = models.Filter(
            must=[
                models.FieldCondition(
                    key="tenant_id",
                    match=models.MatchValue(value=tenant_id)  # 严格匹配租户ID
                )
            ]
        )
        return self.client.search(
            collection_name="enterprise_docs",
            query_vector=query_vector,
            query_filter=search_filter,  # 过滤器确保数据隔离
            limit=top_k
        )