import json, os, time
import numpy as np
from sentence_transformers import SentenceTransformer

def _cosine_sim(a: np.ndarray, b: np.ndarray) -> float:
    an = np.linalg.norm(a)
    bn = np.linalg.norm(b)
    if an == 0 or bn == 0:
        return 0.0
    return float(np.dot(a, b) / (an * bn))

class SemanticMemory:
    """
    Per-agent vector memory: stores (text, category, embedding, ts)
    Fast, JSON-backed, CPU-only.
    """
    def __init__(self, agent_id: str, file_path: str = None, model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
        self.agent_id = agent_id
        self.file_path = file_path or f"semantic_memory_{agent_id}.json"
        self.model = SentenceTransformer(model_name)
        self._init_store()

    def _init_store(self):
        if not os.path.exists(self.file_path):
            with open(self.file_path, "w") as f:
                json.dump({"entries": []}, f)

    def _load(self):
        with open(self.file_path, "r") as f:
            return json.load(f)

    def _save(self, obj):
        with open(self.file_path, "w") as f:
            json.dump(obj, f, indent=2)

    def add(self, text: str, category: str):
        store = self._load()
        emb = self.model.encode([text])[0].tolist()  # list[float]
        store["entries"].append({
            "text": text,
            "category": category,
            "embedding": emb,
            "ts": int(time.time())
        })
        self._save(store)

    def query(self, query_text: str, category: str = None, top_k: int = 5):
        store = self._load()
        entries = store.get("entries", [])
        if category:
            entries = [e for e in entries if e.get("category") == category]
        if not entries:
            return []

        q = self.model.encode([query_text])[0]
        scored = []
        for e in entries:
            v = np.array(e["embedding"], dtype=np.float32)
            s = _cosine_sim(q, v)
            scored.append((s, e))
        scored.sort(key=lambda x: x[0], reverse=True)
        return [e for (s, e) in scored[:top_k] if s > 0.35]  # simple confidence cutoff