"""
Lab 2.9: Context-Aware Agent -- SOLUTION

Implements sliding window + summarization for long agent sessions.
"""


class ContextManager:
    def __init__(self, max_recent_turns: int = 5, system_prompt: str = ""):
        self.system_prompt = system_prompt
        self.max_recent_turns = max_recent_turns
        self.summary = ""
        self.recent_messages = []
    
    def add_message(self, role: str, content: str):
        self.recent_messages.append({"role": role, "content": content})
        
        max_messages = self.max_recent_turns * 2
        if len(self.recent_messages) > max_messages:
            oldest = self.recent_messages[:2]
            self.recent_messages = self.recent_messages[2:]
            
            if self.summary:
                self.summary += " | "
            self.summary += f"[{oldest[0]['role']}]: {oldest[0]['content'][:100]}..."
            if len(oldest) > 1:
                self.summary += f" -> [{oldest[1]['role']}]: {oldest[1]['content'][:100]}..."
    
    def build_context(self) -> list[dict]:
        context = []
        if self.system_prompt:
            context.append({"role": "system", "content": self.system_prompt})
        if self.summary:
            context.append({"role": "system", "content": f"Session summary (earlier context): {self.summary}"})
        context.extend(self.recent_messages)
        return context
    
    def summarize(self, messages: list[dict]) -> str:
        parts = []
        for m in messages:
            role = m.get("role", "unknown")
            content = m.get("content", "")[:80]
            parts.append(f"[{role}]: {content}")
        return " | ".join(parts)


if __name__ == "__main__":
    cm = ContextManager(system_prompt="You are a helpful assistant.")
    
    for i in range(20):
        user_msg = f"User message {i+1} asking about topic {'A B C D'[i % 4]}"
        asst_msg = f"Response {i+1}: Here is information about topic {'A B C D'[i % 4]}"
        cm.add_message("user", user_msg)
        cm.add_message("assistant", asst_msg)
    
    context = cm.build_context()
    print(f"Context messages: {len(context)}")
    for m in context:
        role = m["role"]
        content = m["content"][:80]
        print(f"  [{role}]: {content}...")