import asyncio
from typing import Any, Dict, List, Type, Union
from sqlalchemy import desc, func
from sqlalchemy.future import select

from server.dna.constants import ARTIST, TRACK
from server.dna.helpers.base import BaseSearchesHelper
from server.dna.models import RecentSearches
from server.dna.utils import QueryWrapper
from server.artist.helpers.search_by_name import ArtistSearchByNameHelper, _ArtistSearchByNameHelper
from server.track.helpers.search_by_name import TrackSearchByNameHelper, _TrackSearchByNameHelper


__all__ = ["RecentSearchesHelper"]


class _RecentSearchesHelper(BaseSearchesHelper):
    _data_model = RecentSearches

    def _get_entities_ids(self, recent_searches: List[Dict[str, Any]], entity_type: str) -> List[str]:
        return [
            recent_search["entity_id"]
            for recent_search in recent_searches
            if recent_search["entity_type"] == entity_type
        ]

    async def _prepare_data(
        self, entities_ids: List[str], helper: Union[Type[_ArtistSearchByNameHelper], Type[_TrackSearchByNameHelper]]
    ) -> Union[Dict[str, Dict[str, Any]], List]:

        index = helper._index_name
        data = await self._make_search(entities_ids, index)
        if not data:
            return []

        results = self._update_data(data)
        obj_results = helper.get_data(results)  # type: ignore
        return {result.entity_id: result for result in obj_results}

    async def _get_results(self, recent_searches):
        artist_entitites_ids = self._get_entities_ids(recent_searches, ARTIST)
        track_entitites_ids = self._get_entities_ids(recent_searches, TRACK)

        artists_data, tracks_data = await asyncio.gather(
            *(
                asyncio.wait_for(self._prepare_data(artist_entitites_ids, ArtistSearchByNameHelper), timeout=5),
                asyncio.wait_for(self._prepare_data(track_entitites_ids, TrackSearchByNameHelper), timeout=5),
            )
        )

        if not artists_data and not tracks_data:
            return []

        data = {
            ARTIST: artists_data,
            TRACK: tracks_data,
        }

        return [
            data[recent_search["entity_type"]][recent_search["entity_id"]]
            for recent_search in recent_searches
            if recent_search["entity_id"] in data[recent_search["entity_type"]]
        ]

    async def _get_recent_searches(self, filters, limit=20):
        query = (
            select(
                func.max(self._data_model.id).label("id"),
                (self._data_model.entity_id),
                self._data_model.entity_type,
            )
            .filter_by(**filters)
            .order_by(desc("id"))
            .group_by(self._data_model.entity_id, self._data_model.entity_type)
            .limit(limit)
        )
        return await QueryWrapper.fetchall(query)

    async def get_recent_searches(self, filters):
        recent_searches = await self._get_recent_searches(filters)
        results = await self._get_results(recent_searches)
        return results


RecentSearchesHelper = _RecentSearchesHelper()