class ContextBudget:
    """上下文预算管理器"""

    def __init__(self, total_budget: int):
        self.total_budget = total_budget
        # 固定预算分配（可调整）
        self.allocations = {
            "system_prompt": 0.10,   # 10% 系统指令
            "tools": 0.15,           # 15% 工具定义
            "few_shot": 0.10,        # 10% Few-shot
            "user_input": 0.05,      # 5% 用户输入
            "memory": 0.25,          # 25% 历史记忆（弹性）
            "retrieval": 0.25,       # 25% 检索结果（弹性）
            "state": 0.05,           # 5% 中间状态
            "metadata": 0.01,        # 1% 元数据
            "buffer": 0.04,          # 4% 安全缓冲
        }

    def allocate(self, content_types: dict) -> dict:
        """根据预算分配上下文 token"""
        allocated = {}
        used = 0

        # 第一遍：固定内容
        for ctype in ["system_prompt", "tools", "few_shot", "user_input", "metadata"]:
            budget = int(self.total_budget * self.allocations[ctype])
            content = content_types.get(ctype, "")
            allocated[ctype] = content[:budget]  # 简单截断
            used += len(allocated[ctype])

        # 第二遍：弹性内容按比例分配剩余预算
        remaining = self.total_budget - used - int(self.total_budget * self.allocations["buffer"])
        elastic_total = sum(self.allocations[k] for k in ["memory", "retrieval", "state"])

        for ctype in ["memory", "retrieval", "state"]:
            share = self.allocations[ctype] / elastic_total
            budget = int(remaining * share)
            content = content_types.get(ctype, "")
            allocated[ctype] = content[:budget]

        return allocated

class ContextCompressor:
    """上下文压缩器"""

    COMPRESSION_STRATEGIES = {
        "truncate_oldest": "截断最旧的历史对话",
        "summarize": "LLM 摘要压缩历史对话",
        "semantic_filter": "语义过滤，只保留与当前任务相关的片段",
        "hierarchical_collapse": "层次化折叠，保留摘要、移除细节",
    }

    def compress(self, context: dict, budget: int, strategy: str = "summarize") -> dict:
        """根据策略压缩上下文到目标预算内"""

        current_size = sum(len(v) for v in context.values())
        if current_size <= budget:
            return context

        # 按优先级排序（低优先级先压缩）
        priority_order = ["retrieval", "state", "memory", "few_shot"]

        compressed = dict(context)

        for layer in priority_order:
            if sum(len(v) for v in compressed.values()) <= budget:
                break

            if layer not in compressed:
                continue

            if strategy == "truncate_oldest":
                compressed[layer] = self._truncate_tail(compressed[layer], budget)
            elif strategy == "summarize":
                compressed[layer] = self._summarize_layer(compressed[layer], budget)
            elif strategy == "semantic_filter":
                compressed[layer] = self._filter_by_relevance(
                    compressed[layer],
                    query=compressed.get("user_input", ""),
                    budget=budget,
                )

        return compressed

    def _summarize_layer(self, content: str, target_size: int) -> str:
        """使用 LLM 对某一层内容进行摘要压缩"""
        # 实际实现中调用 LLM API
        summary_prompt = f"""
        Summarize the following content in under {target_size} tokens.
        Preserve key facts, decisions, and action items.
        Remove redundant details and tangential information.

        Content:
        {content}
        """
        # return llm.generate(summary_prompt)
        pass

class BudgetAwareContext:
    """预算感知的上下文管理器"""

    def __init__(self, model_max_tokens: int, safety_margin: float = 0.10):
        self.model_max = model_max_tokens
        self.safety_margin = safety_margin
        # 可用预算 = 总窗口 - 安全缓冲（留给模型输出）
        self.effective_budget = int(model_max_tokens * (1 - safety_margin))

    def build_context(self, components: dict) -> tuple[str, int]:
        """构建不超过预算的上下文"""
        # 组件按优先级排序
        priority = [
            "system_prompt",   # 不可裁剪
            "tools",           # 不可裁剪
            "few_shot",        # 可部分裁剪
            "user_input",      # 不可裁剪
            "memory",          # 可压缩
            "retrieval",       # 可压缩
            "state",           # 可截断
        ]

        # 第一遍：计算固定组件的总大小
        fixed_tokens = 0
        for key in ["system_prompt", "tools", "user_input"]:
            if key in components:
                fixed_tokens += self._count_tokens(components[key])

        if fixed_tokens > self.effective_budget:
            raise ValueError(
                f"固定组件已超出预算: {fixed_tokens} > {self.effective_budget}. "
                "请减少系统指令或工具定义的规模。"
            )

        # 第二遍：分配弹性组件的预算
        remaining = self.effective_budget - fixed_tokens
        elastic_keys = ["few_shot", "memory", "retrieval", "state"]
        elastic_budgets = self._distribute_budget(remaining, elastic_keys, components)

        # 第三遍：裁剪并组装
        context_parts = []
        total_used = 0

        for key in priority:
            if key not in components:
                continue
            if key in elastic_budgets:
                budget = elastic_budgets[key]
                content = self._trim_to_budget(components[key], budget)
            else:
                content = components[key]

            context_parts.append(f"<{key}>\n{content}\n</{key}>")
            total_used += self._count_tokens(content)

        full_context = "\n".join(context_parts)
        return full_context, total_used

    def _distribute_budget(self, total: int, keys: list, components: dict) -> dict:
        """按内容实际大小比例分配弹性预算"""
        sizes = {k: self._count_tokens(components.get(k, "")) for k in keys}
        total_elastic = sum(sizes.values())

        if total_elastic == 0:
            return {k: 0 for k in keys}

        budgets = {}
        for k in keys:
            share = sizes[k] / total_elastic
            budgets[k] = int(total * share)

        return budgets

    def _trim_to_budget(self, content: str, budget: int) -> str:
        """将内容裁剪到预算内（简单的 token 近似：1 token ≈ 1.3 chars for Chinese）"""
        char_budget = int(budget * 1.3)
        if len(content) <= char_budget:
            return content
        return content[:char_budget] + "..."

# ❌ 反模式：扁平、冗余、无结构的上下文

def build_bad_context(user_query: str, chat_history: list, docs: list) -> str:
    """
    问题：
    1. 所有内容拼接为一个扁平字符串，无结构标记
    2. 历史对话全量注入，不做任何过滤或压缩
    3. 检索结果按原始顺序注入，不管相关性排序
    4. 关键指令混在中间，被大量噪声淹没
    5. 没有 token 预算意识，可能超出窗口限制
    """

    context = ""

    # 系统指令被埋在开头，但没有清晰边界
    context += "你是一个有帮助的助手。请根据以下信息回答问题。\n"
    context += "注意：回答要简洁准确。如果不确定，请说不知道。\n"
    context += "另外，请使用中文回答。格式要清晰。\n"
    context += "还有，如果涉及到代码，请使用 Python。\n"

    # 全量注入历史对话——不管是否相关
    for msg in chat_history:  # 可能包含 500+ 轮对话
        context += f"{msg['role']}: {msg['content']}\n"

    # 检索结果直接拼接——没有去重、没有排序
    for doc in docs:  # 可能包含 20+ 个片段，大量重复
        context += doc["text"] + "\n"

    # 用户请求被淹没在大量文本的末尾
    context += f"\n用户问题：{user_query}\n"
    context += "请回答。"

    return context
    # 总大小可能达到 50K+ tokens，但有效密度 < 10%

# ✅ 正模式：层次化、结构化、预算感知的上下文

