# -*- coding: utf-8 -*-
"""context_engine 模块测试"""
import pytest
from abyss.context_engine import (
    ContextEngine,
    SimpleTruncationEngine,
    create_engine,
)


class TestSimpleTruncationEngine:
    """SimpleTruncationEngine 测试"""

    def test_create_default(self):
        engine = SimpleTruncationEngine()
        assert engine.name == "simple"
        assert engine.context_length == 128000
        assert engine.threshold_tokens == int(128000 * 0.75)

    def test_create_custom(self):
        engine = SimpleTruncationEngine(
            context_length=64000,
            threshold_percent=0.8,
            protect_first_n=2,
            protect_last_n=4,
        )
        assert engine.context_length == 64000
        assert engine.threshold_tokens == int(64000 * 0.8)
        assert engine.protect_first_n == 2
        assert engine.protect_last_n == 4

    def test_update_from_response(self):
        engine = SimpleTruncationEngine()
        engine.update_from_response({
            "prompt_tokens": 50000,
            "completion_tokens": 2000,
            "total_tokens": 52000,
        })
        assert engine.last_prompt_tokens == 50000
        assert engine.last_completion_tokens == 2000
        assert engine.last_total_tokens == 52000

    def test_should_compress_below_threshold(self):
        engine = SimpleTruncationEngine(context_length=128000, threshold_percent=0.75)
        engine.update_from_response({"prompt_tokens": 50000})
        assert engine.should_compress() is False

    def test_should_compress_above_threshold(self):
        engine = SimpleTruncationEngine(context_length=128000, threshold_percent=0.75)
        engine.update_from_response({"prompt_tokens": 100000})
        assert engine.should_compress() is True

    def test_should_compress_explicit_tokens(self):
        engine = SimpleTruncationEngine(context_length=128000, threshold_percent=0.75)
        assert engine.should_compress(prompt_tokens=96000) is True
        assert engine.should_compress(prompt_tokens=10000) is False

    def test_compress_preserves_system_prompt(self):
        engine = SimpleTruncationEngine(protect_last_n=2)
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
        ]
        result = engine.compress(messages)
        assert result[0] == messages[0]

    def test_compress_small_list_no_change(self):
        """消息数太少时不压缩。"""
        engine = SimpleTruncationEngine(protect_last_n=8)
        messages = [
            {"role": "system", "content": "sys"},
            {"role": "user", "content": "hello"},
            {"role": "assistant", "content": "hi"},
        ]
        result = engine.compress(messages)
        assert result == messages

    def test_compress_truncates_middle(self):
        engine = SimpleTruncationEngine(
            protect_first_n=1,
            protect_last_n=2,
        )
        messages = [
            {"role": "system", "content": "sys"},
            {"role": "user", "content": "msg1"},
            {"role": "assistant", "content": "msg2"},
            {"role": "user", "content": "msg3"},
            {"role": "assistant", "content": "msg4"},
            {"role": "user", "content": "msg5"},
            {"role": "assistant", "content": "msg6"},
        ]
        result = engine.compress(messages)
        # system + head(1) + truncation_note + tail(2)
        assert len(result) == 5
        assert result[0]["role"] == "system"
        assert result[1] == messages[1]  # head
        assert result[2]["role"] == "user"
        assert "上下文压缩" in result[2]["content"]
        assert result[3] == messages[5]  # tail[-2]
        assert result[4] == messages[6]  # tail[-1]

    def test_compress_increments_counter(self):
        engine = SimpleTruncationEngine(protect_first_n=1, protect_last_n=1)
        messages = [
            {"role": "system", "content": "sys"},
            {"role": "user", "content": "a"},
            {"role": "assistant", "content": "b"},
            {"role": "user", "content": "c"},
            {"role": "assistant", "content": "d"},
            {"role": "user", "content": "e"},
        ]
        engine.compress(messages)
        assert engine.compression_count == 1
        engine.compress(messages)
        assert engine.compression_count == 2

    def test_compress_with_focus_topic(self):
        engine = SimpleTruncationEngine(protect_first_n=1, protect_last_n=1)
        messages = [
            {"role": "system", "content": "sys"},
            {"role": "user", "content": "a"},
            {"role": "assistant", "content": "b"},
            {"role": "user", "content": "c"},
            {"role": "assistant", "content": "d"},
            {"role": "user", "content": "e"},
        ]
        result = engine.compress(messages, focus_topic="数据库迁移")
        assert "数据库迁移" in result[2]["content"]

    def test_compress_does_not_modify_original(self):
        engine = SimpleTruncationEngine(protect_last_n=1)
        messages = [
            {"role": "system", "content": "sys"},
            {"role": "user", "content": "a"},
            {"role": "assistant", "content": "b"},
            {"role": "user", "content": "c"},
        ]
        original = [dict(m) for m in messages]
        engine.compress(messages)
        assert messages == original  # 原始未被修改

    def test_update_model(self):
        engine = SimpleTruncationEngine(context_length=128000)
        engine.update_model(64000)
        assert engine.context_length == 64000
        assert engine.threshold_tokens == int(64000 * 0.75)

    def test_get_status(self):
        engine = SimpleTruncationEngine(context_length=128000)
        engine.update_from_response({"prompt_tokens": 96000})
        status = engine.get_status()
        assert status["last_prompt_tokens"] == 96000
        assert status["context_length"] == 128000
        assert status["usage_percent"] == 75
        assert status["compression_count"] == 0

    def test_on_session_reset(self):
        engine = SimpleTruncationEngine()
        engine.update_from_response({"prompt_tokens": 96000})
        engine.compression_count = 3
        engine.on_session_reset()
        assert engine.last_prompt_tokens == 0
        assert engine.compression_count == 0


class TestCreateEngine:
    """工厂函数测试"""

    def test_create_simple(self):
        engine = create_engine("simple", context_length=64000)
        assert isinstance(engine, SimpleTruncationEngine)
        assert engine.name == "simple"

    def test_create_unknown(self):
        with pytest.raises(ValueError, match="未知上下文引擎"):
            create_engine("nonexistent")
