# -*- coding: utf-8 -*-
"""
可插拔上下文压缩引擎

定义 ContextEngine 抽象基类，允许不同压缩策略替换默认实现。
默认 SimpleTruncationEngine 在接近 token 上限时按 FIFO 方式截断旧消息，
始终保留 system prompt 和最近 N 条消息。

设计借鉴 Hermes Agent 的 agent/context_engine.py：
- 基类定义 lifecycle + compress 接口
- 默认实现是可替换的
- 配置驱动选择引擎
"""

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional


class ContextEngine(ABC):
    """上下文压缩引擎基类。所有引擎必须实现此接口。"""

    # -- 身份标识 --

    # name 为具体属性，子类在 __init__ 中赋值，
    # 不用 @abstractmethod 避免与类属性赋值冲突。
    name: str = ""

    # -- Token 状态 --

    last_prompt_tokens: int = 0
    last_completion_tokens: int = 0
    last_total_tokens: int = 0
    threshold_tokens: int = 0
    context_length: int = 0
    compression_count: int = 0

    # -- 压缩参数 --

    threshold_percent: float = 0.75  # 超过 context * 75% 触发压缩
    protect_first_n: int = 1         # 始终保护前 N 条非 system 消息
    protect_last_n: int = 6          # 始终保护最后 N 条消息

    # -- 核心接口 --

    @abstractmethod
    def update_from_response(self, usage: Dict[str, Any]) -> None:
        """从 API 响应更新 token 使用量。

        每次 LLM 调用后调用，传入 usage dict：
        - prompt_tokens / completion_tokens / total_tokens
        """

    @abstractmethod
    def should_compress(self, prompt_tokens: Optional[int] = None) -> bool:
        """返回 True 表示本回合应触发压缩。"""

    @abstractmethod
    def compress(
        self,
        messages: List[Dict[str, Any]],
        current_tokens: Optional[int] = None,
        focus_topic: Optional[str] = None,
        aggressive: bool = False,
        emergency: bool = False,
    ) -> List[Dict[str, Any]]:
        """压缩消息列表，返回压缩后的消息列表。

        Args:
            messages: 当前完整消息列表
            current_tokens: 当前估算 token 数
            focus_topic: 可选，优先保留与此主题相关的上下文
            aggressive: Phase 1 激进压缩，减少保护条数以删除更多中间消息
            emergency: Phase 2 紧急压缩，仅保留 system + 最近 1 轮对话。
                aggressive 与 emergency 同时为 True 时以 emergency 优先。
        """

    # -- 可选：会话生命周期 --

    def on_session_start(self, session_id: str, **kwargs) -> None:
        """新会话开始时调用。"""

    def on_session_end(self, session_id: str) -> None:
        """会话结束时调用。"""

    def on_session_reset(self) -> None:
        """会话重置时调用（/clear）。重置压缩计数和 token 跟踪。"""
        self.last_prompt_tokens = 0
        self.last_completion_tokens = 0
        self.last_total_tokens = 0
        self.compression_count = 0

    # -- 可选：状态查询 --

    def get_status(self) -> Dict[str, Any]:
        """返回状态字典供显示/日志使用。"""
        last_prompt = self.last_prompt_tokens if self.last_prompt_tokens > 0 else 0
        return {
            "last_prompt_tokens": last_prompt,
            "threshold_tokens": self.threshold_tokens,
            "context_length": self.context_length,
            "usage_percent": (
                min(100, last_prompt / self.context_length * 100)
                if self.context_length else 0
            ),
            "compression_count": self.compression_count,
        }

    def update_model(self, context_length: int) -> None:
        """模型切换时更新 context_length 并重新计算阈值。"""
        self.context_length = context_length
        self.threshold_tokens = int(context_length * self.threshold_percent)