from dataclasses import dataclass
from typing import Optional


@dataclass
class ContextConfig:
    """上下文配置"""
    model_max_tokens: int = 8192
    safety_margin: float = 0.10
    max_history_turns: int = 5
    max_retrieval_docs: int = 3
    max_retrieval_tokens_per_doc: int = 500


class ContextBuilder:
    """层次化上下文构建器"""

    def __init__(self, config: ContextConfig):
        self.config = config
        self.budget = int(config.model_max_tokens * (1 - config.safety_margin))

    def build(self,
              system_prompt: str,
              tools_schema: Optional[str],
              user_query: str,
              chat_history: list,
              retrieved_docs: list,
              few_shot: Optional[str] = None,
              tool_result: Optional[str] = None) -> str:
        """构建结构化上下文"""

        parts = []

        # ===== Layer 1: 系统层（首区，最高优先级） =====
        parts.append(self._section("system", system_prompt))

        if tools_schema:
            parts.append(self._section("tools", tools_schema))

        if few_shot:
            parts.append(self._section("examples", few_shot))

        # ===== Layer 2: 会话层（中区，弹性） =====
        recent_history = self._select_relevant_history(
            chat_history, user_query, max_turns=self.config.max_history_turns
        )
        if recent_history:
            parts.append(self._section("history", recent_history))

        # ===== Layer 3: 任务层（中-尾区，按需） =====
        if retrieved_docs:
            docs_text = self._format_retrieval(retrieved_docs)
            parts.append(self._section("retrieval", docs_text))

        if tool_result:
            parts.append(self._section("tool_result", tool_result))

        # ===== Layer 4: 用户输入（尾区，最高注意力） =====
        parts.append(self._section("query", user_query))

        # ===== 组装并验证预算 =====
        context = "\n\n".join(parts)
        token_count = self._estimate_tokens(context)

        if token_count > self.budget:
            context = self._compress_to_budget(context, self.budget)

        return context

    def _section(self, name: str, content: str) -> str:
        """带结构标记的段落"""
        return f"<{name}>\n{content.strip()}\n</{name}>"

    def _select_relevant_history(self, history: list, query: str, max_turns: int) -> str:
        """选择与当前查询最相关的历史对话"""
        # 策略 1：取最近的 N 轮（简单但有效）
        recent = history[-max_turns * 2:]  # 每轮 = user + assistant

        # 策略 2（进阶）：语义相似度筛选
        # relevant = semantic_search(query, history, top_k=max_turns)

        return "\n".join(
            f"{msg['role']}: {msg['content']}" for msg in recent
        )

    def _format_retrieval(self, docs: list) -> str:
        """格式化检索结果（去重、排序、截断）"""
        # 去重
        seen = set()
        unique_docs = []
        for doc in docs:
            text = doc["text"][:self.config.max_retrieval_tokens_per_doc]
            if text not in seen:
                seen.add(text)
                unique_docs.append(doc)

        # 按相关性排序并取 top-K
        sorted_docs = sorted(unique_docs, key=lambda d: d.get("score", 0), reverse=True)
        top_docs = sorted_docs[:self.config.max_retrieval_docs]

        return "\n---\n".join(
            f"[文档 {i+1}] (相关性: {doc.get('score', 0):.2f})\n{doc['text']}"
            for i, doc in enumerate(top_docs)
        )

    def _estimate_tokens(self, text: str) -> int:
        """粗略的 token 估算（1 token ≈ 1.3 中文字符）"""
        return int(len(text) / 1.3)

    def _compress_to_budget(self, context: str, budget: int) -> str:
        """压缩到预算内（保留首尾，压缩中间）"""
        # 保留系统层和查询（首尾），压缩中间部分
        # 实际实现中可调用 LLM 摘要
        char_budget = int(budget * 1.3)
        if len(context) <= char_budget:
            return context
        return context[:char_budget] + "\n\n[... 中间内容已压缩以适配上下文窗口 ...]"

"""
Context Building Pipeline — 上下文构建管线实现
每条数据流经 6 个阶段，每个阶段有独立的算子和质量门控。
"""

from dataclasses import dataclass, field
from typing import Any
from enum import Enum
import hashlib
import time
import logging
from collections import deque
from abc import ABC, abstractmethod

logger = logging.getLogger(__name__)


class ContextStage(Enum):
    COLLECT = "collect"
    FILTER = "filter"
    COMPRESS = "compress"
    RANK = "rank"
    INJECT = "inject"
    VALIDATE = "validate"


@dataclass
class ContextItem:
    """管线中流动的基本单元"""
    id: str
    source: str           # 来源：rag / memory / tool / session
    content: str
    relevance_score: float = 0.0
    timestamp: float = 0.0
    metadata: dict = field(default_factory=dict)
    token_count: int = 0

    def __hash__(self):
        return hash(self.id)


@dataclass
class PipelineMetrics:
    """管线各阶段的度量指标"""
    stage: str
    items_in: int = 0
    items_out: int = 0
    duration_ms: float = 0.0
    quality_gate_passed: bool = True
    details: dict = field(default_factory=dict)


class StageOperator(ABC):
    """管线阶段的抽象基类"""
    @abstractmethod
    def execute(self, items: list[ContextItem], params: dict) -> list[ContextItem]:
        ...

    @abstractmethod
    def quality_gate(self, items_in: list[ContextItem], items_out: list[ContextItem]) -> tuple[bool, dict]:
        ...


# ===== 阶段 1：信息收集 =====

class Collector(StageOperator):
    """从多个信息源收集候选上下文项"""

    def __init__(self, sources: list[Any]):
        self.sources = sources  # RAG retriever, MemoryStore, ToolRegistry 等

    def execute(self, items: list[ContextItem], params: dict) -> list[ContextItem]:
        query = params.get("query", "")
        collected = []
        seen_ids = set()

        for source in self.sources:
            try:
                raw_items = source.retrieve(query, **params.get("source_params", {}))
                for item in raw_items:
                    # 去重：相同 content 的 MD5 视为同一项
                    content_hash = hashlib.md5(item.content.encode()).hexdigest()
                    if content_hash not in seen_ids:
                        seen_ids.add(content_hash)
                        ctx_item = ContextItem(
                            id=content_hash[:12],
                            source=source.name,
                            content=item.content,
                            timestamp=time.time(),
                            metadata={"source_name": source.name}
                        )
                        collected.append(ctx_item)
            except Exception as e:
                logger.warning(f"Source {source.name} failed: {e}")

        return collected

    def quality_gate(self, items_in, items_out):
        # 收集阶段：至少有一个源成功
        passed = len(items_out) > 0
        return passed, {"count": len(items_out)}


# ===== 阶段 2：过滤 =====

class Filter(StageOperator):
    """基于多维度策略过滤不相关信息"""

    def __init__(self, min_score: float = 0.3, time_decay_halflife: float = 86400 * 7):
        self.min_score = min_score
        self.halflife = time_decay_halflife

    def execute(self, items: list[ContextItem], params: dict) -> list[ContextItem]:
        now = time.time()
        filtered = []

        for item in items:
            # 1. 语义相关性过滤
            score = self._compute_relevance(item, params)

            # 2. 时间衰减
            age_seconds = now - item.timestamp
            time_weight = 0.5 ** (age_seconds / self.halflife)
            weighted_score = score * time_weight

            # 3. 源可信度
            source_weight = params.get("source_weights", {}).get(item.source, 1.0)
            final_score = weighted_score * source_weight

            if final_score >= self.min_score:
                item.relevance_score = final_score
                filtered.append(item)

        return filtered

    def _compute_relevance(self, item: ContextItem, params: dict) -> float:
        """计算原始相关性得分（简化版，实际可用 embedding 相似度）"""
        query = params.get("query", "")
        if not query:
            return 0.5

        query_terms = set(query.lower().split())
        content_terms = set(item.content.lower().split())
        overlap = query_terms & content_terms

        # Jaccard 相似度
        union = query_terms | content_terms
        return len(overlap) / len(union) if union else 0.0

    def quality_gate(self, items_in, items_out):
        n_in, n_out = len(items_in), len(items_out)
        if n_in == 0:
            return True, {"ratio": 1.0}
        ratio = n_out / n_in
        passed = 0.3 <= ratio <= 0.7
        return passed, {"ratio": ratio, "kept": n_out}


