"""Library of the project."""
# pylint: disable=wrong-import-position

# %% IMPORTS

__import__("pysqlite3")

import os
import sys

# https://docs.trychroma.com/troubleshooting#sqlite
sys.modules["sqlite3"] = sys.modules.pop("pysqlite3")

import chromadb
from chromadb.utils import embedding_functions

# %% CONFIGS

DATABASE_COLLECTION = "resume"
DATABASE_PATH = "database"

EMBEDDING_MODEL = "text-embedding-ada-002"
EMBEDDING_TOKENIZER = "cl100k_base"

OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]

# %% TYPINGS

Collection = chromadb.Collection

# %% FUNCTIONS


def get_database_client(path: str) -> chromadb.API:
    """Get a persistent client to the Chroma DB."""
    settings = chromadb.Settings(
        allow_reset=True,
        anonymized_telemetry=False,
    )
    return chromadb.PersistentClient(path=path, settings=settings)


def get_embedding_function(
    model_name: str = EMBEDDING_MODEL, api_key: str = OPENAI_API_KEY
) -> embedding_functions.EmbeddingFunction:
    """Get the embedding function for Chroma DB collections."""
    return embedding_functions.OpenAIEmbeddingFunction(
        model_name=model_name,
        api_key=api_key,
    )