from langchain.text_splitter import (
    RecursiveCharacterTextSplitter,
    TokenTextSplitter,
    MarkdownTextSplitter
)

class ChunkingStrategy:
    """分块策略"""
    
    @staticmethod
    def fixed_size_chunking(text: str, chunk_size: int = 1000, overlap: int = 200):
        """固定大小分块"""
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=overlap,
            separators=["\n\n", "\n", "。", ".", " ", ""]
        )
        return splitter.split_text(text)
    
    @staticmethod
    def token_based_chunking(text: str, chunk_size: int = 500, overlap: int = 50):
        """基于 Token 的分块"""
        splitter = TokenTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=overlap
        )
        return splitter.split_text(text)
    
    @staticmethod
    def semantic_chunking(text: str, embeddings):
        """语义分块 - 基于语义边界"""
        from langchain_experimental.text_splitter import SemanticChunker
        splitter = SemanticChunker(embeddings)
        return splitter.split_text(text)

---

#### 核心挑战：Lost in the Middle (迷失在中间)

研究表明，当上下文非常长时，模型对**开头**和**结尾**的信息记忆最深刻，而位于**中间**的信息最容易被忽视。

#### 应对策略
1.  **重排序 (Re-ranking)**：将检索到的最相关的文档块放在 Prompt 的最前面或最后面。
2.  **信息密度优化**：在 Prompt 中明确指出：“请特别注意文档中间关于 XXX 的描述”。
3.  **多轮精炼**：不要一次性喂给模型，而是分批处理。

---

#### 长文档处理的三大模式

当文档长度超过模型窗口时，通常采用以下三种模式：

#### 1. Stuffing (填充模式)
将所有内容直接塞进 Prompt。
- **优点**：模型拥有完整上下文，理解最透彻。
- **缺点**：受限于窗口大小，成本高。

#### 2. Map-Reduce (映射-归约模式)
- **Map**：将文档分块，对每一块分别进行总结或提取。
- **Reduce**：将所有块的总结再次汇总，生成最终答案。
- **适用场景**：生成超长文档的摘要。

#### 3. Refine (迭代精炼模式)
- **过程**：先处理第一块，得到初始答案；然后将答案和第二块一起发给模型，要求模型根据新信息“精炼”答案；以此类推。
- **优点**：答案会越来越精准。
- **缺点**：无法并行，速度慢。

