src/memory/store.py

from __future__ import annotations

import json
from pathlib import Path

from src.schemas import MemoryFact, UserMemory, now_iso


class MemoryStore:
    def __init__(self, base_dir: str | Path = "data") -> None:
        self.base_dir = Path(base_dir)
        self.path = self.base_dir / "memory.json"
        self.events_path = self.base_dir / "memory_events.jsonl"
        self.base_dir.mkdir(parents=True, exist_ok=True)
        (self.base_dir / "runs").mkdir(exist_ok=True)
        if not self.path.exists():
            self.save(UserMemory())
        if not self.events_path.exists():
            self.events_path.write_text("", encoding="utf-8")

    def load(self) -> UserMemory:
        try:
            return UserMemory.from_dict(json.loads(self.path.read_text(encoding="utf-8")))
        except (json.JSONDecodeError, FileNotFoundError):
            return UserMemory()

    def save(self, memory: UserMemory) -> None:
        self.path.write_text(json.dumps(memory.to_dict(), indent=2), encoding="utf-8")

    def append_event(self, action: str, fact: MemoryFact) -> None:
        event = {"ts": now_iso(), "action": action, "fact": fact.to_dict()}
        with self.events_path.open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(event) + "\n")

    def add_candidates(self, facts: list[MemoryFact], auto_accept: bool = False) -> list[MemoryFact]:
        memory = self.load()
        existing = {(fact.kind, fact.text.lower(), fact.status) for fact in memory.facts}
        added: list[MemoryFact] = []
        for fact in facts:
            fact.status = "accepted" if auto_accept else "candidate"
            key = (fact.kind, fact.text.lower(), fact.status)
            if key in existing:
                continue
            memory.facts.append(fact)
            self.append_event("add_candidate", fact)
            added.append(fact)
        self.save(memory)
        return added

    def update_status(self, fact_id: str, status: str) -> UserMemory:
        memory = self.load()
        for fact in memory.facts:
            if fact.id == fact_id:
                fact.status = status
                fact.updated_at = now_iso()
                self.append_event(f"status:{status}", fact)
                break
        self.save(memory)
        return memory

    def edit_fact(self, fact_id: str, text: str, kind: str | None = None) -> UserMemory:
        memory = self.load()
        for fact in memory.facts:
            if fact.id == fact_id:
                fact.text = text.strip()
                if kind:
                    fact.kind = kind.strip()
                fact.updated_at = now_iso()
                self.append_event("edit", fact)
                break
        self.save(memory)
        return memory

    def import_json(self, raw_json: str) -> UserMemory:
        memory = UserMemory.from_dict(json.loads(raw_json))
        self.save(memory)
        return memory