import os
import chromadb
from chromadb.utils import embedding_functions

class ChromaStorage:
    def __init__(self, db_path: str = './src/chroma_db', collection_name: str = 'my_collection'):
       
        self.client = chromadb.PersistentClient(path=db_path)
        self.collection = self._get_or_create_collection(collection_name)

    def _get_or_create_collection(self, name: str):
        try:
            return self.client.get_collection(
                name=name
            )
        except Exception:
            return self.client.create_collection(
                name=name,
                embedding_function=self.openai_ef
            )

    def add_batch(self, documents: list, metadatas: list, ids: list, batch_size: int = 200):
        for i in range(0, len(documents), batch_size):
            docs = documents[i:i + batch_size]
            metas = metadatas[i:i + batch_size]
            batch_ids = ids[i:i + batch_size]
            self.collection.add(documents=docs, metadatas=metas, ids=batch_ids)

    def query(self, query_text: str, k: int = 1) -> list:
        return self.collection.query(query_texts=[query_text], n_results=k)

    def delete_all(self):
        self.client.reset()