"""agent.py: a memory-augmented chat agent.

Every turn it does three memory things that a plain chatbot does not:
  1. RECALL relevant long-term facts about the user and put them in the system prompt.
  2. Send the recent conversation (short-term window, trimmed to a budget) plus the new message.
  3. After replying, SAVE the turn to short-term memory (and optionally store a new long-term fact).

It also supports CHECKPOINTING: save_state / load_state let the agent pause and resume later,
even in a brand new process, with its memory intact.

The model call is injectable (pass client=...) so the whole thing runs offline with a mock.

Run (venv active, key in .env):
    python agent.py
"""

import os
import json
from dotenv import load_dotenv
import anthropic
from memory import ShortTermMemory, LongTermMemory

load_dotenv()
MODEL = "claude-opus-4-8"


class MemoryAgent:
    def __init__(self, client=None, token_budget=120):
        self.client = client or anthropic.Anthropic()
        self.short = ShortTermMemory(token_budget=token_budget)
        self.long = LongTermMemory()

    def _system_prompt(self, user_msg):
        recalled = self.long.recall(user_msg)
        if recalled:
            facts = "\n".join(f"- {f}" for f in recalled)
            return f"You are a helpful assistant. What you remember about the user:\n{facts}"
        return "You are a helpful assistant. You have no stored facts about the user yet."

    def chat(self, user_msg, remember_fact=None):
        """One turn. Returns the assistant's reply. Optionally store a durable fact."""
        if remember_fact:
            self.long.remember(remember_fact)

        system = self._system_prompt(user_msg)
        messages = self.short.window() + [{"role": "user", "content": user_msg}]
        resp = self.client.messages.create(model=MODEL, max_tokens=400, system=system, messages=messages)
        reply = resp.content[0].text

        self.short.add("user", user_msg)          # grow short-term AFTER building the window
        self.short.add("assistant", reply)
        return reply

    # ---- checkpointing: pause and resume with memory intact ------------------
    def save_state(self, path):
        with open(path, "w") as fh:
            json.dump({"short": self.short.turns, "long": self.long.facts,
                       "budget": self.short.token_budget}, fh)

    def load_state(self, path):
        with open(path) as fh:
            state = json.load(fh)
        self.short = ShortTermMemory(token_budget=state.get("budget", 120))
        self.short.turns = state["short"]
        self.long = LongTermMemory()
        self.long.facts = state["long"]
        return self


if __name__ == "__main__":
    agent = MemoryAgent()
    print(agent.chat("My name is Sam and I love hiking.", remember_fact="The user's name is Sam and they love hiking."))
    print(agent.chat("What outdoor activity might I enjoy this weekend?"))
    agent.save_state("agent_state.json")
    print("Saved state to agent_state.json (resume later with load_state).")