# emb_cache.py — one shared MiniLM instance + a text->vector cache (in-memory + disk).
# The fast metric (emb_judge) and the dedup guard both embed the same gold refs / drafts
# many times; caching by text turns repeated encodes into dict lookups, which is what makes
# offline replay and tuning run in seconds instead of minutes.
import os
import json
import threading

import numpy as np

from humomni.core.env_util import load_env

load_env()   # set HF_TOKEN / HF_HUB_OFFLINE from .env BEFORE sentence_transformers/huggingface_hub load

# Persist as a plain .npy matrix + a .json key list (NO pickle / allow_pickle) so loading
# never executes arbitrary code, even though this cache is only ever produced locally.
PATH = "cache/emb_cache"   # -> emb_cache.npy + emb_cache.keys.json
_model = None
_cache = {}
_lock = threading.Lock()


def _model_():
    global _model
    if _model is None:
        from sentence_transformers import SentenceTransformer
        _model = SentenceTransformer("all-MiniLM-L6-v2")
    return _model


def load_disk(path=PATH):
    npy, kj = path + ".npy", path + ".keys.json"
    if os.path.exists(npy) and os.path.exists(kj):
        vecs = np.load(npy)                       # plain float matrix, no pickle
        keys = json.load(open(kj, encoding="utf-8"))
        for k, v in zip(keys, vecs):
            _cache.setdefault(k, v)


def save_disk(path=PATH):
    if not _cache:
        return
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    keys = list(_cache.keys())
    vecs = np.stack([_cache[k] for k in keys]).astype(np.float32)
    np.save(path + ".npy", vecs)
    json.dump(keys, open(path + ".keys.json", "w", encoding="utf-8"))


def embed_many(texts):
    """Return unit-norm vectors for texts (order preserved), encoding only cache misses."""
    miss = [t for t in texts if t not in _cache]
    if miss:
        uniq = list(dict.fromkeys(miss))
        with _lock:
            still = [t for t in uniq if t not in _cache]
            if still:
                vecs = _model_().encode(still, normalize_embeddings=True,
                                        batch_size=256, show_progress_bar=False)
                for t, v in zip(still, np.asarray(vecs, dtype=np.float32)):
                    _cache[t] = v
    return [_cache[t] for t in texts]


def embed(text):
    return embed_many([text])[0]


def cos(a, b):
    va, vb = embed_many([a, b])
    return float(np.dot(va, vb))