"""
Step 04 练习题答案

⚠️ 先自己思考，再看答案！
⚠️ 答案不是唯一的，这里只是其中一种实现
"""

import math
import re

from step_04_memory.concept import MemorySystem, Message, ShortTermMemory


# ═══════════════════════════════════════════════════════════════════════════════
# 练习 1 答案：Token 估算
# ═══════════════════════════════════════════════════════════════════════════════

_CN_PATTERN = re.compile(r"[\u4e00-\u9fff]")
_EN_WORD = re.compile(r"[A-Za-z]+")


def estimate_tokens_for_text(text: str) -> int:
    cn = len(_CN_PATTERN.findall(text))
    en = len(_EN_WORD.findall(text))
    return math.ceil(cn * 1.5 + en * 1.3)


def install_estimate_tokens() -> None:
    def estimate_tokens(self: ShortTermMemory) -> int:
        return sum(estimate_tokens_for_text(m.content) for m in self.messages)

    ShortTermMemory.estimate_tokens = estimate_tokens


# ═══════════════════════════════════════════════════════════════════════════════
# 练习 2 答案：基于 Token 阈值的自动压缩
# ═══════════════════════════════════════════════════════════════════════════════

def install_maybe_compress() -> None:
    def maybe_compress(self: ShortTermMemory, max_tokens: int = 800) -> bool:
        if not hasattr(self, "estimate_tokens"):
            raise RuntimeError("请先调用 install_estimate_tokens()")
        if self.estimate_tokens() <= max_tokens:
            return False
        summary = self.summarize_older(keep_recent=5)
        if not summary:
            return False
        # 保留最近 5 条，把摘要作为 system message 放最前
        recent = self.messages[-5:]
        self.messages = [Message(role="system", content=f"[历史摘要]\n{summary}")] + recent
        return True

    ShortTermMemory.maybe_compress = maybe_compress


# ═══════════════════════════════════════════════════════════════════════════════
# 练习 3 答案：把 MemorySystem 接到 SimpleAgent
# ═══════════════════════════════════════════════════════════════════════════════

def install_memory_to_agent() -> None:
    from step_03_simple_agent.concept import SimpleAgent

    orig_init = SimpleAgent.__init__

    def patched_init(self, *args, **kwargs):
        orig_init(self, *args, **kwargs)
        self.memory = MemorySystem()

    SimpleAgent.__init__ = patched_init

    orig_process = SimpleAgent.process

    def patched_process(self, user_input: str) -> str:
        # 把记忆上下文注入 state
        self.state["context"] = self.memory.get_context()
        response = orig_process(self, user_input)
        # 记录本轮对话
        self.memory.add_message("user", user_input)
        self.memory.add_message("assistant", response)
        return response

    SimpleAgent.process = patched_process


# ═══════════════════════════════════════════════════════════════════════════════
# 测试
# ═══════════════════════════════════════════════════════════════════════════════

def test_answers():
    print("\n" + "=" * 60)
    print("Step 04 练习答案测试")
    print("=" * 60)

    install_estimate_tokens()
    install_maybe_compress()

    mem = MemorySystem()
    for i in range(20):
        mem.short_term.add("user", f"第 {i} 轮对话内容，包含中文与 english words " * 5)
    print(f"\n📝 练习 1: 注入 20 条后估算 token = {mem.short_term.estimate_tokens()}")
    compressed = mem.short_term.maybe_compress(max_tokens=200)
    print(f"  maybe_compress() = {compressed}, 压缩后消息数 = {len(mem.short_term.messages)}")
    print(f"  压缩后估算 token = {mem.short_term.estimate_tokens()}")

    print("\n📝 练习 3: SimpleAgent 接入 Memory")
    try:
        install_memory_to_agent()
        from step_03_simple_agent.concept import SimpleAgent

        agent = SimpleAgent()
        agent.process("1 + 2")
        print(f"  agent.memory 工作正常，消息数 = {len(agent.memory.short_term.messages)}")
    except Exception as e:
        print(f"  接入失败（可忽略，需在 step_03 父目录运行）: {e}")


if __name__ == "__main__":
    test_answers()