import torch
from sentence_transformers import SentenceTransformer
from collections import OrderedDict
import constants

print("AI Model Load हो रहा है...")
model = SentenceTransformer(constants.MODEL_NAME, trust_remote_code=True)
if torch.cuda.is_available():
    model = model.to("cuda")

# ⚡ Bounded Cache (CPU RAM storage to save VRAM)
class BoundedEmbeddingCache:
    def __init__(self, maxsize=2000):
        self.cache = OrderedDict()
        self.maxsize = maxsize

    def get(self, text):
        if text in self.cache:
            self.cache.move_to_end(text)
            return self.cache[text]
        return None

    def put(self, text, emb):
        self.cache[text] = emb
        self.cache.move_to_end(text)
        if len(self.cache) > self.maxsize:
            self.cache.popitem(last=False)

EMBEDDING_CACHE = BoundedEmbeddingCache(maxsize=2000)

def get_batch_embeddings(texts):
    uncached_texts = [t for t in texts if EMBEDDING_CACHE.get(t) is None]
    if uncached_texts:
        with torch.no_grad():
            embs = model.encode(uncached_texts, convert_to_tensor=True, normalize_embeddings=True)
            for t, e in zip(uncached_texts, embs):
                EMBEDDING_CACHE.put(t, e.cpu())
                
    return torch.stack([EMBEDDING_CACHE.get(t) for t in texts]).to(model.device)