# core/embedder.py
# CRITICAL: Must set BEFORE importing chromadb to suppress telemetry
import os
os.environ["ANONYMIZED_TELEMETRY"] = "False"

import hashlib
import chromadb
from chromadb.config import Settings

class LocalEmbedder:
    def __init__(self, persist_dir=None):
        self.persist_dir = persist_dir or "/data/invicta_data/vectors"
        os.makedirs(self.persist_dir, exist_ok=True)
        self.client = chromadb.Client(Settings(
            persist_directory=self.persist_dir,
            anonymized_telemetry=False
        ))
        self.collection = self.client.get_or_create_collection("invicta_embeddings")
        self._embedding_function = None
        self._embedding_ok = None

    def _ensure_embedding_function(self):
        """Lazy-load the embedding model so startup isn't blocked by a download."""
        if self._embedding_ok is None:
            try:
                from chromadb.utils import embedding_functions
                self._embedding_function = embedding_functions.DefaultEmbeddingFunction()
                self._embedding_ok = True
                print("✅ Embedding function loaded")
            except Exception as e:
                print(f"⚠️ Embedding function unavailable: {e}")
                self._embedding_ok = False
        return self._embedding_ok

    def _doc_id(self, text):
        return hashlib.md5(text.strip().lower().encode()).hexdigest()

    def store(self, text, metadata=None):
        if not text or not text.strip():
            return
        if not self._ensure_embedding_function():
            return
        doc_id = self._doc_id(text)
        meta = metadata or {}
        existing = self.collection.get(ids=[doc_id])
        if existing and existing.get("ids") and len(existing["ids"]) > 0:
            return
        self.collection.add(
            ids=[doc_id],
            documents=[text],
            metadatas=[meta],
            embeddings=[self._embedding_function([text])[0]]
        )

    def find_similar(self, query, top_k=5):
        if not self._ensure_embedding_function():
            return []
        embedding = self._embedding_function([query])[0]
        results = self.collection.query(
            query_embeddings=[embedding],
            n_results=top_k,
            include=["documents", "distances"]
        )
        docs = []
        if results and results.get("documents"):
            for doc, dist in zip(results["documents"][0], results["distances"][0]):
                if dist < 0.8:
                    docs.append(doc)
        return docs

    def count(self):
        return self.collection.count()