# memory.py — lightweight conversation memory
from __future__ import annotations
from dataclasses import dataclass, asdict
from typing import List, Dict, Optional
import json, time, os

@dataclass
class Turn:
    role: str          # "user" | "assistant" | "tool"
    content: str
    meta: Dict = None
    ts: float = 0.0

    def to_dict(self):
        d = asdict(self)
        d["ts"] = self.ts or time.time()
        d["meta"] = self.meta or {}
        return d

class ConversationMemory:
    """
    Simple memory with:
      - buffer: last N turns (short-term)
      - summary: rolling abstractive summary (long-term)
      - entities: key entities/ids spotted so far
    """
    def __init__(self, path: str, buffer_size: int = 12):
        self.path = path
        self.buffer_size = buffer_size
        self.buffer: List[Turn] = []
        self.summary: str = ""
        self.entities: Dict[str, List[str]] = {}  # e.g., {"competition_id": ["2313", ...]}
        self._load()

    # ------------ persistence ------------
    def _load(self):
        if not os.path.exists(self.path): return
        with open(self.path, "r") as f:
            data = json.load(f)
        self.buffer = [Turn(**t) for t in data.get("buffer", [])]
        self.summary = data.get("summary", "")
        self.entities = data.get("entities", {})

    def _save(self):
        os.makedirs(os.path.dirname(self.path), exist_ok=True)
        with open(self.path, "w") as f:
            json.dump({
                "buffer": [t.to_dict() for t in self.buffer],
                "summary": self.summary,
                "entities": self.entities,
            }, f, ensure_ascii=False, indent=2)

    # ------------ public API ------------
    def add_turn(self, role: str, content: str, meta: Optional[Dict]=None):
        self.buffer.append(Turn(role=role, content=content, meta=meta or {}, ts=time.time()))
        if len(self.buffer) > self.buffer_size:
            self.buffer = self.buffer[-self.buffer_size:]
        self._save()

    def get_context(self) -> Dict:
        """What to feed into prompts/tools."""
        return {
            "summary": self.summary,
            "recent": [{"role": t.role, "content": t.content} for t in self.buffer[-self.buffer_size:]],
            "entities": self.entities,
        }

    def update_summary(self, llm_summarize_fn):
        """
        Call with a function that maps (summary, recent) -> new_summary.
        Only do this occasionally (e.g., every 6–10 user turns).
        """
        if not self.buffer:
            return
        recent_text = "\n".join(
            f"{t.role.upper()}: {t.content}" for t in self.buffer[-self.buffer_size:]
        )
        prompt = (
            "You are a diligent note-taker. Update the long-term summary of this conversation.\n"
            "Keep it under 150 words. Capture tasks, preferences, important grant IDs/themes, and open questions.\n\n"
            f"EXISTING SUMMARY:\n{self.summary or '(none)'}\n\n"
            f"RECENT TURNS:\n{recent_text}\n\n"
            "Return ONLY the updated summary text."
        )
        new_sum = llm_summarize_fn(prompt).strip()
        if new_sum:
            self.summary = new_sum
            self._save()

    def add_entity(self, kind: str, value: str):
        if not value: return
        arr = self.entities.setdefault(kind, [])
        if value not in arr:
            arr.append(value)
            self._save()