# -*- coding: utf-8 -*-
"""TokenCounter 精确 token 计数测试

测试 TokenCounter 类和 SimpleTruncationEngine 的 token 计数集成。
借鉴 zhikuncode-main 的 TokenizerService，使用 tiktoken 精确计算 token 数，
替换原本基于消息条数的估算（长消息单条即可能超限却未触发压缩的问题）。
"""
from unittest.mock import patch, MagicMock

import pytest

from abyss.context_engine import (
    SimpleTruncationEngine,
    TokenCounter,
)


# ---------- TokenCounter 基本计数 ----------


class TestTokenCounterBasic:
    """TokenCounter 基本计数测试"""

    def test_count_text_english(self):
        """英文文本应返回合理 token 数（hello world 约 2-3 token）。"""
        counter = TokenCounter()
        tokens = counter.count_text("hello world")
        assert 2 <= tokens <= 3

    def test_count_text_empty(self):
        """空字符串应返回 0。"""
        counter = TokenCounter()
        assert counter.count_text("") == 0

    def test_count_text_chinese(self):
        """中文文本 token 数应与英文不同（每字约 1-2 token）。"""
        counter = TokenCounter()
        tokens = counter.count_text("你好世界")
        assert tokens > 0
        english_tokens = counter.count_text("hello world")
        # 中文与英文 token 数应不同
        assert tokens != english_tokens


# ---------- count_messages 精确计数 ----------


class TestCountMessages:
    """count_messages 精确计数测试"""

    def test_count_messages_positive(self):
        """消息列表应返回正 token 数。"""
        counter = TokenCounter()
        messages = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "Hello"},
        ]
        assert counter.count_messages(messages) > 0

    def test_count_messages_includes_role_overhead(self):
        """count_messages 应比逐条 count_text 之和多（含 role 开销）。"""
        counter = TokenCounter()
        messages = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "Hello"},
        ]
        # 仅统计 content 的 token 数
        sum_of_content = sum(
            counter.count_text(m["content"]) for m in messages
        )
        total = counter.count_messages(messages)
        # 含 role 字段开销和每条 4 token 的开销，应严格大于
        assert total > sum_of_content

    def test_count_messages_empty(self):
        """空消息列表应返回 0。"""
        counter = TokenCounter()
        assert counter.count_messages([]) == 0


# ---------- 缓存命中 ----------


class TestTokenCounterCache:
    """TokenCounter encoder 缓存测试"""

    def test_encoder_loaded_only_once(self):
        """同一实例多次调用 count_text 应只加载一次 encoder。"""
        with patch("tiktoken.get_encoding") as mock_get:
            mock_encoder = MagicMock()
            mock_encoder.encode.return_value = [1, 2, 3]
            mock_get.return_value = mock_encoder

            counter = TokenCounter()
            counter.count_text("hello")
            counter.count_text("world")
            counter.count_text("foo bar")

            # get_encoding 只应被调用一次（缓存命中）
            assert mock_get.call_count == 1

    def test_encoder_lazy_load(self):
        """encoder 应懒加载：构造时不立即加载，首次计数时才加载。"""
        with patch("tiktoken.get_encoding") as mock_get:
            mock_encoder = MagicMock()
            mock_encoder.encode.return_value = [1]
            mock_get.return_value = mock_encoder

            counter = TokenCounter()
            # 构造完成，未调用任何计数方法，不应加载 encoder
            assert mock_get.call_count == 0

            counter.count_text("hello")
            assert mock_get.call_count == 1


# ---------- SimpleTruncationEngine 集成 ----------


class TestEngineTokenIntegration:
    """SimpleTruncationEngine 与 TokenCounter 集成测试"""

    def test_update_from_response_uses_api_value(self):
        """update_from_response 应优先使用 API 提供的 prompt_tokens。"""
        engine = SimpleTruncationEngine()
        engine.update_from_response({"prompt_tokens": 1000})
        assert engine.last_prompt_tokens == 1000

    def test_should_compress_local_estimation_long_message(self):
        """无 API usage 时，单条超长消息应触发压缩（即使消息条数少）。"""
        engine = SimpleTruncationEngine(
            context_length=1000,
            threshold_percent=0.75,
        )
        # 单条约 9000 字符的英文消息，token 数远超 750 阈值
        long_content = "The quick brown fox jumps over the lazy dog. " * 200
        engine._messages = [
            {"role": "user", "content": long_content},
        ]
        assert engine.should_compress() is True

    def test_should_compress_local_estimation_short_message(self):
        """无 API usage 时，短消息不应触发压缩。"""
        engine = SimpleTruncationEngine(
            context_length=1000,
            threshold_percent=0.75,
        )
        engine._messages = [
            {"role": "user", "content": "hi"},
        ]
        assert engine.should_compress() is False

    def test_estimate_tokens_returns_positive(self):
        """estimate_tokens 应返回正数。"""
        engine = SimpleTruncationEngine()
        messages = [{"role": "user", "content": "hello world"}]
        assert engine.estimate_tokens(messages) > 0

    def test_estimate_tokens_uses_token_counter(self):
        """estimate_tokens 应与 TokenCounter.count_messages 结果一致。"""
        engine = SimpleTruncationEngine()
        messages = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "Hello"},
        ]
        assert engine.estimate_tokens(messages) == engine._token_counter.count_messages(messages)


# ---------- 阈值判断 ----------


class TestThresholdJudgment:
    """阈值判断测试"""

    def test_threshold_tokens_calculated(self):
        """threshold_tokens 应为 context_length * threshold_percent。"""
        engine = SimpleTruncationEngine(
            context_length=1000,
            threshold_percent=0.75,
        )
        assert engine.threshold_tokens == 750

    def test_compress_triggered_at_threshold(self):
        """消息 token 数 >= 阈值（750）时应触发压缩。"""
        engine = SimpleTruncationEngine(
            context_length=1000,
            threshold_percent=0.75,
        )
        # 构造 >= 750 token 的消息（约 4500 字符英文，~1200 token）
        long_content = "The quick brown fox jumps over the lazy dog. " * 100
        engine._messages = [
            {"role": "user", "content": long_content},
        ]
        assert engine.should_compress() is True

    def test_compress_not_triggered_below_threshold(self):
        """消息 token 数 < 阈值时不应触发压缩。"""
        engine = SimpleTruncationEngine(
            context_length=1000,
            threshold_percent=0.75,
        )
        engine._messages = [
            {"role": "user", "content": "short message"},
        ]
        assert engine.should_compress() is False
