import json
from typing import Any

from strands.agent.conversation_manager.sliding_window_conversation_manager import SlidingWindowConversationManager

DEFAULT_MAX_BYTES = 300000
DEFAULT_MAX_WINDOW = 400


class ByteWindowConversationManager(SlidingWindowConversationManager):
    def __init__(
            self,
            max_bytes: int = DEFAULT_MAX_BYTES,
            max_window: int = DEFAULT_MAX_WINDOW,
            **kwargs: Any
    ):
        super().__init__(max_window=max_window, **kwargs)
        self.max_bytes = max_bytes
        
    def _byte_window(self, messages: list) -> int:
        """计算保留的消息数量"""
        total = 0
        keep = 0
        for m in reversed(messages):
            total += len(json.dumps(m, ensure_ascii=False, default=str).encode("utf-8"))
            if keep >= 1 and total > self.max_bytes:
                break
            keep += 1
        return max(keep, 2)
    
    def apply_management(self, agent: "Any", **kwargs: Any) -> None:
        """应用消息管理策略，确保消息总字节数不超过 max_bytes。"""
        messages = agent.messages
        eff = min(self.window_size, self._byte_window(messages))
        
        if len(messages) <= eff:
            return
        
        saved = self.window_size
        self.window_size = eff
        try:
            self.reduce_context(agent)
        finally:
            self.window_size = saved