from server.db.constants import ARTIST_SEARCH_INDEX
from server.db.es_connector import ES_INSTANCE


class ESRepo:
    @staticmethod
    async def search(query: dict, offset: int, limit: int, sort: dict = None, index: str = ARTIST_SEARCH_INDEX) -> list:
        _, results, scroll_id = await ES_INSTANCE.search(index=index, query=query, sort=sort, from_=offset, size=limit)

        return [
            result["_source"].update({"id": result["_id"], "score": result["_score"]}) or result["_source"]
            for result in results
        ]

    @staticmethod
    async def count(query: dict, index: str) -> int:
        result = await ES_INSTANCE.count(index=index, query=query)
        return result["count"]

    @staticmethod
    async def get(id: str, index: str) -> dict:
        result: dict = await ES_INSTANCE.get(index=index, id=id)

        return result.get("_source", {})