import chromadb as db
from chromadb.utils import embedding_functions
from chromadb.config import Settings
from chromadb.api.client import Client
from chromadb.api.models.Collection import Collection



class ChromaDB:
    """
    A class representing a Chroma database.
    """

    def __init__(self, host: str = "192.168.1.10", port: int = 8001) -> None:
        """
        Initializes a ChromaDB object running on specified port.

        Args:
            host (str, optional): The host address of the Chroma database. Defaults to "192.168.1.10".
            port (int, optional): The port number of the Chroma database. Defaults to 8001.
        """
        self.client: Client = db.HttpClient(
            settings=Settings(anonymized_telemetry=False),
            host=host,
            port=port,
        )
        huggingface_ef = embedding_functions.HuggingFaceEmbeddingFunction(
        api_key='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',
        model_name="KBLab/sentence-bert-swedish-cased"
)
        self.embedding_function: embedding_functions = huggingface_ef
    
    def print_collections(self):
        """
        Prints all collections in the database.
        """
        collections = self.client.list_collections()
        for collection in collections:
            print(collection.name)

# Initialize the ChromaDB object
chroma = ChromaDB()

if __name__ == '__main__':

    chroma.client.delete_collection('mala_persons')
    col = chroma.client.get_or_create_collection('mala_persons')

    print(col.count())