from agently.utils import LazyImport, FunctionShifter

LazyImport.import_package("chromadb")

from typing import Callable, TypedDict, Any, TYPE_CHECKING

from chromadb.api.types import EmbeddingFunction, DefaultEmbeddingFunction

if TYPE_CHECKING:
    from chromadb.api.types import (
        Documents,
        Embedding,
        Embeddings,
        QueryResult,
        Schema,
        CollectionMetadata,
        DataLoader,
        Loadable,
    )
    from chromadb.api.collection_configuration import CreateCollectionConfiguration
    from chromadb.api import ClientAPI
    from agently.core import BaseAgent


class ChromaDataDictOptional(TypedDict, total=False):
    metadata: dict[Any, Any]
    id: Any
    embedding: "Embedding"


class ChromaDataDict(ChromaDataDictOptional):
    document: str


class ChromaData:
    def __init__(
        self,
        original_data: ChromaDataDict | list[ChromaDataDict],
        *,
        embedding_function: "Callable[[str | list[str]], Embeddings] | None" = None,
        agent: "BaseAgent | None" = None,
    ):
        self._original_data = original_data if isinstance(original_data, list) else [original_data]
        if embedding_function:
            self._embedding_function = embedding_function
        elif agent:

            def embedding_function_by_agent(texts: str | list[str]) -> "Embeddings":
                return agent.input(texts).start()

            self._embedding_function = embedding_function_by_agent
        else:
            self._embedding_function = None
        self._chroma_data = {
            "documents": [],
            "metadatas": [],
            "ids": [],
            "embeddings": [],
        }
        if self._embedding_function is not None:
            documents = [data["document"] for data in self._original_data]
            embeddings = self._embedding_function(documents)
            self._chroma_data["embeddings"] = embeddings
            for index, data in enumerate(self._original_data):
                data["embedding"] = embeddings[index]
        else:
            del self._chroma_data["embeddings"]
        for index, data in enumerate(self._original_data):
            self._chroma_data["documents"].append(data["document"])
            self._chroma_data["metadatas"].append(data["metadata"] if "metadata" in data else None)
            self._chroma_data["ids"].append(data["id"] if "id" in data else str(index))

    def get_kwargs(self):
        return self._chroma_data

    def get_original_data(self):
        return self._original_data


class ChromaResults:
    def __init__(
        self,
        *,
        queries: str | list[str],
        results: "QueryResult",
    ):
        if isinstance(queries, str):
            queries = [queries]
        if len(queries) != len(results["documents"] or []):
            raise ValueError(
                f"The length of queries does not equal the length of results.\nQueries: { queries }\nResults: { results }"
            )
        self._chroma_results = results
        self._results = {}
        for query_index, query in enumerate(queries):
            for result_index, id in enumerate(results["ids"][query_index]):
                result = {
                    "id": id,
                    "document": results["documents"][query_index][result_index] if results["documents"] else None,
                    "metadata": results["metadatas"][query_index][result_index] if results["metadatas"] else None,
                    "distance": results["distances"][query_index][result_index] if results["distances"] else None,
                }
                self._results[query] = result

    def get_original_results(self):
        return self._chroma_results

    def get(self):
        return self._results


class ChromaEmbeddingFunction(EmbeddingFunction):
    def __init__(
        self,
        *,
        embedding_function: "Callable[[list[str]], Embeddings] | None" = None,
        agent: "BaseAgent | None" = None,
    ):
        if embedding_function:
            self.embedding_function = embedding_function
        elif agent:

            def embedding_function_by_agent(texts: list[str]) -> "Embeddings":
                return agent.input(texts).start()

            self.embedding_function = embedding_function_by_agent
        else:
            raise ValueError(
                f"ChromaEmbeddingFunction() requires at least one definition for 'embedding_function' or 'agent'."
            )

    def __call__(self, documents: "Documents") -> "Embeddings":
        return self.embedding_function([document for document in documents])


class ChromaCollection:
    def __init__(
        self,
        conn: "ClientAPI",
        collection_name: str,
        *,
        schema: "Schema | None" = None,
        configuration: "CreateCollectionConfiguration | None" = None,
        metadata: None = None,
        embedding_function: EmbeddingFunction | None = None,
        agent: "BaseAgent | None" = None,
        data_loader: "DataLoader[Loadable] | None" = None,
        get_or_create: bool = False,
    ):
        if embedding_function:
            self._embedding_function = embedding_function
        elif agent:
            self._embedding_function = ChromaEmbeddingFunction(agent=agent)
        else:
            self._embedding_function = DefaultEmbeddingFunction()

        self._collection = conn.create_collection(
            collection_name,
            schema=schema,
            configuration=configuration,
            metadata=metadata,
            embedding_function=embedding_function,
            data_loader=data_loader,
            get_or_create=get_or_create,
        )

    def add(self): ...
