from typing import Any, Dict, Iterable, Optional, Tuple

import backoff
from elasticsearch import Elasticsearch, TransportError
from more_itertools import ichunked

from config import ELASTIC_SEARCH_URL
from constants import BACKOFF_TIMEOUT, ES_DEFAULT_SCROLL_SIZE

from .db_connector import DBConnector


__all__ = ["ElasticsearchConnector", "ElasticsearchError"]


class ElasticsearchError(Exception):
    pass


class ElasticsearchConnector(DBConnector):
    url = ELASTIC_SEARCH_URL

    max_page_size: int = 2000
    request_timeout: int = 60

    def close(self):
        self._conn.close()

    def _get_connection(self):
        return Elasticsearch(hosts=[self.url])

    def check_db_connection(self) -> bool:
        return self._conn.ping()

    @backoff.on_exception(backoff.expo, TransportError, max_time=BACKOFF_TIMEOUT)
    def put_alias(self, index: str, name: str):
        self._conn.indices.put_alias(index=index, name=name, request_timeout=self.request_timeout)

    @backoff.on_exception(backoff.expo, TransportError, max_time=BACKOFF_TIMEOUT)
    def index_create(self, index: str, param: dict):
        self._conn.indices.create(index=index, body=param)

    @backoff.on_exception(backoff.expo, TransportError, max_time=BACKOFF_TIMEOUT)
    def index_delete(self, index: str):
        self._conn.indices.delete(index=index, ignore_unavailable=True)

    @backoff.on_exception(backoff.expo, TransportError, max_time=BACKOFF_TIMEOUT)
    def add_documents(self, dframe: list, refresh: bool = False):
        results = self._conn.bulk(body=dframe, refresh=refresh, request_timeout=self.request_timeout)

        if results.get("errors"):
            errors = []
            for result in results["items"]:
                if result["create"]["status"] != 201:
                    errors.append(result)

            raise ElasticsearchError(f"Bulk indexing failed. {len(errors)} errors: {errors}")

    @backoff.on_exception(backoff.expo, TransportError, max_time=BACKOFF_TIMEOUT)
    def refresh(self, index: str):
        self._conn.indices.refresh(index=index, request_timeout=self.request_timeout)

    @backoff.on_exception(backoff.expo, TransportError, max_time=BACKOFF_TIMEOUT)
    def count(self, index: str, body: dict) -> dict:
        result = self._conn.count(index=index, body=body, request_timeout=self.request_timeout)
        return result

    @backoff.on_exception(backoff.expo, TransportError, max_time=BACKOFF_TIMEOUT)
    def scroll(self, scroll_id: str) -> Tuple[list, Optional[str]]:
        result: dict = self._conn.scroll(scroll_id=scroll_id, request_timeout=self.request_timeout)

        return result["hits"]["hits"], scroll_id

    @backoff.on_exception(backoff.expo, TransportError, max_time=BACKOFF_TIMEOUT)
    def search(self, index: str, body: dict, size: int = ES_DEFAULT_SCROLL_SIZE) -> Tuple[list, Optional[str]]:
        if size > 10000:
            raise ValueError("Size + from must be less than 10000")
        result: dict = self._conn.search(
            index=index, body=body, size=size, scroll="1m", request_timeout=self.request_timeout
        )

        return result["hits"]["hits"], result.get("_scroll_id")

    def index(self, index: str, records: Iterable[Dict[str, Any]]) -> int:
        inserted = 0
        for chunk in ichunked(records, self.max_page_size):
            dframe = []

            for record in chunk:
                dframe.append({"create": {"_index": index, "_id": record["id"]}})
                dframe.append(record["value"])
                inserted += 1

            self.add_documents(dframe)

        return inserted