# ===== 阶段 3：压缩 =====

class Compressor(StageOperator):
    """对上下文项进行压缩，减少 token 消耗"""

    def __init__(self, max_tokens_per_item: int = 500, compression_target: float = 0.5):
        self.max_tokens = max_tokens_per_item
        self.target_ratio = compression_target

    def execute(self, items: list[ContextItem], params: dict) -> list[ContextItem]:
        compressed = []
        for item in items:
            original_len = len(item.content)

            if len(item.content.split()) > self.max_tokens:
                # 策略 1：提取关键句（简化版）
                sentences = self._split_sentences(item.content)
                key_sentences = self._extract_key_sentences(sentences, item.relevance_score)
                new_content = " ".join(key_sentences)
            else:
                new_content = item.content

            compression_ratio = len(new_content) / original_len if original_len > 0 else 1.0

            compressed_item = ContextItem(
                id=item.id,
                source=item.source,
                content=new_content,
                relevance_score=item.relevance_score,
                timestamp=item.timestamp,
                metadata={**item.metadata, "compression_ratio": compression_ratio}
            )
            compressed.append(compressed_item)

        return compressed

    def _split_sentences(self, text: str) -> list[str]:
        import re
        return [s.strip() for s in re.split(r'[。！？.!?\n]+', text) if s.strip()]

    def _extract_key_sentences(self, sentences: list[str], relevance: float) -> list[str]:
        """提取关键句：优先保留包含数字、专业术语、结论性词汇的句子"""
        scored = []
        for s in sentences:
            s_score = 0
            # 结论信号
            if any(w in s for w in ['因此', '所以', '结论', '综上', '因此可以', 'in conclusion', 'therefore']):
                s_score += 3
            # 数据信号
            if any(c.isdigit() for c in s):
                s_score += 2
            # 长度信号（太短信息量低，太长可能冗余）
            word_count = len(s.split())
            if 5 <= word_count <= 50:
                s_score += 1

            scored.append((s_score, s))

        scored.sort(reverse=True)
        # 取 top N，保证压缩率
        max_keep = max(1, int(len(sentences) * self.target_ratio))
        return [s for _, s in scored[:max_keep]]

    def quality_gate(self, items_in, items_out):
        if not items_in:
            return True, {}
        ratios = [item.metadata.get("compression_ratio", 1.0) for item in items_out]
        avg_ratio = sum(ratios) / len(ratios) if ratios else 1.0
        passed = avg_ratio <= self.target_ratio
        return passed, {"avg_ratio": avg_ratio}


# ===== 阶段 4：排序 =====

class Ranker(StageOperator):
    """对上下文项进行排序，确保最重要的信息位于最优位置"""

    def execute(self, items: list[ContextItem], params: dict) -> list[ContextItem]:
        # 多维度评分
        for item in items:
            item.relevance_score = self._composite_score(item, params)

        # 主排序：相关性降序
        ranked = sorted(items, key=lambda x: x.relevance_score, reverse=True)

        # 多样性重排：避免连续出现同质化内容
        if params.get("diversity_rerank", False):
            ranked = self._diversity_rerank(ranked, params.get("diversity_lambda", 0.5))

        return ranked

    def _composite_score(self, item: ContextItem, params: dict) -> float:
        base = item.relevance_score
        # 时间衰减加成
        recency_bonus = params.get("recency_weight", 0.1)
        age_hours = (time.time() - item.timestamp) / 3600
        recency = recency_bonus / (1 + age_hours)

        # 位置先验（用户提供的元数据中的优先级）
        prior = item.metadata.get("priority", 0) * 0.05

        return base + recency + prior

    def _diversity_rerank(self, items: list[ContextItem], lam: float) -> list[ContextItem]:
        """MMR（Maximal Marginal Relevance）简化实现"""
        if len(items) <= 1:
            return items

        selected = [items[0]]
        remaining = list(items[1:])

        while remaining:
            best_idx = 0
            best_score = -float('inf')

            for i, item in enumerate(remaining):
                # 与已选项目的最大相似度
                max_sim = max(
                    self._content_similarity(item, sel)
                    for sel in selected
                ) if selected else 0

                score = lam * item.relevance_score - (1 - lam) * max_sim
                if score > best_score:
                    best_score = score
                    best_idx = i

            selected.append(remaining.pop(best_idx))

        return selected

    def _content_similarity(self, a: ContextItem, b: ContextItem) -> float:
        """简化的内容相似度（余弦近似）"""
        words_a = set(a.content.lower().split())
        words_b = set(b.content.lower().split())
        intersection = words_a & words_b
        union = words_a | words_b
        return len(intersection) / len(union) if union else 0.0

    def quality_gate(self, items_in, items_out):
        if not items_out:
            return True, {}
        avg_score = sum(i.relevance_score for i in items_out) / len(items_out)
        passed = avg_score >= 0.7
        return passed, {"avg_score": avg_score}


# ===== 阶段 5：注入 =====

class Injector(StageOperator):
    """将排序后的上下文项注入到模板中"""

    def __init__(self, token_budget: int = 8192, template: str = None):
        self.token_budget = token_budget
        self.template = template or self._default_template()

    def execute(self, items: list[ContextItem], params: dict) -> list[ContextItem]:
        """裁剪以适应 token 预算，返回最终注入的项"""
        total_tokens = 0
        injected = []

        for item in items:
            item_tokens = self._estimate_tokens(item.content)
            if total_tokens + item_tokens <= self.token_budget:
                total_tokens += item_tokens
                injected.append(item)
            else:
                # 预算耗尽，剩余项被裁剪
                logger.info(f"Token budget reached: {total_tokens}/{self.token_budget}, "
                           f"truncated {len(items) - len(injected)} items")
                break

        return injected

    def _estimate_tokens(self, text: str) -> int:
        """粗略的 token 估算（中英文混合）"""
        # 中文字符 ~1 token/char, 英文 ~1 token/0.75 words
        cn_chars = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
        en_words = len([w for w in text.split() if not ('\u4e00' <= w[0] <= '\u9fff' if w else False)])
        return cn_chars + int(en_words * 1.33)

    def render(self, items: list[ContextItem], params: dict) -> str:
        """渲染最终上下文"""
        context_sections = []
        for i, item in enumerate(items, 1):
            context_sections.append(
                f"[参考片段 {i}] (来源: {item.source}, 相关度: {item.relevance_score:.2f})\n"
                f"{item.content}"
            )

        return self.template.format(
            query=params.get("query", ""),
            context="\n\n".join(context_sections),
            n_fragments=len(items)
        )

    @staticmethod
    def _default_template() -> str:
        return (
            "你是一个智能助手。请基于以下参考信息回答用户问题。\n\n"
            "用户问题：{query}\n\n"
            "以下是 {n_fragments} 个参考片段：\n\n"
            "{context}\n\n"
            "请综合以上信息，给出准确、简洁的回答。"
        )

    def quality_gate(self, items_in, items_out):
        total = sum(self._estimate_tokens(i.content) for i in items_out)
        passed = total <= self.token_budget
        return passed, {"total_tokens": total, "budget": self.token_budget}


