import json
import time
from typing import Dict, Optional, Sequence, Callable
from chromadb.api import API
from chromadb.api.types import GetResult, QueryResult
from chromadb.api.models.Collection import Collection

import re

# mimics s3 bucket requirements for naming
def is_valid_index_name(index_name):
    if len(index_name) < 3 or len(index_name) > 63:
        return False
    if not re.match("^[a-z0-9][a-z0-9._-]*[a-z0-9]$", index_name):
        return False
    if ".." in index_name:
        return False
    if re.match("^[0-9]{1,3}.[0-9]{1,3}.[0-9]{1,3}.[0-9]{1,3}$", index_name):
        return False
    return True


class LocalAPI(API):
    def __init__(self, settings, db):
        self._db = db

    def heartbeat(self):
        return int(1000 * time.time_ns())

    #
    # COLLECTION METHODS
    #
    def create_collection(
        self,
        name: str,
        metadata: Optional[Dict] = None,
        embedding_fn: Optional[Callable] = None,
    ) -> Collection:
        if not is_valid_index_name(name):
            raise ValueError("Invalid index name: %s" % name)  # NIT: tell the user why

        self._db.create_collection(name, metadata)
        return Collection(client=self, name=name, embedding_fn=embedding_fn)

    def get_collection(
        self,
        name: str,
    ) -> Collection:
        self._db.get_collection(name)
        return Collection(client=self, name=name)

    def _get_collection_db(self, name: str) -> int:
        return self._db.get_collection(name)

    def list_collections(self) -> Sequence[Collection]:
        collections = []
        db_collections = self._db.list_collections()
        for db_collection in db_collections:
            collections.append(Collection(client=self, name=db_collection[1]))
        return collections

    def modify(
        self,
        current_name: str,
        new_name: Optional[str] = None,
        new_metadata: Optional[Dict] = None,
    ):
        # NIT: make sure we have a valid name like we do in create
        if new_name is not None:
            if not is_valid_index_name(new_name):
                raise ValueError("Invalid index name: %s" % new_name)

        self._db.update_collection(current_name, new_name, new_metadata)

    def delete_collection(self, name: str) -> int:
        return self._db.delete_collection(name)

    #
    # ITEM METHODS
    #
    def _add(
        self,
        ids,
        collection_name: str,
        embeddings,
        metadatas=None,
        documents=None,
        increment_index=True,
    ):

        number_of_embeddings = len(embeddings)

        # fill in empty objects if not provided
        if metadatas is None:
            if isinstance(embeddings[0], list):
                metadatas = [{} for _ in range(number_of_embeddings)]
            else:
                metadatas = {}

        if ids is None:
            if isinstance(embeddings[0], list):
                ids = [None for _ in range(number_of_embeddings)]
            else:
                ids = None

        if documents is None:
            if isinstance(embeddings[0], list):
                documents = [None for _ in range(number_of_embeddings)]
            else:
                documents = None

        # convert all metadatas values to strings : TODO: handle this better
        # this is currently here because clickhouse-driver does not support json
        if isinstance(embeddings[0], list):
            for m in metadatas:
                for k, v in m.items():
                    m[k] = str(v)
        else:
            for k, v in metadatas.items():
                metadatas[k] = str(v)

        # convert to array for downstream processing
        if not isinstance(embeddings[0], list):
            embeddings = [embeddings]
            metadatas = [metadatas]
            documents = [documents]
            ids = [ids]

        collection_uuid = self._db.get_collection_uuid_from_name(collection_name)
        added_uuids = self._db.add(
            collection_uuid, embedding=embeddings, metadata=metadatas, documents=documents, ids=ids
        )

        if increment_index:
            self._db.add_incremental(collection_uuid, added_uuids, embeddings)

        return True  # NIT: should this return the ids of the succesfully added items?

    def _update(self, collection_name: str, embedding, metadata=None):

        number_of_embeddings = len(embedding)

        if metadata is None:
            metadata = [{} for _ in range(number_of_embeddings)]

        # convert all metadata values to strings : TODO: handle this better
        # this is currently here because clickhouse-driver does not support json
        for m in metadata:
            for k, v in m.items():
                m[k] = str(v)

        collection_uuid = self._db.get_collection_uuid_from_name(collection_name)

        # find the uuids of the embeddings where the metadata matches
        # then update the embeddings for that
        # then update the index position for that embedding
        for item in metadata:
            uuid = self._db.get_uuid(collection_uuid, item)
            if uuid is not None:
                self._db.update(collection_uuid, uuid, embedding, metadata)

        # added_uuids = self._db.add(collection_uuid, embedding, metadata)
        # self._db.add_incremental(collection_uuid, added_uuids, embedding)

        return True

    def _db_result_to_get_result(self, db_result) -> GetResult:
        query_result = GetResult(embeddings=[], documents=[], ids=[], metadatas=[])
        for entry in db_result:
            query_result["embeddings"].append(entry[2])
            query_result["documents"].append(entry[3])
            query_result["ids"].append(entry[4])
            query_result["metadatas"].append(json.loads(entry[5]))
        return query_result

    def _get(
        self,
        collection_name,
        ids=None,
        where=None,
        sort=None,
        limit=None,
        offset=None,
        page=None,
        page_size=None,
    ):

        if where is None:
            where = {}

        if page and page_size:
            offset = (page - 1) * page_size
            limit = page_size

        return self._db_result_to_get_result(
            self._db.get(
                collection_name=collection_name,
                ids=ids,
                where=where,
                sort=sort,
                limit=limit,
                offset=offset,
            )
        )

    def _delete(self, collection_name, ids=None, where=None):

        if where is None:
            where = {}

        deleted_uuids = self._db.delete(collection_name=collection_name, where=where, ids=ids)
        return deleted_uuids

    def _count(self, collection_name):

        return self._db.count(collection_name=collection_name)

    def reset(self):
        self._db.reset()
        return True

    def _query(self, collection_name, query_embeddings, n_results=10, where={}):
        uuids, distances = self._db.get_nearest_neighbors(
            collection_name=collection_name,
            where=where,
            embeddings=query_embeddings,
            n_results=n_results,
        )

        query_result = QueryResult(embeddings=[], documents=[], ids=[], metadatas=[], distances=[])
        for i in range(len(uuids)):
            embeddings = []
            documents = []
            ids = []
            metadatas = []
            db_result = self._db.get_by_ids(uuids[i])

            for entry in db_result:
                embeddings.append(entry[2])
                documents.append(entry[3])
                ids.append(entry[4])
                metadatas.append(json.loads(entry[5]))

            query_result["embeddings"].append(embeddings)
            query_result["documents"].append(documents)
            query_result["ids"].append(ids)
            query_result["metadatas"].append(metadatas)
            query_result["distances"].append(distances[i].tolist())

        return query_result

    def raw_sql(self, raw_sql):

        return self._db.raw_sql(raw_sql)

    def create_index(self, collection_name):

        collection_uuid = self._db.get_collection_uuid_from_name(collection_name)
        self._db.create_index(collection_uuid=collection_uuid)
        return True

    def _peek(self, collection_name, n=10):
        return self._get(collection_name=collection_name, limit=n)
