"""
clankerDiffusion — learned secondary memory (side store).

The model controls a long-term store through special tokens it emits while
generating:

    <mem_write>KEY<mem_kv>BLOB</mem_kv>      -> persist BLOB under KEY
    <mem_read>KEY</mem_read>                -> recall BLOB for KEY
    <mem_evict>KEY</mem_evict>              -> delete KEY

BLOBs are packed into the token stream as <mem_kv>...</mem_kv> so the model can
read what it wrote. At inference, `MemoryStore` intercepts these tokens, performs
the read/write against an external store (dict / JSONL / vector index), and
replaces the in-stream content with the resolved value (or an <unk> if missing).

This is "knowledge brought in and out of memory": the weights never hold the
facts; only the keys/pointers do. Train 8k, but the store is unbounded, so the
effective secondary memory is far larger than the 128k normal window.

Training note: the data generator emits syntactically-valid mem_write / mem_read
/ mem_evict spans using this same module, so the model learns the format and
when to invoke it.
"""
import re
import json
import os


class MemoryStore:
    def __init__(self, path=None):
        self.path = path
        self.store = {}
        if path and os.path.exists(path):
            try:
                with open(path, "r", encoding="utf-8") as f:
                    self.store = json.load(f)
            except Exception:
                self.store = {}

    # ---- token-stream helpers (used by data gen + inference) ----
    def write(self, key, blob):
        self.store[key] = blob
        self._maybe_persist()
        return blob

    def read(self, key):
        return self.store.get(key)

    def evict(self, key):
        return self.store.pop(key, None)

    def _maybe_persist(self):
        if self.path:
            try:
                with open(self.path, "w", encoding="utf-8") as f:
                    json.dump(self.store, f)
            except Exception:
                pass

    # ---- stream rewriting for inference ----
    @staticmethod
    def _split_segments(text):
        """Yield (is_mem, content) chunks where is_mem=True spans a full
        <mem_*>...</mem_*> block we can act on."""
        pat = re.compile(
            r"<mem_write>(.*?)<mem_kv>(.*?)</mem_kv>|"
            r"<mem_read>(.*?)</mem_read>|"
            r"<mem_evict>(.*?)</mem_evict>",
            re.DOTALL)
        pos = 0
        for m in pat.finditer(text):
            if m.start() > pos:
                yield (False, text[pos:m.start()])
            if m.group(1) is not None:
                key, blob = m.group(1), m.group(2)
                yield (True, ("write", key, blob))
            elif m.group(3) is not None:
                yield (True, ("read", m.group(3)))
            elif m.group(4) is not None:
                yield (True, ("evict", m.group(4)))
            pos = m.end()
        if pos < len(text):
            yield (False, text[pos:])

    def resolve(self, text):
        """Rewrite a generated stream: execute memory ops, replacing the op
        with its resolved <mem_kv> value (or a miss marker). Returns the new
        text and a list of (op, key) that were performed (for logging)."""
        out = []
        ops = []
        for is_mem, chunk in self._split_segments(text):
            if not is_mem:
                out.append(chunk)
                continue
            op = chunk[0]
            if op == "write":
                _, key, blob = chunk
                self.write(key, blob)
                ops.append(("write", key))
                out.append(f"<mem_kv>{blob}</mem_kv>")
            elif op == "read":
                key = chunk[1]
                blob = self.read(key)
                ops.append(("read", key))
                out.append(f"<mem_kv>{blob}</mem_kv>" if blob is not None
                           else "<mem_kv><unk></mem_kv>")
            elif op == "evict":
                key = chunk[1]
                self.evict(key)
                ops.append(("evict", key))
                out.append("")
        return "".join(out), ops