# ===== 阶段 6：验证 =====

class Validator(StageOperator):
    """验证最终上下文的完整性和安全性"""

    def execute(self, items: list[ContextItem], params: dict) -> list[ContextItem]:
        # 安全检查
        validated = []
        for item in items:
            if self._is_safe(item.content, params):
                validated.append(item)
            else:
                logger.warning(f"Item {item.id} failed safety check, excluded")
        return validated

    def _is_safe(self, content: str, params: dict) -> bool:
        """基础安全检查"""
        blocked_patterns = params.get("blocked_patterns", [])
        return not any(p in content for p in blocked_patterns)

    def quality_gate(self, items_in, items_out):
        return True, {"excluded": len(items_in) - len(items_out)}


# ===== 管线编排器 =====

class ContextPipeline:
    """上下文构建管线的编排器"""

    def __init__(self, stages: dict[ContextStage, StageOperator], token_budget: int = 8192):
        self.stages = stages
        self.token_budget = token_budget
        self.metrics: list[PipelineMetrics] = []

    def build(self, query: str, params: dict = None) -> str:
        """执行完整管线，返回渲染后的上下文字符串"""
        params = params or {}
        params["query"] = query
        items: list[ContextItem] = []
        self.metrics = []

        # 管线执行顺序
        for stage in ContextStage:
            operator = self.stages.get(stage)
            if not operator:
                continue

            start = time.time()
            items_in_count = len(items)

            items = operator.execute(items, params)
            passed, details = operator.quality_gate(
                [ContextItem(id=f"dummy_{i}") for i in range(items_in_count)],
                items
            )

            duration = (time.time() - start) * 1000
            metric = PipelineMetrics(
                stage=stage.value,
                items_in=items_in_count,
                items_out=len(items),
                duration_ms=duration,
                quality_gate_passed=passed,
                details=details
            )
            self.metrics.append(metric)

            if not passed and stage != ContextStage.COLLECT:
                logger.warning(f"Quality gate failed at stage {stage.value}: {details}")

        # 渲染
        injector = self.stages.get(ContextStage.INJECT)
        if injector and hasattr(injector, 'render'):
            return injector.render(items, params)
        return "\n".join(item.content for item in items)

    def get_metrics_report(self) -> str:
        """生成管线执行报告"""
        lines = ["=== Context Pipeline Report ==="]
        total_in = self.metrics[0].items_out if self.metrics else 0
        for m in self.metrics:
            gate = "✅" if m.quality_gate_passed else "❌"
            lines.append(
                f"  [{m.stage:^10}] {m.items_in:4d} → {m.items_out:4d} items  "
                f"| {m.duration_ms:6.1f}ms  | {gate}"
            )
        return "\n".join(lines)

# 构建管线
pipeline = ContextPipeline(
    stages={
        ContextStage.COLLECT:   Collector(sources=[rag_retriever, memory_store]),
        ContextStage.FILTER:    Filter(min_score=0.3, time_decay_halflife=86400*7),
        ContextStage.COMPRESS:  Compressor(max_tokens_per_item=300, compression_target=0.5),
        ContextStage.RANK:      Ranker(),
        ContextStage.INJECT:    Injector(token_budget=4096),
        ContextStage.VALIDATE:  Validator(),
    },
    token_budget=4096
)

# 执行
context = pipeline.build(
    query="如何优化 PostgreSQL 在大规模并发写入场景下的性能？",
    params={
        "source_weights": {"rag": 1.0, "memory": 0.6, "tool": 0.8},
        "diversity_rerank": True,
        "diversity_lambda": 0.6,
        "recency_weight": 0.1,
    }
)

print(pipeline.get_metrics_report())

"""
Context Layer Architecture — 上下文分层架构实现
每层有独立的生命周期、更新策略和预算控制。
"""

from dataclasses import dataclass, field
from typing import Protocol, Any
from abc import ABC, abstractmethod
import json
import time


@dataclass
class LayerBudget:
    """每层的预算配置"""
    min_tokens: int       # 保证最小存在
    max_tokens: int       # 硬性上限
    priority: int         # 优先级（数字越小越高）
    trimmable: bool       # 是否可裁剪


class ContextLayer(ABC):
    """上下文层的抽象基类"""

    def __init__(self, name: str, budget: LayerBudget):
        self.name = name
        self.budget = budget
        self.content: str = ""
        self.token_count: int = 0
        self.updated_at: float = time.time()

    @abstractmethod
    def refresh(self, params: dict) -> None:
        """更新层内容"""
        ...

    @abstractmethod
    def estimate_tokens(self) -> int:
        ...

    def trim_to_budget(self, max_tokens: int) -> None:
        """裁剪内容以适应预算"""
        if self.token_count <= max_tokens:
            return
        # 简单截断策略（实际应使用智能压缩）
        words = self.content.split()
        keep_ratio = max_tokens / self.token_count
        keep_count = max(10, int(len(words) * keep_ratio))
        self.content = " ".join(words[:keep_count]) + " ...[truncated]"
        self.token_count = self.estimate_tokens()

    def to_dict(self) -> dict:
        return {
            "layer": self.name,
            "tokens": self.token_count,
            "content": self.content,
            "updated_at": self.updated_at,
        }


# ===== L0 系统层 =====

class SystemLayer(ContextLayer):
    """L0 系统层：模型角色定义、行为约束、全局安全规则

    这是最稳定的层，通常在会话初始化时设置一次。
    内容应尽可能精确——模糊的系统指令是 prompt injection 的主要入口。
    """

    def __init__(self, budget: LayerBudget = None):
        budget = budget or LayerBudget(
            min_tokens=200, max_tokens=1500, priority=1, trimmable=False
        )
        super().__init__("L0_system", budget)
        self.role = ""
        self.constraints: list[str] = []
        self.safety_rules: list[str] = []

    def configure(self, role: str, constraints: list[str], safety_rules: list[str]):
        self.role = role
        self.constraints = constraints
        self.safety_rules = safety_rules
        self.updated_at = time.time()

    def refresh(self, params: dict) -> None:
        sections = []

        if self.role:
            sections.append(f"## 角色\n{self.role}")

        if self.constraints:
            sections.append("## 行为约束\n" + "\n".join(
                f"- {c}" for c in self.constraints
            ))

        if self.safety_rules:
            sections.append("## 安全规则\n" + "\n".join(
                f"- {r}" for r in self.safety_rules
            ))

        self.content = "\n\n".join(sections)
        self.token_count = self.estimate_tokens()

    def estimate_tokens(self) -> int:
        return len(self.content.encode("utf-8")) // 3  # 粗略估算


# ===== L1 任务层 =====