```python
# LangChain 中的实现示例
from langchain.chains.summarize import load_summarize_chain

chain = load_summarize_chain(llm, chain_type="map_reduce")
summary = chain.run(docs)

import re

class StructuredChunker:
    """基于文档结构的分块"""
    
    def __init__(self, max_chunk_size: int = 2000):
        self.max_size = max_chunk_size
    
    def chunk_by_sections(self, text: str) -> list:
        """按章节分块"""
        # 识别标题模式
        section_patterns = [
            r'^#{1,6}\s+.+$',           # Markdown 标题
            r'^\d+\.\s+.+$',             # 数字编号
            r'^第[一二三四五六七八九十]+[章节部分]\s*.+$',  # 中文章节
            r'^[A-Z][A-Z\s]+:?\s*$',     # 大写标题
        ]
        
        combined_pattern = '|'.join(f'({p})' for p in section_patterns)
        
        sections = []
        current_section = {"title": "Introduction", "content": ""}
        
        for line in text.split('\n'):
            if re.match(combined_pattern, line, re.MULTILINE):
                if current_section["content"].strip():
                    sections.append(current_section)
                current_section = {"title": line.strip(), "content": ""}
            else:
                current_section["content"] += line + "\n"
        
        if current_section["content"].strip():
            sections.append(current_section)
        
        # 处理过大的章节
        final_chunks = []
        for section in sections:
            if len(section["content"]) > self.max_size:
                # 进一步分割
                sub_chunks = self._split_large_section(section)
                final_chunks.extend(sub_chunks)
            else:
                final_chunks.append(section)
        
        return final_chunks
    
    def _split_large_section(self, section: dict) -> list:
        """分割过大的章节"""
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.max_size,
            chunk_overlap=100
        )
        
        sub_texts = splitter.split_text(section["content"])
        
        return [
            {
                "title": f"{section['title']} (Part {i+1})",
                "content": sub
            }
            for i, sub in enumerate(sub_texts)
        ]

from llmlingua import PromptCompressor

# 初始化压缩器
compressor = PromptCompressor(
    model_name="microsoft/llmlingua-2-bert-base-multilingual-cased-v1",
    use_llmlingua2=True
)

def compress_long_context(context: str, instruction: str, question: str):
    """
    使用 LLMLingua-2 压缩长上下文
    """
    results = compressor.compress_prompt(
        context=[context],
        instruction=instruction,
        question=question,
        target_token=2000,
        rank_method="longllmlingua" # 针对长文本优化
    )
    
    print(f"原始 Token: {results['origin_tokens']}")
    print(f"压缩后 Token: {results['compressed_tokens']}")
    print(f"压缩率: {results['ratio']}")
    
    return results['compressed_prompt']

# Anthropic Context Caching 示例 (伪代码)
import anthropic

client = anthropic.Anthropic()

response = client.beta.prompt_caching.messages.create(
    model="claude-3-5-sonnet-20240620",
    max_tokens=1024,
    system=[
        {
            "type": "text",
            "text": "你是一个法律专家，请基于以下长文档回答问题...",
            "cache_control": {"type": "ephemeral"} # 标记为可缓存
        },
        {
            "type": "text",
            "text": LONG_LEGAL_DOCUMENT, # 超过 100k tokens 的文档
            "cache_control": {"type": "ephemeral"}
        }
    ],
    messages=[{"role": "user", "content": "这份合同的违约条款是什么？"}]
)

class ExtractionCompressor:
    """基于信息抽取的压缩"""
    
    def __init__(self):
        self.llm = ChatOpenAI(model="gpt-4o-mini")
    
    def extract_key_info(self, text: str, query: str) -> dict:
        """抽取关键信息"""
        prompt = f"""从文本中抽取与问题相关的关键信息。

问题：{query}

文本：{text}

请返回 JSON 格式：
key_facts"""
        
        response = self.llm.invoke(prompt)
        
        import json
        try:
            return json.loads(response.content)
        except:
            return {"summary": response.content}
    
    def compress_documents(
        self,
        documents: list,
        query: str
    ) -> str:
        """压缩文档列表"""
        all_facts = []
        all_quotes = []
        all_entities = set()
        
        for doc in documents:
            info = self.extract_key_info(doc.page_content, query)
            all_facts.extend(info.get("key_facts", []))
            all_quotes.extend(info.get("relevant_quotes", []))
            all_entities.update(info.get("entities", []))
        
        # 构建压缩上下文
        compressed = f"""关键事实：
{chr(10).join(['- ' + f for f in all_facts[:10]])}

相关实体：{', '.join(list(all_entities)[:20])}

关键引用：
{chr(10).join(['> ' + q for q in all_quotes[:5]])}"""
        
        return compressed

class MapReduceSummarizer:
    """Map-Reduce 摘要"""
    
    def __init__(self):
        self.llm = ChatOpenAI(model="gpt-4o-mini")
    
    def map_phase(self, chunks: list) -> list:
        """Map 阶段：并行摘要"""
        summaries = []
        
        for chunk in chunks:
            prompt = f"""请为以下文本生成简洁摘要（100字以内）：

{chunk}

摘要："""
            
            response = self.llm.invoke(prompt)
            summaries.append(response.content)
        
        return summaries
    
    def reduce_phase(self, summaries: list) -> str:
        """Reduce 阶段：合并摘要"""
        combined = "\n".join([f"- {s}" for s in summaries])
        
        prompt = f"""将以下多个摘要合并为一个连贯的总结：

{combined}

合并摘要："""
        
        response = self.llm.invoke(prompt)
        return response.content
    
    def summarize(self, document: str, chunk_size: int = 2000) -> str:
        """完整摘要流程"""
        # 分块
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=100
        )
        chunks = splitter.split_text(document)
        
        # Map
        summaries = self.map_phase(chunks)
        
        # Reduce（可能需要多轮）
        while len(summaries) > 5:
            # 分组 reduce
            groups = [summaries[i:i+5] for i in range(0, len(summaries), 5)]
            summaries = [self.reduce_phase(g) for g in groups]
        
        # 最终 reduce
        return self.reduce_phase(summaries)

# 使用
summarizer = MapReduceSummarizer()
summary = summarizer.summarize(very_long_document)

class RefineSummarizer:
    """迭代精炼摘要"""
    
    def __init__(self):
        self.llm = ChatOpenAI(model="gpt-4o")
    
    def summarize(self, document: str, chunk_size: int = 3000) -> str:
        """迭代摘要"""
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=chunk_size,
            chunk_overlap=200
        )
        chunks = splitter.split_text(document)
        
        # 初始摘要
        current_summary = self._initial_summary(chunks[0])
        
        # 迭代精炼
        for chunk in chunks[1:]:
            current_summary = self._refine_summary(current_summary, chunk)
        
        return current_summary
    
    def _initial_summary(self, text: str) -> str:
        """初始摘要"""
        prompt = f"""请为以下文本生成详细摘要：

{text}

摘要："""
        
        response = self.llm.invoke(prompt)
        return response.content
    
    def _refine_summary(self, existing_summary: str, new_text: str) -> str:
        """精炼摘要"""
        prompt = f"""已有摘要：
{existing_summary}

新增内容：
{new_text}

请更新摘要，整合新内容中的重要信息："""
        
        response = self.llm.invoke(prompt)
        return response.content

# 使用
refiner = RefineSummarizer()
summary = refiner.summarize(very_long_document)

class SlidingWindowMemory:
    """滑动窗口记忆"""
    
    def __init__(self, max_tokens: int = 4000):
        self.max_tokens = max_tokens
        self.messages = []
        self.encoder = tiktoken.encoding_for_model("gpt-4")
    
    def add_message(self, role: str, content: str):
        """添加消息"""
        self.messages.append({"role": role, "content": content})
        self._trim_to_limit()
    
    def _count_tokens(self, messages: list) -> int:
        """计算 token 数"""
        total = 0
        for msg in messages:
            total += len(self.encoder.encode(msg["content"])) + 4
        return total
    
    def _trim_to_limit(self):
        """裁剪到限制内"""
        while self._count_tokens(self.messages) > self.max_tokens:
            if len(self.messages) > 2:
                # 保留系统消息和最新消息
                self.messages.pop(1)
            else:
                break
    
    def get_messages(self) -> list:
        """获取消息列表"""
        return self.messages.copy()

class SmartWindowMemory:
    """智能窗口记忆 - 带摘要"""
    
    def __init__(self, max_tokens: int = 4000, summary_threshold: int = 10):
        self.max_tokens = max_tokens
        self.summary_threshold = summary_threshold
        self.messages = []
        self.summary = ""
        self.llm = ChatOpenAI(model="gpt-4o-mini")
    
    def add_message(self, role: str, content: str):
        """添加消息"""
        self.messages.append({"role": role, "content": content})
        
        if len(self.messages) > self.summary_threshold:
            self._create_summary()
    
    def _create_summary(self):
        """创建历史摘要"""
        # 取前半部分消息生成摘要
        to_summarize = self.messages[:len(self.messages)//2]
        
        history = "\n".join([
            f"{m['role']}: {m['content']}"
            for m in to_summarize
        ])
        
        prompt = f"""总结以下对话的要点：

{history}

摘要："""
        
        response = self.llm.invoke(prompt)
        
        if self.summary:
            self.summary = f"{self.summary}\n{response.content}"
        else:
            self.summary = response.content
        
        # 移除已摘要的消息
        self.messages = self.messages[len(self.messages)//2:]
    
    def get_context(self) -> list:
        """获取上下文"""
        context = []
        
        if self.summary:
            context.append({
                "role": "system",
                "content": f"之前对话摘要：{self.summary}"
            })
        
        context.extend(self.messages)
        
        return context

# 使用
memory = SmartWindowMemory(max_tokens=4000)
memory.add_message("user", "你好")
memory.add_message("assistant", "你好！有什么可以帮助你的？")
# ... 更多消息

context = memory.get_context()

class LongDocumentProcessor:
    """长文档处理器"""
    
    def __init__(self, max_context: int = 8000):
        self.max_context = max_context
        self.llm = ChatOpenAI(model="gpt-4o")
    
    def process_document(
        self,
        document: str,
        question: str
    ) -> str:
        """处理长文档回答问题"""
        # 1. 分块
        chunks = self._chunk_document(document)
        
        # 2. 相关性过滤
        relevant_chunks = self._filter_relevant(chunks, question)
        
        # 3. 如果还是太长，进行压缩
        context = self._prepare_context(relevant_chunks, question)
        
        # 4. 生成答案
        return self._generate_answer(context, question)
    
    def _chunk_document(self, document: str) -> list:
        """分块"""
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=1500,
            chunk_overlap=200
        )
        return splitter.split_text(document)
    
    def _filter_relevant(
        self,
        chunks: list,
        question: str,
        top_k: int = 10
    ) -> list:
        """过滤相关块"""
        embeddings = OpenAIEmbeddings()
        
        question_emb = embeddings.embed_query(question)
        chunk_embs = embeddings.embed_documents(chunks)
        
        # 计算相似度
        similarities = []
        for i, emb in enumerate(chunk_embs):
            sim = np.dot(question_emb, emb) / (
                np.linalg.norm(question_emb) * np.linalg.norm(emb)
            )
            similarities.append((i, sim))
        
        # 排序并返回 top_k
        similarities.sort(key=lambda x: x[1], reverse=True)
        
        return [chunks[i] for i, _ in similarities[:top_k]]
    
    def _prepare_context(self, chunks: list, question: str) -> str:
        """准备上下文"""
        context = "\n\n".join(chunks)
        
        encoder = tiktoken.encoding_for_model("gpt-4")
        tokens = len(encoder.encode(context))
        
        if tokens > self.max_context:
            # 压缩
            compressor = ContextCompressor()
            context = compressor.compress_prompt(
                context,
                question,
                target_ratio=self.max_context / tokens
            )
        
        return context
    
    def _generate_answer(self, context: str, question: str) -> str:
        """生成答案"""
        prompt = f"""基于以下文档内容回答问题。

文档内容：
{context}

问题：{question}

答案："""
        
        response = self.llm.invoke(prompt)
        return response.content

# 使用
processor = LongDocumentProcessor(max_context=8000)
answer = processor.process_document(very_long_doc, "文档的主要结论是什么？")

class MiddleAwareRetriever:
    """解决 "Lost in the Middle" 问题"""
    
    def __init__(self):
        self.llm = ChatOpenAI(model="gpt-4o")
    
    def reorder_documents(self, documents: list) -> list:
        """重排文档顺序 - 重要的放两端"""
        if len(documents) <= 2:
            return documents
        
        # 假设已按相关性排序
        # 将最相关的交替放在首尾
        reordered = []
        left = []
        right = []
        
        for i, doc in enumerate(documents):
            if i % 2 == 0:
                left.append(doc)
            else:
                right.append(doc)
        
        # 首部放最相关的
        reordered = left + right[::-1]
        
        return reordered
    
    def query_with_sections(
        self,
        question: str,
        documents: list
    ) -> str:
        """分段查询"""
        # 将文档分组
        group_size = 3
        groups = [
            documents[i:i+group_size]
            for i in range(0, len(documents), group_size)
        ]
        
        partial_answers = []
        
        for i, group in enumerate(groups):
            context = "\n\n".join([d.page_content for d in group])
            
            prompt = f"""基于以下内容（第 {i+1} 部分）回答问题。
如果这部分没有相关信息，回答 "无相关信息"。

内容：
{context}

问题：{question}

回答："""
            
            response = self.llm.invoke(prompt)
            if "无相关信息" not in response.content:
                partial_answers.append(response.content)
        
        # 合并答案
        if not partial_answers:
            return "未找到相关信息"
        
        combined = "\n".join(partial_answers)
        
        final_prompt = f"""综合以下多个回答，生成最终答案：

{combined}

问题：{question}

最终答案："""
        
        response = self.llm.invoke(final_prompt)
        return response.content

def allocate_token_budget(
    total_budget: int,
    query_tokens: int,
    num_documents: int
) -> dict:
    """Token 预算分配"""
    # 预留输出空间
    output_reserve = min(2000, total_budget // 4)
    
    # 可用于上下文
    context_budget = total_budget - query_tokens - output_reserve
    
    # 每个文档的预算
    per_doc = context_budget // num_documents
    
    return {
        "total": total_budget,
        "context": context_budget,
        "per_document": per_doc,
        "output_reserve": output_reserve
    }