"""Client for using Elastiknn."""

from typing import Iterable, Tuple, Dict

from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk

from .api import *


class ElastiKnnClient(object):

    def __init__(self, es: Elasticsearch = None):
        """Wrapper on the official `elasticsearch.Elasticsearch` client for making Elastiknn requests.

        The client is fairly primitive in that it assumes the vector is basically the only thing stored in each doc.
        For more complex use-cases you should use this client as an example for constructing the mapping, indexing,
        and search requests in your own code.

        Parameters
        ----------
        es : `elasticsearch.Elasticsearch` client.
            This client is used internally to make all requests.
            Defaults to a client pointing at http://localhost:9200.
        """
        if es is None:
            self.es = Elasticsearch(["http://localhost:9200"])
        else:
            self.es = es

    def put_mapping(self, index: str, field: str, mapping: Mapping.Base):
        """
        Update the mapping at the given index and field to store an Elastiknn vector.

        Parameters
        ----------
        index : string
            Index containing the given field. Should already exist.
        field : string
            Field that should use the given mapping.
        mapping : instance of `Mapping.Base`
            Mapping object defining the vector's storage properties.

        Returns
        -------
        Dict
            Json response as a dict. Successful request returns `{"acknowledged": true}`.
        """
        body = {
            "properties": {
                field: mapping.to_dict()
            }
        }
        return self.es.transport.perform_request("PUT", f"/{index}/_mapping", body=body)

    def index(self, index: str, field: str, vecs: Iterable[Vec.Base], ids: List[str] = None, refresh: bool = False) -> Tuple[int, List[Dict]]:
        """Index (i.e. store) the given vectors at the given index and field with the optional ids.

        Parameters
        ----------
        index : string
            Index where the vectors are stored.
        field : string
            Field containing the vector in each document.
        vecs : List of `Vec.Base`
            Vectors that should be indexed.
        ids : List of strings
            Optional list of ids associated with the given vectors.
            If not given, the ids are generated by Elasticsearch.
        refresh : bool
            Whether to refresh before returning. Set to true if you want to immediately run queries after indexing.

        Returns
        -------
        Int
            Number of vectors successfully indexed.
        List of Dicts
            Error responses for the vectors that were not indexed.
        """
        if ids is None or ids == []:
            ids = [None for _ in vecs]

        def gen():
            for vec, _id in zip(vecs, ids):
                d = { "_op_type": "index", "_index": index, field: vec.to_dict()}
                if _id:
                    d["_id"] = str(_id)
                elif "_id" in d:
                    del d["_id"]
                yield d

        res = bulk(self.es, gen())
        if refresh:
            self.es.indices.refresh(index=index)
        return res

    def nearest_neighbors(self, index: str, query: NearestNeighborsQuery.Base, k: int = 10, fetch_source: bool = False) -> Dict:
        """Build and execute a nearest neighbors query against the given index.

        Parameters
        ----------
        index : string
            Index to run the search against.
        query : NearestNeighborsQuery.Base
            Query object defining the query properties.
        k: int
            Number of hits to return.
        fetch_source : bool
            Whether to return the `_source` of the document. It's generally faster to _not_ return the `_source`.

        Returns
        -------
        Dict
            Standard Elasticsearch search response parsed as a dict.
        """
        body = {
            "query": {
                "elastiknn_nearest_neighbors": query.to_dict()
            }
        }
        return self.es.search(index, body=body, size=k, _source=fetch_source)