class TaskLayer(ContextLayer):
    """L1 任务层：当前任务描述、输出格式要求、评估标准

    随每个任务变化。好的任务描述应该包含：
    - 明确的输入/输出契约
    - 可验证的完成标准
    - 格式模板或 Schema
    """

    def __init__(self, budget: LayerBudget = None):
        budget = budget or LayerBudget(
            min_tokens=100, max_tokens=800, priority=1, trimmable=False
        )
        super().__init__("L1_task", budget)
        self.description = ""
        self.output_format = ""
        self.evaluation_criteria: list[str] = []

    def configure(self, description: str, output_format: str = "",
                  evaluation_criteria: list[str] = None):
        self.description = description
        self.output_format = output_format
        self.evaluation_criteria = evaluation_criteria or []
        self.updated_at = time.time()

    def refresh(self, params: dict) -> None:
        sections = []

        sections.append(f"## 当前任务\n{self.description}")

        if self.output_format:
            sections.append(f"## 输出格式\n```\n{self.output_format}\n```")

        if self.evaluation_criteria:
            sections.append("## 评估标准\n" + "\n".join(
                f"- {c}" for c in self.evaluation_criteria
            ))

        self.content = "\n\n".join(sections)
        self.token_count = self.estimate_tokens()

    def estimate_tokens(self) -> int:
        return len(self.content.encode("utf-8")) // 3


# ===== L2 知识层 =====

class KnowledgeLayer(ContextLayer):
    """L2 知识层：RAG 检索结果、领域知识库、实时数据

    这是上下文中的"事实来源"。内容质量直接影响回答的准确性。
    需要配合 RAG 管线，按查询动态构建。
    """

    def __init__(self, budget: LayerBudget = None):
        budget = budget or LayerBudget(
            min_tokens=200, max_tokens=4096, priority=2, trimmable=True
        )
        super().__init__("L2_knowledge", budget)
        self.retriever = None  # 外部检索器
        self.fragments: list[dict] = []

    def set_retriever(self, retriever):
        self.retriever = retriever

    def refresh(self, params: dict) -> None:
        query = params.get("query", "")
        top_k = params.get("top_k", 5)

        if self.retriever and query:
            self.fragments = self.retriever.retrieve(query, top_k=top_k)

        sections = []
        for i, frag in enumerate(self.fragments, 1):
            source = frag.get("source", "unknown")
            content = frag.get("content", "")
            score = frag.get("score", 0.0)
            sections.append(
                f"### 知识片段 {i}\n"
                f"来源: {source} | 相关度: {score:.2f}\n"
                f"{content}"
            )

        self.content = "\n\n".join(sections) if sections else "（无外部知识）"
        self.token_count = self.estimate_tokens()

    def estimate_tokens(self) -> int:
        return len(self.content.encode("utf-8")) // 3


# ===== L3 记忆层 =====

class MemoryLayer(ContextLayer):
    """L3 记忆层：历史对话摘要、用户画像、偏好与习惯

    这是最动态的层。包含：
    - 短期记忆：当前会话的对话历史（滑动窗口）
    - 长期记忆：跨会话的用户偏好和关键事实
    """

    def __init__(self, budget: LayerBudget = None):
        budget = budget or LayerBudget(
            min_tokens=100, max_tokens=3072, priority=3, trimmable=True
        )
        super().__init__("L3_memory", budget)
        self.short_term: list[dict] = []  # 近期对话
        self.long_term: dict = {}          # 持久化记忆
        self.window_size: int = 10         # 滑动窗口大小

    def add_turn(self, role: str, content: str):
        self.short_term.append({"role": role, "content": content, "ts": time.time()})
        # 保持窗口大小
        if len(self.short_term) > self.window_size:
            self.short_term = self.short_term[-self.window_size:]

    def set_long_term(self, memory: dict):
        self.long_term = memory
        self.updated_at = time.time()

    def refresh(self, params: dict) -> None:
        sections = []

        # 短期记忆：最近对话
        if self.short_term:
            st_lines = ["## 近期对话\n"]
            for turn in self.short_term[-self.window_size:]:
                st_lines.append(f"**{turn['role']}**: {turn['content']}")
            sections.append("\n".join(st_lines))

        # 长期记忆：用户画像
        if self.long_term:
            lt_lines = ["## 用户信息\n"]
            for key, value in self.long_term.items():
                lt_lines.append(f"- **{key}**: {value}")
            sections.append("\n".join(lt_lines))

        self.content = "\n\n".join(sections) if sections else "（无记忆数据）"
        self.token_count = self.estimate_tokens()

    def estimate_tokens(self) -> int:
        return len(self.content.encode("utf-8")) // 3


# ===== L4 工具层 =====

class ToolLayer(ContextLayer):
    """L4 工具层：可用工具列表、使用规范、输出约束

    描述模型可以使用的工具，包括：
    - 工具名称和功能描述
    - 参数 Schema（JSON Schema 格式）
    - 使用约束和最佳实践
    """

    def __init__(self, budget: LayerBudget = None):
        budget = budget or LayerBudget(
            min_tokens=100, max_tokens=2048, priority=2, trimmable=True
        )
        super().__init__("L4_tools", budget)
        self.tools: list[dict] = []

    def register_tool(self, name: str, description: str, schema: dict,
                      constraints: list[str] = None):
        self.tools.append({
            "name": name,
            "description": description,
            "schema": schema,
            "constraints": constraints or [],
        })
        self.updated_at = time.time()

    def refresh(self, params: dict) -> None:
        sections = ["## 可用工具\n"]

        for tool in self.tools:
            sections.append(
                f"### {tool['name']}\n"
                f"{tool['description']}\n\n"
                f"**参数 Schema:**\n```json\n"
                f"{json.dumps(tool['schema'], indent=2, ensure_ascii=False)}\n```"
            )
            if tool["constraints"]:
                sections.append(
                    "**约束:**\n" + "\n".join(f"- {c}" for c in tool["constraints"])
                )

        self.content = "\n\n".join(sections)
        self.token_count = self.estimate_tokens()

    def estimate_tokens(self) -> int:
        return len(self.content.encode("utf-8")) // 3


# ===== 分层上下文组装器 =====

class LayeredContextAssembler:
    """组装各层上下文为最终的上下文窗口"""

    def __init__(self, layers: list[ContextLayer], total_budget: int = 8192):
        self.layers = layers
        self.total_budget = total_budget

    def assemble(self, params: dict = None) -> str:
        """按优先级组装上下文，超出预算时裁剪低优先级层"""
        params = params or {}

        # 1. 刷新所有层
        for layer in self.layers:
            layer.refresh(params)

        # 2. 按优先级排序
        sorted_layers = sorted(self.layers, key=lambda l: l.budget.priority)

        # 3. 分配预算
        remaining = self.total_budget
        assigned = {}

        # 先分配不可裁剪层
        for layer in sorted_layers:
            if not layer.budget.trimmable:
                cost = min(layer.token_count, layer.budget.max_tokens)
                assigned[layer.name] = layer.budget.max_tokens
                remaining -= cost
            else:
                assigned[layer.name] = layer.budget.min_tokens
                remaining -= layer.budget.min_tokens

        # 将剩余预算按 min_tokens 比例分配给可裁剪层
        if remaining > 0:
            trimmable = [l for l in sorted_layers if l.budget.trimmable]
            total_min = sum(l.budget.min_tokens for l in trimmable)
            if total_min > 0:
                for layer in trimmable:
                    extra = int(remaining * (layer.budget.min_tokens / total_min))
                    assigned[layer.name] += extra

        # 4. 裁剪 + 渲染
        sections = []
        for layer in sorted_layers:
            max_for_layer = assigned.get(layer.name, layer.budget.min_tokens)
            layer.trim_to_budget(max_for_layer)
            sections.append(layer.content)

        return "\n\n---\n\n".join(sections)

    def get_budget_report(self) -> str:
        """报告各层的 token 使用情况"""
        lines = ["=== Layer Budget Report ==="]
        total = 0
        for layer in self.layers:
            pct = layer.token_count / self.total_budget * 100
            lines.append(
                f"  {layer.name:15s}: {layer.token_count:5d} tokens  ({pct:5.1f}%)"
            )
            total += layer.token_count
        lines.append(f"  {'TOTAL':15s}: {total:5d} tokens  ({total/self.total_budget*100:.1f}%)")
        return "\n".join(lines)