class TokenCounter:
    """使用 tiktoken 精确计算消息列表的 token 数。

    借鉴 zhikuncode-main 的 TokenizerService，替换基于消息条数的估算。
    缓存模型 encoder，避免重复加载；懒加载，构造时不触发 import。
    """

    def __init__(self, model: str = "cl100k_base"):
        """初始化，指定编码模型。延迟到首次计数时才加载 encoder。"""
        self._model = model
        self._encoder = None  # 懒加载，首次使用时才赋值

    def _get_encoder(self):
        """懒加载 tiktoken encoder，缓存复用避免重复加载。"""
        if self._encoder is None:
            import tiktoken
            self._encoder = tiktoken.get_encoding(self._model)
        return self._encoder

    def count_text(self, text: str) -> int:
        """计算单段文本的 token 数。空文本返回 0。"""
        if not text:
            return 0
        return len(self._get_encoder().encode(text))

    def count_messages(self, messages: list) -> int:
        """计算消息列表的总 token 数（含 role 字段开销）。

        每条消息按 4 token role 开销 + role 字段 + content 字段累加，
        参考 OpenAI 官方约定的近似估算方式。
        """
        total = 0
        for msg in messages:
            role = msg.get("role", "")
            content = msg.get("content", "") or ""
            # 每条消息约 4 token 的 role 开销（OpenAI 约定）
            total += 4 + self.count_text(role) + self.count_text(content)
        return total


