"Connect to ChromaDB"
import logging
import os
import sys
from datetime import datetime

import chromadb
import chromadb.errors as errors
import chromadb.utils.embedding_functions as embedding_functions
from chromadb.config import Settings
from dotenv import load_dotenv

sys.path.append(".")

from configs.config import OPENAI_EF

load_dotenv()


class ChromaDB:
    "Class to query ChromaDB"

    def __init__(self, persist_directory: str = None):

        server_address = os.getenv("chroma_server")
        server_port = server_address.split(":")[1]
        server_ip = server_address.split(":")[0]

        if server_ip != "localhost":
            self.chroma_client = chromadb.HttpClient(
                host=server_ip,
                port=server_port,
                settings=Settings(
                    chroma_client_auth_provider="chromadb.auth.token_authn.TokenAuthClientProvider",
                    chroma_client_auth_credentials=os.getenv("chroma_server_api_key"),
                ),
            )
        else:
            if persist_directory:
                self.chroma_client = chromadb.PersistentClient(path=persist_directory)
            else:
                self.chroma_client = chromadb.PersistentClient(path="./chroma")

        # if persist_directory:
        #     self.chroma_client = chromadb.PersistentClient(path=persist_directory)
        # else:
        #     self.chroma_client = chromadb.PersistentClient(
        #         path="./chroma", server=http_link
        #     )
        self.openai_ef = OPENAI_EF

        self.product_collection = self.chroma_client.get_or_create_collection(
            "company_product_description",
            embedding_function=self.openai_ef,
        )
        self.overview_collection = self.chroma_client.get_or_create_collection(
            "company_overview_description", embedding_function=self.openai_ef
        )
        self.fund_collection = self.chroma_client.get_or_create_collection(
            "funds", embedding_function=self.openai_ef
        )

    def retrive_by_collection(
        self, query_items: list, collection: chromadb.Collection, n_results=10, **kwargs
    ):
        "Get company product description"

        if "where" not in kwargs:
            kwargs["where"] = None
        if "where_document" not in kwargs:
            kwargs["where_document"] = None

        result_items = collection.query(
            query_texts="\n".join(query_items),
            where=kwargs["where"],
            where_document=kwargs["where_document"],
            n_results=n_results,
        )

        return result_items

    def create_where_filter(self, industry_filter, **kwargs):
        """ "
        Create where filter for ChromaDB query
        kwargs: key value pairs
            - key: metadata field
            - value: dict
                Key: filtering (
                    $eq - equal to (string, int, float)
                    $ne - not equal to (string, int, float)
                    $gt - greater than (int, float)
                    $gte - greater than or equal to (int, float)
                    $lt - less than (int, float)
                    $lte - less than or equal to (int, float)
                )
                Value: value to filter by

        """
        where_filter = {}

        # clean-up kwargs to remove any empty values and None values
        kwargs = {k: v for k, v in kwargs.items() if v is not None}

        if len(kwargs) == 0 and industry_filter is not None:
            return industry_filter

        if len(kwargs) == 0 and industry_filter is None:
            return None
        if len(kwargs) == 1 and industry_filter is None:
            return kwargs

        if industry_filter is not None:
            where_filter = {"$and": [industry_filter]}
        else:
            where_filter = {"$and": []}

        for key, value in kwargs.items():
            where_filter["$and"].append({key: value})

        return where_filter

    def add_item(self, collection, item, document_name):
        "Add item into collection"
        if isinstance(self.chroma_client.heartbeat(), int) is False:
            raise ValueError("ChromaDB is not available")

        try:
            collection.add(
                ids=[item["root_url"]],
                documents=[item[document_name]],
                metadatas=[
                    {
                        "industry": item["industry"],
                        "sector": item["sector"],
                        "public": item["public"],
                        "company_url": item["root_url"],
                    }
                ],
            )

        except errors.DuplicateIDError as e:
            print(f"Duplicate ID: {e}")
            collection.update(
                ids=[item["root_url"]],
                documents=[item[document_name]],
                metadatas=[
                    {
                        "industry": item["industry"],
                        "sector": item["sector"],
                        "public": item["public"],
                        "company_url": item["root_url"],
                    }
                ],
            )


if __name__ == "__main__":

    cdb = ChromaDB()

    cdb.chroma_client.list_collections()
    # cdb.chroma_client.create_collection(
    #     "funds",
    #     embedding_function=cdb.openai_ef,
    # )