# 1. 构建各层
system_layer = SystemLayer()
system_layer.configure(
    role="你是一位资深的数据库架构师，专注于分布式系统的性能优化。",
    constraints=[
        "回答必须基于事实，不确定时明确说明",
        "涉及代码时，给出可直接运行的示例",
        "避免泛泛而谈，给出具体参数和基准数据",
    ],
    safety_rules=[
        "不得给出未经测试的生产环境配置建议",
        "涉及安全配置时，强调需要安全审计",
    ]
)

task_layer = TaskLayer()
task_layer.configure(
    description="分析用户提供的 PostgreSQL 慢查询日志，识别性能瓶颈并给出优化建议。",
    output_format="JSON: {\"bottleneck\": str, \"root_cause\": str, \"solutions\": list[str]}",
    evaluation_criteria=["准确性", "可操作性", "性能影响预估"],
)

knowledge_layer = KnowledgeLayer()
knowledge_layer.set_retriever(rag_retriever)  # 外部 RAG 检索器

memory_layer = MemoryLayer()
memory_layer.set_long_term({
    "用户角色": "数据库工程师",
    "技术栈": "PostgreSQL 15, Go",
    "偏好": "倾向给出基准测试数据而非经验判断",
})
memory_layer.add_turn("user", "最近写入延迟很高，帮我看一下执行计划")
memory_layer.add_turn("assistant", "请提供慢查询日志...")

tool_layer = ToolLayer()
tool_layer.register_tool(
    name="sql_explain",
    description="获取 SQL 查询的执行计划",
    schema={
        "type": "object",
        "properties": {
            "query": {"type": "string", "description": "SQL 查询语句"},
            "format": {"type": "string", "enum": ["text", "json", "yaml"]},
        },
        "required": ["query"],
    },
    constraints=["仅用于 SELECT 语句分析", "不执行数据修改"],
)

# 2. 组装
assembler = LayeredContextAssembler(
    layers=[system_layer, task_layer, knowledge_layer, memory_layer, tool_layer],
    total_budget=8192
)

context = assembler.assemble(params={"query": "PostgreSQL 写入延迟", "top_k": 5})
print(assembler.get_budget_report())

"""
Context Lifecycle Management — 上下文生命周期管理
实现 TTL 策略、相关性衰减和状态转换。
"""

from dataclasses import dataclass, field
from enum import Enum
import time
import math
import heapq
from typing import Optional


class ContextState(Enum):
    CREATED = "created"
    ACTIVE = "active"
    UPDATING = "updating"
    STALE = "stale"
    ARCHIVED = "archived"


@dataclass
class LifecycleConfig:
    """生命周期配置"""
    half_life_seconds: float     # 半衰期
    initial_ttl: float           # 初始 TTL（秒）
    active_extension: float      # 被引用时延长的 TTL
    stale_threshold: float       # 相关性低于此值进入 STALE
    max_active_count: int        # 最大活跃项数
    archive_ttl: float           # 归档后的保留时间


@dataclass
class ContextLifecycleEntry:
    """生命周期跟踪条目"""
    id: str
    content: str
    state: ContextState = ContextState.CREATED
    created_at: float = field(default_factory=time.time)
    last_accessed_at: float = field(default_factory=time.time)
    reference_count: int = 0
    relevance_score: float = 1.0
    version: int = 1
    metadata: dict = field(default_factory=dict)

    # 衰减计算用
    initial_relevance: float = 1.0
    half_life: float = 604800  # 默认 7 天

    @property
    def age_seconds(self) -> float:
        return time.time() - self.created_at

    @property
    def last_access_age(self) -> float:
        return time.time() - self.last_accessed_at

    def decayed_relevance(self) -> float:
        """计算衰减后的相关性得分"""
        if self.half_life <= 0:
            return self.initial_relevance
        decay_factor = 0.5 ** (self.age_seconds / self.half_life)
        return self.initial_relevance * decay_factor

    def touch(self):
        """被引用：更新访问时间和引用计数"""
        self.last_accessed_at = time.time()
        self.reference_count += 1


class ContextLifecycleManager:
    """上下文生命周期管理器

    核心职责：
    1. 跟踪每个上下文项的生命周期状态
    2. 基于衰减模型计算实时相关性
    3. 自动触发状态转换
    4. 提供淘汰候选列表
    """

    def __init__(self, config: LifecycleConfig):
        self.config = config
        self.entries: dict[str, ContextLifecycleEntry] = {}
        self._active_heap: list[tuple[float, str]] = []  # 最小堆：(-相关性, id)

    def create(self, item_id: str, content: str,
               initial_relevance: float = 1.0,
               half_life: float = None,
               metadata: dict = None) -> ContextLifecycleEntry:
        """创建新的上下文项"""
        half_life = half_life or self.config.half_life_seconds
        entry = ContextLifecycleEntry(
            id=item_id,
            content=content,
            state=ContextState.CREATED,
            initial_relevance=initial_relevance,
            relevance_score=initial_relevance,
            half_life=half_life,
            metadata=metadata or {},
        )
        self.entries[item_id] = entry
        return entry

    def activate(self, item_id: str) -> bool:
        """将项激活：注入上下文窗口"""
        entry = self.entries.get(item_id)
        if not entry:
            return False

        # 状态转换：CREATED/STALE → ACTIVE
        if entry.state in (ContextState.CREATED, ContextState.STALE):
            entry.state = ContextState.ACTIVE
            entry.touch()
            # 推入活跃堆
            heapq.heappush(self._active_heap, (-entry.decayed_relevance(), item_id))
            return True
        return False

    def reference(self, item_id: str) -> bool:
        """项被模型引用：延长 TTL，更新相关性"""
        entry = self.entries.get(item_id)
        if not entry or entry.state != ContextState.ACTIVE:
            return False

        entry.touch()
        # 引用加成：适度提升相关性（避免无限放大）
        entry.initial_relevance = min(1.0, entry.initial_relevance * 1.05)

        # 更新堆中的优先级
        self._update_heap_priority(item_id)
        return True

    def update(self, item_id: str, new_content: str) -> bool:
        """更新上下文项内容"""
        entry = self.entries.get(item_id)
        if not entry:
            return False

        entry.state = ContextState.UPDATING
        entry.content = new_content
        entry.version += 1
        entry.last_accessed_at = time.time()

        # 重新评估相关性
        new_relevance = entry.decayed_relevance()
        if new_relevance < self.config.stale_threshold:
            entry.state = ContextState.STALE
        else:
            entry.state = ContextState.ACTIVE
            entry.relevance_score = new_relevance

        return True

    def tick(self) -> list[str]:
        """推进生命周期：检查所有活跃项，返回需要淘汰的项 ID 列表"""
        to_stale = []
        to_archive = []

        for item_id, entry in list(self.entries.items()):
            if entry.state not in (ContextState.ACTIVE, ContextState.STALE):
                continue

            decayed = entry.decayed_relevance()
            entry.relevance_score = decayed

            if entry.state == ContextState.ACTIVE:
                if decayed < self.config.stale_threshold:
                    entry.state = ContextState.STALE
                    to_stale.append(item_id)

            elif entry.state == ContextState.STALE:
                # STALE 项检查复活或归档
                if entry.reference_count > 0 and entry.last_access_age < 3600:
                    # 最近被引用过 → 复活
                    entry.state = ContextState.ACTIVE
                    entry.initial_relevance = min(1.0, decayed * 1.5)
                    if item_id in to_stale:
                        to_stale.remove(item_id)
                else:
                    # 检查归档条件
                    if (time.time() - entry.last_accessed_at) > self.config.archive_ttl:
                        entry.state = ContextState.ARCHIVED
                        to_archive.append(item_id)

        # 清理活跃堆中的过期项
        self._clean_heap()

        return to_stale, to_archive

    def get_active_entries(self, top_k: int = None) -> list[ContextLifecycleEntry]:
        """获取当前活跃的上下文项，按衰减后相关性排序"""
        active = [
            e for e in self.entries.values()
            if e.state == ContextState.ACTIVE
        ]
        active.sort(key=lambda e: e.decayed_relevance(), reverse=True)

        if top_k:
            active = active[:top_k]
        return active

    def get_eviction_candidates(self, count: int = 1) -> list[ContextLifecycleEntry]:
        """获取淘汰候选：相关性最低的活跃项"""
        active = self.get_active_entries()
        return active[-count:] if len(active) >= count else active

    def archive(self, item_id: str) -> dict:
        """归档上下文项，返回可序列化的归档数据"""
        entry = self.entries.get(item_id)
        if not entry:
            return {}

        entry.state = ContextState.ARCHIVED
        return {
            "id": entry.id,
            "content": entry.content,
            "final_relevance": entry.relevance_score,
            "total_references": entry.reference_count,
            "lifetime_seconds": entry.age_seconds,
            "version": entry.version,
        }

    def summary(self) -> dict:
        """生命周期统计摘要"""
        counts = {s.value: 0 for s in ContextState}
        for entry in self.entries.values():
            counts[entry.state.value] += 1

        avg_relevance = (
            sum(e.decayed_relevance() for e in self.entries.values())
            / len(self.entries)
            if self.entries else 0
        )

        return {
            "total_entries": len(self.entries),
            "state_distribution": counts,
            "avg_relevance": avg_relevance,
            "active_count": counts["active"],
        }

    def _update_heap_priority(self, item_id: str):
        """更新堆中项的优先级（惰性删除 + 重新插入）"""
        self._active_heap = [
            (priority, iid)
            for priority, iid in self._active_heap
            if iid != item_id
        ]
        entry = self.entries.get(item_id)
        if entry and entry.state == ContextState.ACTIVE:
            heapq.heappush(self._active_heap, (-entry.decayed_relevance(), item_id))

    def _clean_heap(self):
        """清理堆中已不活跃或已不存在的项"""
        self._active_heap = [
            (priority, iid)
            for priority, iid in self._active_heap
            if iid in self.entries and self.entries[iid].state == ContextState.ACTIVE
        ]
        heapq.heapify(self._active_heap)

