import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer

# 1) Modèle d'embedding
model = SentenceTransformer("all-MiniLM-L6-v2")

def embed(texts):
    return model.encode(texts).tolist()

# 2) Client Chroma en mode local (en mémoire pour le test)
client = chromadb.PersistentClient(path="db/")

# 3) Création / récupération d’une collection
collection = client.get_or_create_collection(
    name="docs_demo",
    metadata={"hnsw:space": "cosine"}  # métrique de similarité
)

# 4) Jeu d’essai : quelques phrases
documents = [
    "Les chats aiment dormir au soleil.",
    "Les chiens sont des animaux très fidèles.",
    "Les voitures électriques deviennent de plus en plus populaires.",
    "Les félins sauvages chassent principalement la nuit.",
    "Les transports en commun réduisent la pollution en ville."
]

ids = [f"doc_{i}" for i in range(len(documents))]
metadatas = [
    {"type": "animal"},
    {"type": "animal"},
    {"type": "transport"},
    {"type": "animal"},
    {"type": "transport"},
]

# 5) Vectorisation + insertion dans Chroma
embeddings = embed(documents)

collection.add(
    ids=ids,
    documents=documents,
    metadatas=metadatas,
    embeddings=embeddings
)

print("Documents insérés dans la collection.")

# 6) Requête sémantique
query = "animaux qui chassent la nuit"

query_emb = embed([query])

results = collection.query(
    query_embeddings=query_emb,
    n_results=3,
    include=["documents", "metadatas", "distances"]
)

print("\nRequête :", query)
for doc, meta, dist in zip(
    results["documents"][0],
    results["metadatas"][0],
    results["distances"][0]
):
    print(f"- doc: {doc}")
    print(f"  meta: {meta}, distance: {dist:.4f}")

print(len(embeddings[1]))
