import logging
from typing import Optional, Tuple

import backoff
from elasticsearch import AsyncElasticsearch, exceptions

from config import ELASTIC_SEARCH_URL, ELASTIC_SEARCH_ENV

log = logging.getLogger(__name__)


class Elasticsearch:
    MAX_PAGE_SIZE = 1000
    REQUEST_TIMEOUT = 60

    def __init__(self, url: str):
        self.conn: AsyncElasticsearch = AsyncElasticsearch(hosts=[url])

    def get_index(self, index: str) -> str:
        if not index.startswith(ELASTIC_SEARCH_ENV):
            index = f"{ELASTIC_SEARCH_ENV}_{index}"
        return index

    async def put_alias(self, index: str, name: str):
        await self.conn.indices.put_alias(
            index=self.get_index(index), name=self.get_index(name), request_timeout=self.REQUEST_TIMEOUT
        )

    async def index_delete(self, index: str):
        await self.conn.indices.delete(index=self.get_index(index), ignore_unavailable=True)

    async def index_create(self, index: str):
        await self.conn.indices.create(index=self.get_index(index))

    @backoff.on_exception(backoff.expo, Exception)
    async def add_documents(self, dframe: list, refresh: bool = False):
        await self.conn.bulk(body=dframe, refresh=refresh, request_timeout=self.REQUEST_TIMEOUT)

    async def refresh(self, index: str):
        await self.conn.indices.refresh(index=self.get_index(index), request_timeout=self.REQUEST_TIMEOUT)

    async def index(self, index: str, data: list, force: bool = False):
        if not data:
            return

        index = self.get_index(index)

        if force:
            await self.index_delete(index)
            await self.index_create(index)

        dframe = []
        while True:
            record = data.pop()

            dframe.append({"create": {"_index": index, "_id": record.pop("id")}})
            dframe.append(record)

            if len(dframe) == self.MAX_PAGE_SIZE or not data:
                await self.add_documents(dframe)
                if not data:
                    return

                dframe = []

    async def close(self):
        await self.conn.close()

    @backoff.on_exception(backoff.expo, Exception)
    async def count(self, **kwargs):
        kwargs["index"] = self.get_index(kwargs["index"])
        kwargs["body"] = {
            "query": kwargs.pop("query", {}),
        }  # hack for elasticsearch v7, just remove it for lib v8
        result = await self.conn.count(**kwargs, request_timeout=self.REQUEST_TIMEOUT)
        return result

    @backoff.on_exception(backoff.expo, Exception)
    async def scroll(self, scroll_id: str) -> Tuple[int, list, Optional[str]]:
        result: dict = await self.conn.scroll(scroll_id=scroll_id, request_timeout=self.REQUEST_TIMEOUT)

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

    @backoff.on_exception(backoff.expo, Exception)
    async def search(self, **kwargs) -> Tuple[int, list, Optional[str]]:
        if kwargs.get("size", 0) + kwargs.get("from", 0) > 10000:
            raise ValueError("Size + from must be less than 10000")

        kwargs["index"] = self.get_index(kwargs["index"])
        kwargs["body"] = {
            "query": kwargs.pop("query", {}),
            "sort": kwargs.pop("sort", {}),
        }  # hack for elasticsearch v7, just remove it for lib v8
        result: dict = await self.conn.search(**kwargs, request_timeout=self.REQUEST_TIMEOUT)

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

    @backoff.on_exception(backoff.expo, Exception)
    async def get(self, **kwargs) -> dict:
        kwargs["index"] = self.get_index(kwargs["index"])
        try:
            return await self.conn.get(**kwargs, request_timeout=self.REQUEST_TIMEOUT)
        except exceptions.NotFoundError:
            return {}


ES_INSTANCE = Elasticsearch(url=ELASTIC_SEARCH_URL)