# 配置：技术知识类上下文
config = LifecycleConfig(
    half_life_seconds=86400 * 30,    # 30 天半衰期
    initial_ttl=86400 * 7,            # 7 天初始 TTL
    active_extension=86400 * 3,       # 被引用延长 3 天
    stale_threshold=0.1,              # 相关性 < 10% 进入 STALE
    max_active_count=50,              # 最多 50 个活跃项
    archive_ttl=86400 * 30,           # 淘汰后 30 天归档
)

lcm = ContextLifecycleManager(config)

# 创建不同半衰期的上下文项
lcm.create("tech_doc_1", "PostgreSQL 并发写入优化方案...",
           initial_relevance=0.9, half_life=86400 * 30)
lcm.create("user_pref_1", "用户偏好：倾向给出基准测试数据",
           initial_relevance=0.8, half_life=86400 * 180)
lcm.create("dialog_turn_5", "用户：最近写入延迟很高...",
           initial_relevance=0.7, half_life=86400 * 1)  # 对话半衰期 1 天

# 激活
for item_id in ["tech_doc_1", "user_pref_1", "dialog_turn_5"]:
    lcm.activate(item_id)

# 模拟被引用
lcm.reference("tech_doc_1")
lcm.reference("user_pref_1")

# 推进生命周期（模拟时间流逝）
import time
time.sleep(0.1)  # 实际应用中会是真实时间
to_stale, to_archive = lcm.tick()

print(f"进入 STALE: {to_stale}")
print(f"进入 ARCHIVED: {to_archive}")
print(lcm.summary())

"""
Context Compression — 上下文压缩实现
包含：LLMLingua 风格压缩、摘要压缩、关键事件提取
"""

from dataclasses import dataclass, field
from typing import Optional
import re
import json
import time


@dataclass
class CompressionResult:
    original_tokens: int
    compressed_tokens: int
    compression_ratio: float
    content: str
    metadata: dict = field(default_factory=dict)


# ===== LLMLingua 风格压缩 =====

class TokenImportanceCompressor:
    """基于 token 重要性的压缩器（LLMLingua 简化版）

    核心思想：用启发式方法估算每个 token 的信息量，
    保留信息量高的 token，删除信息量低的 token。

    真实 LLMLingua 使用小型 LLM 计算条件概率，
    这里用启发式规则模拟。
    """

    def __init__(self, target_ratio: float = 0.3):
        self.target_ratio = target_ratio  # 目标压缩比例

    def compress(self, text: str) -> CompressionResult:
        original_tokens = self._tokenize(text)
        original_len = len(original_tokens)

        # 1. 计算每个 token 的重要性分数
        scored = []
        for i, token in enumerate(original_tokens):
            score = self._importance_score(token, i, original_tokens)
            scored.append((score, token))

        # 2. 按重要性排序
        scored.sort(key=lambda x: x[0], reverse=True)

        # 3. 保留 top N
        keep_count = max(1, int(original_len * self.target_ratio))
        kept = scored[:keep_count]

        # 4. 恢复原始顺序
        kept_with_indices = []
        for score, token in kept:
            idx = original_tokens.index(token)
            # 处理重复 token
            while (idx, token) in [(t[0], t[1]) for t in kept_with_indices]:
                idx += 1
                if idx >= len(original_tokens):
                    break
            kept_with_indices.append((idx, token))

        kept_with_indices.sort(key=lambda x: x[0])
        compressed_tokens = [t for _, t in kept_with_indices]

        compressed_text = " ".join(compressed_tokens)
        actual_ratio = len(compressed_tokens) / original_len

        return CompressionResult(
            original_tokens=original_len,
            compressed_tokens=len(compressed_tokens),
            compression_ratio=actual_ratio,
            content=compressed_text,
            metadata={"method": "token_importance"}
        )

    def _tokenize(self, text: str) -> list[str]:
        """简单分词（中英文混合）"""
        # 中文按字切分，英文按词切分
        tokens = []
        current = ""
        for char in text:
            if '\u4e00' <= char <= '\u9fff':
                if current:
                    tokens.extend(current.split())
                    current = ""
                tokens.append(char)
            else:
                current += char
        if current:
            tokens.extend(current.split())
        return tokens

    def _importance_score(self, token: str, position: int, all_tokens: list[str]) -> float:
        """启发式 token 重要性评分"""
        score = 1.0

        # 1. 稀有词更重要（TF 逆）
        tf = all_tokens.count(token) / len(all_tokens)
        score += (1.0 - tf) * 2  # 稀有词加分

        # 2. 数字、专有名词更重要
        if token.isdigit() or (token and token[0].isupper()):
            score += 1.5

        # 3. 标点符号重要性低（但保留结构标点）
        if token in '，。！？,.!?':
            score *= 0.3
        elif token in '，、:;(){}[]':
            score *= 0.7

        # 4. 停用词惩罚
        stopwords = {'的', '了', '是', '在', '我', '有', '和', '就', '不', '人', '都',
                     'the', 'a', 'an', 'is', 'are', 'was', 'were', 'of', 'to', 'in'}
        if token.lower() in stopwords:
            score *= 0.4

        # 5. 位置加权：开头和结尾的 token 更重要
        n = len(all_tokens)
        if position < n * 0.1 or position > n * 0.9:
            score *= 1.3

        return score


