"""add_memory_policy.py: your turn, make short-term memory smarter than "drop the oldest".

The solution's ShortTermMemory drops whole old turns when it hits the token budget. That is
simple but it forgets things. A better policy: when the conversation gets too long, SUMMARIZE
the oldest turns into a single short note and keep that note plus the recent turns. This keeps
the gist of old context for far fewer tokens.

Steps:
  1. Implement summarize_old below (a real one calls the model; a free placeholder is fine to start).
  2. Use it in window_with_summary to replace dropped turns with one summary turn.
  3. Run:  python add_memory_policy.py

Worked short-term/long-term code is in ../solution/memory.py.
"""

import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "solution"))
from memory import ShortTermMemory, approx_tokens


# TODO: turn several old turns into one short note.
# Free placeholder: join them. Real version: ask the model for a 1-2 sentence summary (costs tokens).
def summarize_old(turns):
    joined = " ".join(t["content"] for t in turns)
    return f"(earlier conversation summary) {joined[:120]}"


def window_with_summary(mem, recent_keep=2):
    """Keep the last `recent_keep` turns verbatim; fold everything older into one summary turn."""
    if len(mem.turns) <= recent_keep:
        return mem.turns
    old, recent = mem.turns[:-recent_keep], mem.turns[-recent_keep:]
    # TODO: return [summary_turn] + recent, where summary_turn uses summarize_old(old)
    summary_turn = {"role": "system", "content": summarize_old(old)}
    return [summary_turn] + recent


if __name__ == "__main__":
    mem = ShortTermMemory(token_budget=9999)
    for i in range(6):
        mem.add("user", f"Fact {i}: I told you something important number {i}.")
    print("plain drop-oldest window keeps:", [t["content"][:22] for t in mem.window()])
    print("summary policy keeps:", [t["content"][:40] for t in window_with_summary(mem)])
    # Compare token counts: the summary policy should remember MORE history for FEWER tokens.