class CompactEffectivenessValidator:
    """压缩效果验证器（借鉴 zhikuncode-main CompactEffectivenessValidator）。

    70% 安全水位：压缩后 token 数 <= effective_window × 0.70 才算有效。
    - effective_window = context_window - min(max_output_tokens, 20000)
    - 窗口非正时回退到 max(context_window // 2, 10000)
    - tokens_to_release 计算还需释放多少 token 才能达安全水位

    与 SimpleTruncationEngine 配合形成"压缩→验证→不够再压缩"闭环：
    1. 普通压缩后调用 is_effective
    2. 若无效，自动升级到 aggressive 压缩
    3. 仍无效则降级到 emergency 压缩
    4. 与 413 两阶段恢复配合：Phase 1 后验证，无效则直接进 Phase 2
    """

    SAFETY_RATIO = 0.70
    MAX_OUTPUT_CAP = 20000
    FALLBACK_MIN = 10000

    def _effective_window(self, context_window: int, max_output_tokens: int) -> int:
        """计算有效上下文窗口（扣除输出预留后的可用输入空间）。"""
        effective = context_window - min(max_output_tokens, self.MAX_OUTPUT_CAP)
        if effective <= 0:
            effective = max(context_window // 2, self.FALLBACK_MIN)
        return effective

    def is_effective(
        self,
        tokens_after: int,
        context_window: int,
        max_output_tokens: int,
    ) -> bool:
        """判断压缩后 token 数是否低于 70% 安全水位。"""
        if tokens_after <= 0:
            return True
        effective = self._effective_window(context_window, max_output_tokens)
        return tokens_after <= effective * self.SAFETY_RATIO

    def tokens_to_release(
        self,
        tokens_after: int,
        context_window: int,
        max_output_tokens: int,
    ) -> int:
        """计算还需释放多少 token 才能达到安全水位。

        返回 0 表示已有效，返回正数表示需要再压缩多少 token。
        永不返回负数。
        """
        if tokens_after <= 0:
            return 0
        effective = self._effective_window(context_window, max_output_tokens)
        safety = int(effective * self.SAFETY_RATIO)
        release = tokens_after - safety
        return max(0, release)


class SimpleTruncationEngine(ContextEngine):
    """简单截断引擎：在接近 token 上限时按 FIFO 截断旧消息。

    规则：
    1. system prompt 始终保留（第一条消息）
    2. 保留最后 protect_last_n 条消息
    3. 保留前 protect_first_n 条非 system 消息
    4. 中间部分丢弃，插入一条 role='user' 的摘要消息标记截断
    """

    def __init__(
        self,
        context_length: int = 128000,
        threshold_percent: float = 0.75,
        protect_first_n: int = 1,
        protect_last_n: int = 8,
    ):
        self.name = "simple"
        self.context_length = context_length
        self.threshold_percent = threshold_percent
        self.protect_first_n = protect_first_n
        self.protect_last_n = protect_last_n
        self.threshold_tokens = int(context_length * threshold_percent)
        # 精确 token 计数器，用于无 API usage 时本地估算
        self._token_counter = TokenCounter()
        # 当前消息列表快照，供 should_compress 本地估算使用
        self._messages: List[Dict[str, Any]] = []
        # 压缩效果验证器，用于"压缩→验证→不够再压缩"闭环
        self.validator = CompactEffectivenessValidator()
        # 最近一次压缩的升级级别：none/aggressive/emergency
        self.last_escalation_level: str = "none"

    def update_from_response(self, usage: Dict[str, Any]) -> None:
        """更新 token 使用量。兼容 OpenAI usage 格式。"""
        self.last_prompt_tokens = usage.get("prompt_tokens", 0)
        self.last_completion_tokens = usage.get("completion_tokens", 0)
        self.last_total_tokens = usage.get("total_tokens", 0)

    def estimate_tokens(self, messages: list) -> int:
        """本地精确估算消息列表的 token 数（无 API usage 时使用）。"""
        return self._token_counter.count_messages(messages)

    def should_compress(self, prompt_tokens: Optional[int] = None) -> bool:
        """当 token 数超过阈值时触发压缩。

        优先级：显式传入 > API usage > 本地 tiktoken 估算。
        修复原基于消息条数估算导致长消息单条超限未触发压缩的问题。
        """
        if prompt_tokens is not None:
            tokens = prompt_tokens
        elif self.last_prompt_tokens > 0:
            tokens = self.last_prompt_tokens
        else:
            # 无 API usage 时，本地精确估算当前消息列表
            tokens = self.estimate_tokens(self._messages)
        if tokens <= 0 or self.threshold_tokens <= 0:
            return False
        return tokens >= self.threshold_tokens

    def compress(
        self,
        messages: List[Dict[str, Any]],
        current_tokens: Optional[int] = None,
        focus_topic: Optional[str] = None,
        aggressive: bool = False,
        emergency: bool = False,
        max_output_tokens: Optional[int] = None,
        validate_effectiveness: bool = False,
    ) -> List[Dict[str, Any]]:
        """FIFO 截断：保留 head + tail，中间丢弃。

        始终返回新列表，不修改传入的 messages。

        Args:
            aggressive: Phase 1 激进压缩，保护条数减半以删除更多中间消息
            emergency: Phase 2 紧急压缩，仅保留 system + 最近 1 轮对话
            max_output_tokens: 输出 token 预留（用于计算有效输入窗口）
            validate_effectiveness: 压缩后用 CompactEffectivenessValidator 验证
                70% 安全水位，无效则自动升级 aggressive → emergency。
                需同时提供 max_output_tokens 才生效。
        """
        # Phase 2 紧急压缩优先级最高
        if emergency:
            self.last_escalation_level = "emergency"
            return self._compress_emergency(messages)

        # 根据模式决定保护条数：aggressive 减少保护以删除更多中间消息
        protect_first_n = 0 if aggressive else self.protect_first_n
        protect_last_n = max(2, self.protect_last_n // 2) if aggressive else self.protect_last_n

        if len(messages) <= protect_first_n + protect_last_n + 2:
            self.last_escalation_level = "none"
            return list(messages)  # 太小，没必要压缩

        # 确保 system prompt 始终在第一位
        system_msgs = [m for m in messages if m["role"] == "system"]
        non_system = [m for m in messages if m["role"] != "system"]

        if len(non_system) <= protect_first_n + protect_last_n:
            self.last_escalation_level = "none"
            return list(messages)

        head = non_system[:protect_first_n]
        tail = non_system[-protect_last_n:]
        dropped = len(non_system) - protect_first_n - protect_last_n

        # 构建截断标记消息
        if focus_topic:
            truncation_note = (
                f"[上下文压缩] 已移除 {dropped} 条中间消息以控制 token 用量。"
                f"当前焦点：{focus_topic}"
            )
        else:
            truncation_note = (
                f"[上下文压缩] 已移除 {dropped} 条中间消息以控制 token 用量。"
            )

        result = system_msgs + head + [
            {"role": "user", "content": truncation_note}
        ] + tail

        self.compression_count += 1

        # 压缩效果验证：若无效则逐级升级（aggressive → emergency）
        if validate_effectiveness and max_output_tokens is not None and not aggressive:
            result = self._maybe_escalate(
                result,
                messages,
                focus_topic=focus_topic,
                max_output_tokens=max_output_tokens,
            )
        elif aggressive:
            # aggressive 模式压缩后仍验证，无效则降级到 emergency
            if validate_effectiveness and max_output_tokens is not None:
                if not self._is_compress_effective(result, max_output_tokens):
                    self.last_escalation_level = "emergency"
                    result = self._compress_emergency(messages)
                else:
                    self.last_escalation_level = "aggressive"
            else:
                self.last_escalation_level = "aggressive"
        else:
            self.last_escalation_level = "none"

        return result

    def _is_compress_effective(
        self,
        compressed: List[Dict[str, Any]],
        max_output_tokens: int,
    ) -> bool:
        """估算压缩后 token 数，判断是否低于 70% 安全水位。"""
        tokens_after = self.estimate_tokens(compressed)
        return self.validator.is_effective(
            tokens_after=tokens_after,
            context_window=self.context_length,
            max_output_tokens=max_output_tokens,
        )

    def _maybe_escalate(
        self,
        compressed: List[Dict[str, Any]],
        original_messages: List[Dict[str, Any]],
        *,
        focus_topic: Optional[str] = None,
        max_output_tokens: int,
    ) -> List[Dict[str, Any]]:
        """压缩后仍超安全水位时逐级升级。

        1. 普通压缩已有效 -> 保持原结果，level=none
        2. 无效 -> 升级到 aggressive，重新压缩
        3. aggressive 仍无效 -> 降级到 emergency
        """
        if self._is_compress_effective(compressed, max_output_tokens):
            self.last_escalation_level = "none"
            return compressed

        # 升级到 aggressive：保护条数减半
        self.last_escalation_level = "aggressive"
        aggressive_result = self.compress(
            original_messages,
            focus_topic=focus_topic,
            aggressive=True,
            emergency=False,
        )

        if self._is_compress_effective(aggressive_result, max_output_tokens):
            return aggressive_result

        # 仍无效 -> 降级到 emergency
        self.last_escalation_level = "emergency"
        return self._compress_emergency(original_messages)

    def _compress_emergency(
        self,
        messages: List[Dict[str, Any]],
    ) -> List[Dict[str, Any]]:
        """Phase 2 紧急压缩：仅保留 system + 最近 1 轮对话（最后 2 条非 system）。

        场景：413 两阶段恢复的 Phase 2，Phase 1 激进压缩后仍超长时兜底。
        始终返回新列表，不修改传入的 messages。
        """
        system_msgs = [m for m in messages if m["role"] == "system"]
        non_system = [m for m in messages if m["role"] != "system"]
        # 保留最近 1 轮对话（user + assistant = 2 条）
        tail = non_system[-2:] if len(non_system) >= 2 else list(non_system)
        dropped = len(non_system) - len(tail)
        note = {
            "role": "user",
            "content": f"[紧急压缩] 已移除 {dropped} 条历史消息（Phase 2）",
        }
        self.compression_count += 1
        return system_msgs + [note] + tail


def create_engine(name: str = "simple", context_length: int = 128000) -> ContextEngine:
    """根据名称创建上下文引擎实例。"""
    engines = {
        "simple": SimpleTruncationEngine,
    }
    cls = engines.get(name)
    if cls is None:
        raise ValueError(f"未知上下文引擎: {name}，可用: {', '.join(engines)}")
    return cls(context_length=context_length)