# ===== 摘要压缩 =====

class SummaryCompressor:
    """基于规则的摘要压缩器

    实际应用中应调用 LLM API 生成摘要，
    这里提供启发式摘要算法作为基础实现。
    """

    def __init__(self, target_length: int = 300):
        self.target_length = target_length

    def compress(self, text: str) -> CompressionResult:
        original_len = len(text.split())

        # 1. 分句
        sentences = self._split_sentences(text)

        # 2. 句子评分
        scored_sentences = []
        for i, sent in enumerate(sentences):
            score = self._sentence_score(sent, i, sentences)
            scored_sentences.append((score, sent, i))

        # 3. 贪心选择：选 top 句子直到达到目标长度
        scored_sentences.sort(key=lambda x: x[0], reverse=True)

        selected = []
        current_len = 0
        for score, sent, orig_idx in scored_sentences:
            sent_len = len(sent.split())
            if current_len + sent_len <= self.target_length:
                selected.append((orig_idx, sent))
                current_len += sent_len
            else:
                break

        # 4. 恢复原始顺序
        selected.sort(key=lambda x: x[0])
        summary = " ".join(sent for _, sent in selected)

        return CompressionResult(
            original_tokens=original_len,
            compressed_tokens=len(summary.split()),
            compression_ratio=len(summary.split()) / original_len if original_len > 0 else 1.0,
            content=summary,
            metadata={"method": "extractive_summary", "sentences_used": len(selected)}
        )

    def _split_sentences(self, text: str) -> list[str]:
        """分句"""
        import re
        return [
            s.strip()
            for s in re.split(r'([。！？.!?\n]+)', text)
            if s.strip() and not re.match(r'^[。！？.!?\n]+$', s)
        ]

    def _sentence_score(self, sentence: str, index: int, all_sentences: list[str]) -> float:
        """句子重要性评分"""
        score = 0.0

        # 1. 包含关键词加分
        keywords = {'重要', '关键', '结论', '因此', '结果', '发现', '建议',
                    'important', 'key', 'conclusion', 'result', 'find'}
        words = set(sentence.lower().split())
        score += len(words & keywords) * 2

        # 2. 包含数字/数据加分
        if any(c.isdigit() for c in sentence):
            score += 1.5

        # 3. 位置：首句和末句通常更重要
        n = len(all_sentences)
        if index == 0:
            score += 2
        elif index == n - 1:
            score += 1.5

        # 4. 长度适中的句子更可能是关键信息
        word_count = len(sentence.split())
        if 10 <= word_count <= 40:
            score += 1

        return score


# ===== 关键事件提取 =====

@dataclass
class ConversationEvent:
    """对话中的关键事件"""
    event_type: str       # "question", "decision", "issue", "resolution"
    turn_index: int
    speaker: str
    summary: str
    timestamp: float = 0.0
    related_turns: list[int] = field(default_factory=list)


class ConversationEventExtractor:
    """从对话中提取关键事件

    将长对话压缩为离散事件列表，大幅减少 token 消耗。
    """

    def __init__(self, max_events: int = 5):
        self.max_events = max_events

    def extract(self, conversation: list[dict]) -> list[ConversationEvent]:
        """将对话列表压缩为关键事件列表

        Args:
            conversation: [{"role": "user/assistant", "content": "...", "ts": ...}]
        Returns:
            关键事件列表
        """
        events = []
        i = 0

        while i < len(conversation):
            turn = conversation[i]
            content = turn.get("content", "")

            # 检测问题事件
            if turn["role"] == "user" and self._is_question(content):
                event = self._extract_question_event(conversation, i)
                events.append(event)
                i = event.related_turns[-1] + 1 if event.related_turns else i + 1

            # 检测决策事件
            elif self._contains_decision(content):
                event = ConversationEvent(
                    event_type="decision",
                    turn_index=i,
                    speaker=turn["role"],
                    summary=self._summarize(content, max_words=30),
                    timestamp=turn.get("ts", 0),
                    related_turns=[i],
                )
                events.append(event)
                i += 1

            # 检测问题解决
            elif self._contains_resolution(content):
                event = ConversationEvent(
                    event_type="resolution",
                    turn_index=i,
                    speaker=turn["role"],
                    summary=self._summarize(content, max_words=30),
                    timestamp=turn.get("ts", 0),
                    related_turns=[i],
                )
                events.append(event)
                i += 1

            else:
                i += 1

        # 按重要性排序，保留 top K
        events.sort(key=lambda e: self._event_priority(e.event_type), reverse=True)
        return events[:self.max_events]

    def _is_question(self, text: str) -> bool:
        return any(p in text for p in ['？', '?', '什么', '如何', '怎么', '为什么',
                                        'what', 'how', 'why', 'which'])

    def _contains_decision(self, text: str) -> bool:
        return any(w in text for w in ['决定', '选择', '确认', '采用', '方案',
                                        'decide', 'choose', 'adopt', 'confirm'])

    def _contains_resolution(self, text: str) -> bool:
        return any(w in text for w in ['解决', '修复', '完成', '已处理',
                                        'resolve', 'fix', 'done', 'completed'])

    def _extract_question_event(self, conversation: list[dict],
                                 q_index: int) -> ConversationEvent:
        """提取一个问题及其相关回复"""
        question = conversation[q_index]["content"]
        related = [q_index]

        # 查找后续回复（直到下一个用户问题或对话结束）
        for j in range(q_index + 1, min(q_index + 5, len(conversation))):
            if conversation[j]["role"] == "user" and self._is_question(conversation[j]["content"]):
                break
            related.append(j)

        # 压缩为摘要
        all_content = " ".join(conversation[j]["content"] for j in related)
        summary = self._summarize(all_content, max_words=50)

        return ConversationEvent(
            event_type="question",
            turn_index=q_index,
            speaker="user",
            summary=f"Q: {self._summarize(question, max_words=20)}\nA: {summary}",
            timestamp=conversation[q_index].get("ts", 0),
            related_turns=related,
        )

    def _summarize(self, text: str, max_words: int = 30) -> str:
        words = text.split()
        if len(words) <= max_words:
            return text
        # 取前 N 个词 + 最后一句
        prefix = " ".join(words[:max_words - 3])
        return prefix + " ..."

    def _event_priority(self, event_type: str) -> int:
        priority = {
            "decision": 5,
            "resolution": 4,
            "question": 3,
            "issue": 2,
        }
        return priority.get(event_type, 1)

    def render_events(self, events: list[ConversationEvent]) -> str:
        """将事件列表渲染为可读文本"""
        lines = ["## 对话关键事件摘要\n"]
        for event in events:
            icon = {"question": "❓", "decision": "📌", "resolution": "✅", "issue": "⚠️"}
            lines.append(f"{icon.get(event.event_type, '•')} [{event.event_type}] {event.summary}")
        return "\n".join(lines)