"""上下文计量与压缩检测。

## 为什么换掉 1.5 的做法

1.5 用 tiktoken 逐条消息**本地估算** token，再拿估算值去做限额与压缩判断。
估算与实际用量必然有偏差，偏差还会随消息结构（工具结果、结构化 content 块）
变化——1.5 为此专门开了 ``fix/context-usage-accounting`` 分支，并留了
``CTX_USAGE_V2`` 开关切换口径。根子在于：**用量是 provider 算的，本地再怎么
估都只是近似，不该让估算值参与决策。**

## 2.0 的做法

**真实 usage 优先**：从 ``AIMessage.usage_metadata``（langchain 统一后的
字段）累加 provider 返回的真实 token 数；只有 provider 完全不返回 usage 时
才回落到本地估算，并且在用量流水上标注 ``usage_source=estimate``——口径
不混，事后能分清哪些数字是可信的。

两个不同的量，不要混淆：
- ``input_tokens`` 累加值 = 本轮请求**计费**口径（多次模型调用之和）
- ``context_tokens`` = 最近一次调用的 input_tokens = **当前上下文占用**，
  压缩判断用它

## 压缩检测

不依赖框架钩子，而是看两个客观信号：
1. 流里出现 ``lc_source == "summarization"`` 的消息（框架的摘要标注）
2. 上下文占用**断崖式下跌**（降到上一次的 60% 以下）——压缩把历史换成了摘要

同时把这个 agent 见过的消息文本留在缓冲区里，压缩发生时由上层把缓冲区
落盘存档，这样「被压缩掉的原文」还能追溯。1.5 的压缩默认关着、也没有
原文存档，出了问题无法复盘。
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Iterable

from obs import get_logger

logger = get_logger(__name__)

# 占用跌破上次的 60% 视为发生压缩
_COMPACTION_DROP_RATIO = 0.6

try:  # tiktoken 是可选依赖：装了就用来估算，没装走字符启发式
    import tiktoken

    _ENCODER = tiktoken.get_encoding("cl100k_base")
except Exception:  # pragma: no cover - 取决于环境
    _ENCODER = None


def estimate_tokens(text: str) -> int:
    """估算一段文本的 token 数。仅在 provider 不返回 usage 时作为兜底。"""
    if not text:
        return 0
    if _ENCODER is not None:
        try:
            return len(_ENCODER.encode(text, disallowed_special=()))
        except Exception:  # pragma: no cover
            pass
    # 无 tiktoken：中日韩字符按 1.5 token/字、其余按 4 字符/token 粗估。
    # 宁可高估——低估会导致该压缩时不压缩，直接撑爆模型窗口。
    cjk = sum(1 for ch in text if "\u4e00" <= ch <= "\u9fff" or "\u3000" <= ch <= "\u303f")
    other = len(text) - cjk
    return int(cjk * 1.5 + other / 4) + 1


def extract_usage(message: Any) -> dict[str, int] | None:
    """从 AIMessage 里取 provider 返回的真实 usage。

    依次尝试 langchain 统一的 ``usage_metadata`` 与各 provider 的
    ``response_metadata`` 原始字段。取不到返回 None（**不返回 0**——
    「用量为 0」与「没有用量信息」必须能区分，否则估算兜底会失效）。
    """
    meta = getattr(message, "usage_metadata", None)
    if isinstance(meta, dict) and meta:
        input_tokens = meta.get("input_tokens")
        output_tokens = meta.get("output_tokens")
        if input_tokens is not None or output_tokens is not None:
            details = meta.get("input_token_details") or {}
            cache_read = 0
            if isinstance(details, dict):
                cache_read = int(details.get("cache_read") or details.get("cached_tokens") or 0)
            return {
                "input_tokens": int(input_tokens or 0),
                "output_tokens": int(output_tokens or 0),
                "cache_read_tokens": cache_read,
            }

    response_meta = getattr(message, "response_metadata", None)
    if isinstance(response_meta, dict):
        for key in ("token_usage", "usage"):
            usage = response_meta.get(key)
            if not isinstance(usage, dict):
                continue
            input_tokens = usage.get("prompt_tokens", usage.get("input_tokens"))
            output_tokens = usage.get("completion_tokens", usage.get("output_tokens"))
            if input_tokens is None and output_tokens is None:
                continue
            details = usage.get("prompt_tokens_details") or {}
            cache_read = 0
            if isinstance(details, dict):
                cache_read = int(details.get("cached_tokens") or 0)
            return {
                "input_tokens": int(input_tokens or 0),
                "output_tokens": int(output_tokens or 0),
                "cache_read_tokens": cache_read,
            }
    return None


def is_summarization_message(message: Any) -> bool:
    """判断一条消息是不是摘要中间件产出的。"""
    for attr in ("lc_source", "name"):
        if getattr(message, attr, None) == "summarization":
            return True
    for attr in ("response_metadata", "additional_kwargs"):
        meta = getattr(message, attr, None)
        if isinstance(meta, dict) and meta.get("lc_source") == "summarization":
            return True
    return False


@dataclass
class ContextSnapshot:
    """一轮请求结束时的上下文状态，用于落 session_context_usage。"""

    input_tokens: int
    output_tokens: int
    cache_read_tokens: int
    context_tokens: int
    usage_source: str
    model_calls: int
    compacted: bool
    compact_count: int
    archive_text: str | None = None
    archive_tokens: int = 0


@dataclass
class ContextMeter:
    """一次请求的上下文计量器。

    刻意做成纯逻辑、不做 I/O：它在流式循环里被高频调用，落盘与落库交给
    sse_core 的完成段统一处理，避免在事件循环里做同步文件写。
    """

    window: int = 0
    input_tokens: int = 0
    output_tokens: int = 0
    cache_read_tokens: int = 0
    context_tokens: int = 0
    model_calls: int = 0
    compacted: bool = False
    compact_count: int = 0
    _saw_provider_usage: bool = False
    _estimated_only: bool = False
    _last_context: int = 0
    # 压缩前原文的缓冲区：累积本 agent 见过的文本，压缩发生时交给上层落盘
    _transcript: list[str] = field(default_factory=list)
    _archive_pending: bool = False

    # -------- 写入 --------

    def note_text(self, text: str) -> None:
        """把一段文本纳入缓冲区（用户消息、助手正文、工具结果都算）。

        缓冲区只在内存里增长，一轮请求结束后随实例释放——它唯一的用途是
        「万一发生压缩，把被丢掉的内容存一份」。
        """
        if text:
            self._transcript.append(text)

    def note_message(self, message: Any) -> None:
        """观测一条消息：抽取真实 usage、累积文本、更新上下文占用。"""
        usage = extract_usage(message)
        if usage is not None:
            self._saw_provider_usage = True
            self.input_tokens += usage["input_tokens"]
            self.output_tokens += usage["output_tokens"]
            self.cache_read_tokens += usage["cache_read_tokens"]
            self.model_calls += 1
            # 最近一次调用的 input 就是「下一轮要带多少上下文进去」的近似
            self._update_context(usage["input_tokens"])

        if is_summarization_message(message):
            self._mark_compacted(reason="lc_source 标记")

    def note_estimated(self, input_tokens: int, output_tokens: int = 0) -> None:
        """在没有真实 usage 时登记估算值。

        一旦走过这条路，本轮用量的 ``usage_source`` 就标为 estimate——
        整条流水口径统一降级，不出现「一半真实一半估算」的混合数字。
        """
        if self._saw_provider_usage:
            # 已经有真实用量了，估算值只会污染口径，丢弃
            return
        self._estimated_only = True
        self.input_tokens += max(int(input_tokens), 0)
        self.output_tokens += max(int(output_tokens), 0)
        self._update_context(int(input_tokens))

    def _update_context(self, current_input: int) -> None:
        """更新上下文占用，并顺带检测「断崖式下跌」。"""
        if self._last_context > 0 and current_input > 0:
            if current_input < self._last_context * _COMPACTION_DROP_RATIO:
                self._mark_compacted(
                    reason=f"占用从 {self._last_context} 跌到 {current_input}"
                )
        self._last_context = current_input
        self.context_tokens = current_input

    def _mark_compacted(self, reason: str) -> None:
        if self.compacted:
            # 一节请求只记一次压缩，避免同一事件被两个信号重复计数
            return
        self.compacted = True
        self.compact_count += 1
        self._archive_pending = True
        logger.info("检测到上下文压缩（%s），当前占用=%d", reason, self.context_tokens)

    # -------- 读取 --------

    @property
    def usage_source(self) -> str:
        """provider | estimate——整轮的口径标注。"""
        return "estimate" if (self._estimated_only and not self._saw_provider_usage) else "provider"

    @property
    def total_tokens(self) -> int:
        return self.input_tokens + self.output_tokens

    def fill_ratio(self) -> float:
        """当前上下文占窗口的比例，用于判断距离压缩还有多远。"""
        if self.window <= 0:
            return 0.0
        return min(self.context_tokens / self.window, 1.0)

    def estimate_from_buffer(self) -> int:
        """用缓冲区里的文本估算总用量。

        仅在 provider 完全不返回 usage 时作为最后兜底。注意它估的是
        「对话文本量」，不含 system prompt、工具 schema、few-shot 等
        固定开销，所以**必然低估**——据此把 usage_source 标成 estimate，
        让下游知道这个数字不可尽信。
        """
        return sum(estimate_tokens(text) for text in self._transcript)

    def snapshot(self) -> ContextSnapshot:
        """产出可落库的快照。"""
        archive_text = None
        archive_tokens = 0
        if self._archive_pending and self._transcript:
            archive_text = "\n\n---\n\n".join(self._transcript)
            archive_tokens = estimate_tokens(archive_text)
            self._archive_pending = False
        return ContextSnapshot(
            input_tokens=self.input_tokens,
            output_tokens=self.output_tokens,
            cache_read_tokens=self.cache_read_tokens,
            context_tokens=self.context_tokens,
            usage_source=self.usage_source,
            model_calls=self.model_calls,
            compacted=self.compacted,
            compact_count=self.compact_count,
            archive_text=archive_text,
            archive_tokens=archive_tokens,
        )

    def summary_fields(self) -> dict[str, Any]:
        """给 SSE 事件用的精简字段。"""
        return {
            "context_tokens": self.context_tokens,
            "context_window": self.window,
            "fill_ratio": round(self.fill_ratio(), 4),
            "compacted": self.compacted,
            "usage_source": self.usage_source,
        }


def count_messages_tokens(messages: Iterable[Any]) -> int:
    """对一组消息做整体估算（压缩存档时标注「压缩前占用」用）。"""
    total = 0
    for message in messages:
        content = getattr(message, "content", None)
        if isinstance(content, str):
            total += estimate_tokens(content)
        elif isinstance(content, list):
            for block in content:
                if isinstance(block, dict) and isinstance(block.get("text"), str):
                    total += estimate_tokens(block["text"])
    return total