import time
import random
import asyncio
from typing import Iterator
from datetime import date, datetime, timedelta
from functools import partial
from itertools import chain
from collections import defaultdict

from pony.orm import get, select, db_session
from pony.orm.core import CacheIndexError

from unsplash.errors import UnsplashError, UnsplashConnectionError

from ... import logger
from ...sql import SQL
from ...util import normalize_features
from ...cache import City, Genre, Artist, Country, Playlist, ImageMixin, format_param
from ...asynch import LimitedAsCompletedError, limited_as_completed
from ...constants import TimeRange


class RecommenderMixin:
    USER_LIST = ("particledetector", "thesoundsofspotify")

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    @staticmethod
    def _get_upsert_genres_query(genres, conn):
        genre_names = ",\n".join(f"({format_param(genre)})" for genre in genres)
        return conn.execute(
            f"""INSERT INTO genres (name)
                VALUES {genre_names}
                ON CONFLICT DO NOTHING"""
        )

    @staticmethod
    def _get_upsert_countries_query(countries, conn):
        country_values = ",\n".join(
            f"({format_param(country.alpha_2)}, {format_param(country.name)})"
            for country in countries
        )
        return conn.execute(
            f"""INSERT INTO countries (code, name)
            VALUES {country_values}
            ON CONFLICT DO NOTHING"""
        )

    @staticmethod
    def _get_upsert_cities_query(cities, conn):
        city_values = ",\n".join(
            f"({format_param(city)}, {format_param(country)})"
            for city, country in cities
        )
        return conn.execute(
            f"""INSERT INTO cities (name, country)
                VALUES {city_values}
                ON CONFLICT DO NOTHING"""
        )

    @staticmethod
    def _get_upsert_playlists_queries(playlist_dicts, iso_countries):
        playlist_items = defaultdict(list)
        for p in playlist_dicts:
            if p.get("country"):
                p["country"] = iso_countries[p["country"]].alpha_2
            playlist_items[tuple(p.keys())].append(p.values())

        queries = []
        for columns, value_list in playlist_items.items():
            column_string = ", ".join(columns)
            param_string = ",".join(
                f"({', '.join(format_param(p) for p in values)})"
                for values in value_list
            )
            queries.append(
                f"""INSERT INTO playlists AS p ({column_string})
                VALUES {param_string}
                ON CONFLICT DO NOTHING"""
            )

        return queries

    @staticmethod
    async def _get_unsplash_image_requests(conn, genres, countries, cities):
        image_requests = []
        if genres:
            image_genres = {
                r[0]
                for r in await conn.fetch(
                    "SELECT DISTINCT genre FROM images WHERE genre IS NOT NULL"
                )
            }
            new_genres = set(genres) - image_genres
            image_requests += [
                partial(Genre.get_image_fields, image_key=genre, genre=genre)
                for genre in new_genres
            ]

        if countries:
            image_countries = {
                r[0]
                for r in await conn.fetch(
                    "SELECT DISTINCT country FROM images WHERE country IS NOT NULL"
                )
            }
            country_dict = {c.alpha_2: c.name for c in countries.values()}
            new_countries = country_dict.keys() - image_countries
            image_requests += [
                partial(
                    Country.get_image_fields, image_key=country_dict[code], country=code
                )
                for code in new_countries
            ]
        if cities:
            image_cities = {
                r[0]
                for r in await conn.fetch(
                    "SELECT DISTINCT city FROM images WHERE city IS NOT NULL"
                )
            }
            city_country_mapping = dict(cities)
            new_cities = city_country_mapping.keys() - image_cities
            image_requests += [
                partial(
                    City.get_image_fields,
                    image_key=f"{city} {city_country_mapping[city]}",
                    city=city,
                )
                for city in new_cities
            ]

        return image_requests

    async def _upsert_remaining_images(self, reqs, initial_reqs):
        minutes = 62 - datetime.utcnow().minute
        delay = (minutes) * 60
        logger.info("Retrying image upsert in %d minutes", minutes)
        await asyncio.sleep(delay)
        await self._upsert_images(reqs, initial_reqs=initial_reqs)

    async def _upsert_images(self, reqs, conn=None, initial_reqs=None):
        concurrency_limit = 100
        async with self.async_db_session(conn=conn) as conn:
            if not isinstance(reqs, Iterator):
                reqs_iterator = (fetch() for fetch in reqs)
            else:
                reqs_iterator = reqs

            try:
                async for resp in limited_as_completed(
                    reqs_iterator,
                    concurrency_limit,
                    ignore_exceptions=(UnsplashError, UnsplashConnectionError),
                ):
                    if not resp:
                        continue

                    image_fields, updated_fields = resp
                    await ImageMixin.upsert_unsplash_image(
                        conn, image_fields, **updated_fields
                    )
            except LimitedAsCompletedError as exc:
                for future in exc.remaining_futures:
                    future.cancel()

                await asyncio.sleep(2)
                for future in exc.remaining_futures:
                    if future.done():
                        try:
                            _ = future.exception()
                        except:
                            pass

                remaining = len(list(reqs_iterator)) + concurrency_limit
                initial_reqs = initial_reqs or reqs

                if exc.original_exc:
                    logger.exception(exc.original_exc)
                logger.warning(
                    "Remaining images to fetch: %d/%d", remaining, len(initial_reqs)
                )

                asyncio.get_event_loop().create_task(
                    self._upsert_remaining_images(
                        (fetch() for fetch in initial_reqs[-remaining:]), initial_reqs
                    )
                )

    # pylint: disable=too-many-locals
    async def fetch_playlists_pg(self, conn=None):
        async with self.async_db_session(conn=conn) as dbconn:
            existing_playlists = {
                p[0] for p in await dbconn.fetch("SELECT id FROM playlists")
            }
            user_ids_str = ",\n".join(f"('{user_id}')" for user_id in self.USER_LIST)
            await dbconn.execute(
                f"""
                INSERT INTO spotify_users (id)
                VALUES {user_ids_str}
                ON CONFLICT DO NOTHING
            """
            )

        playlist_dicts = []
        for user in self.USER_LIST:
            user_playlists = await self.user_playlists(user)
            async for playlist in user_playlists.iterall(ignore_exceptions=True):
                if not playlist:
                    continue

                logger.info("Got %s", playlist.name)
                if playlist.id not in existing_playlists:
                    playlist_dicts.append(Playlist.from_dict_pg(playlist))

                existing_playlists.add(playlist.id)

        genres = [p["genre"] for p in playlist_dicts if p.get("genre")]
        countries = [p["country"] for p in playlist_dicts if p.get("country")]
        iso_countries = {
            country: Country.get_iso_country(country) for country in countries
        }
        cities = [(p["city"], p["country"]) for p in playlist_dicts if p.get("city")]

        async with self.async_db_session(conn=conn) as dbconn:
            if genres:
                await self._get_upsert_genres_query(genres, dbconn)
            if countries:
                await self._get_upsert_countries_query(iso_countries.values(), dbconn)
            if cities:
                await self._get_upsert_cities_query(cities, dbconn)

            playlist_queries = self._get_upsert_playlists_queries(
                playlist_dicts, iso_countries
            )
            for query in playlist_queries:
                await dbconn.execute(query)

            image_requests = await self._get_unsplash_image_requests(
                dbconn, genres, iso_countries, cities
            )
            await self._upsert_images(image_requests, conn=dbconn)

    async def fetch_playlists(self):
        with db_session:
            fetched_ids = set(select(p.id for p in Playlist))
            for user in self.USER_LIST:
                user_playlists = await self.user_playlists(user)
                async for playlist in user_playlists.iterall(ignore_exceptions=True):
                    if not playlist:
                        continue

                    logger.info("Got %s", playlist.name)
                    if playlist.id not in fetched_ids:
                        Playlist.from_dict(playlist)
                    fetched_ids.add(playlist.id)

    async def fetch_user_top(self, time_range):
        with db_session:
            self.user.top_artists.clear()
            top_artists = await self.current_user_top_artists(
                limit=50, time_range=time_range
            )
            async for artist in top_artists.iterall():
                if self.is_disliked_artist(artist):
                    continue

                if Artist.exists(id=artist.id):
                    self.user.top_artists.add(Artist[artist.id])
                else:
                    try:
                        artist = await Artist.from_dict_async(artist)
                    except CacheIndexError:
                        artist = Artist[artist.id]
                    self.user.top_artists.add(artist)
            self.user.top_genres = (
                self.user.top_artists.genres.distinct().keys()
                - self.user.disliked_genres
            )
            if self.user.top_expires_at is None:
                self.user.top_expires_at = {}
            self.user.top_expires_at[TimeRange(time_range).value] = time.mktime(
                (date.today() + timedelta(days=1)).timetuple()
            )

    def genre_playlist(
        self, genre, popularity=Playlist.Popularity.SOUND
    ):  # pylint: disable=no-self-use
        with db_session:
            if not Playlist.exists():
                raise Exception('You have to call "fetch_playlists" first.')

            return Playlist.get(
                genre=genre, popularity=Playlist.Popularity(popularity).value
            )

    async def top_artists_pg(
        self, time_range=TimeRange.SHORT_TERM, ignore=None, limit=None, conn=None
    ):
        async with self.async_db_session(conn=conn) as conn:
            dislikes_stmt = await conn.prepare(SQL.user_artist_genre_dislikes)
            dislikes = await dislikes_stmt.fetch(self.user_id)
            disliked_artists = {row[0][:-3] for row in dislikes if row[0][-2:] == "AR"}
            disliked_genres = {row[0][:-3] for row in dislikes if row[0][-2:] == "GE"}
            top_artists = await self.current_user_top_artists(
                limit=50, time_range=time_range
            )
            ignore = set(ignore or [])
            top_artists = [
                artist
                for artist in top_artists
                if not self.is_disliked_artist(
                    artist, disliked_artists, disliked_genres
                )
                and artist.id not in ignore
            ]
            if limit:
                top_artists = random.sample(top_artists, min(limit, len(top_artists)))

            return top_artists

    async def top_genres_pg(
        self, time_range=TimeRange.SHORT_TERM, ignore=None, limit=None, conn=None
    ):
        artists = await self.top_artists_pg(time_range=time_range, conn=conn)
        genres = set(chain.from_iterable(artist.genres for artist in artists))

        if ignore:
            genres -= set(ignore)
        if limit:
            genres = random.sample(genres, min(limit, len(genres)))

        return list(genres)

    async def top_artists(self, time_range=TimeRange.SHORT_TERM):
        with db_session:
            if self.user.top_expired(time_range):
                await self.fetch_user_top(time_range)
            return self.user.top_artists

    async def top_genres(self, time_range=TimeRange.SHORT_TERM):
        with db_session:
            if self.user.top_expired(time_range):
                await self.fetch_user_top(time_range)
            return self.user.top_genres

    async def order_by(self, features, tracks):
        audio_features = await self.audio_features(tracks=tracks)
        track_ids = [a.id for a in audio_features]
        audio_features = [
            {feature: a[feature] for feature in features.keys()} for a in audio_features
        ]
        audio_features = normalize_features(audio_features, track_ids)
        for feature, direction in features.items():
            audio_features[feature] *= direction
        audio_features["total"] = audio_features.sum(axis=1)
        return audio_features.sort_values("total").index.tolist()

    async def fill_with_related_artists(self, artists, limit=5):
        tries = 5
        if len(artists) >= limit:
            return artists

        artist_set = set(artists)
        artist_list = list(artists)
        while len(artist_set) < limit and tries:
            tries -= 1
            related_artists = {
                a.id
                for a in (await self.artist_related_artists(random.choice(artist_list)))
                if self.is_not_disliked_artist(a)
            }
            related_artists_limit = min(
                random.randint(1, limit - len(artist_set)), len(related_artists)
            )
            artist_set |= set(random.sample(related_artists, related_artists_limit))
        return artist_set

    async def recommend_by_top_artists(
        self,
        artist_limit=2,
        track_limit=100,
        use_related=True,
        features_order=None,
        time_range=TimeRange.SHORT_TERM,
        **kwargs,
    ):
        """Get a list of recommended songs.

        Returns:
            list: List of tracks
        """
        with db_session:
            top_artists = await self.top_artists(time_range=time_range)
            artists = top_artists.select().without_distinct().random(artist_limit)
            if use_related:
                artists = await self.fill_with_related_artists([a.id for a in artists])
            tracks = await self.recommendations(
                seed_artists=artists, limit=track_limit, **kwargs
            )
            tracks = list(filter(self.is_not_disliked_track, tracks))
            if features_order:
                tracks = await self.order_by(features_order, tracks)
            return tracks

    def is_disliked_artist(self, artist, disliked_artists=None, disliked_genres=None):
        if disliked_artists is not None and disliked_genres is not None:
            return artist.id in disliked_artists or bool(
                set(artist.genres or []) & disliked_genres
            )
        with db_session:
            return artist.id in set(
                self.user.disliked_artists.id.distinct().keys()
            ) or bool(set(artist.genres or []) & set(self.user.disliked_genres))

    def is_not_disliked_artist(self, artist):
        with db_session:
            return not self.is_disliked_artist(artist)

    def is_not_disliked_track(self, track):
        with db_session:
            return all(self.is_not_disliked_artist(a) for a in track.artists)

    def disliked_artists(self):
        with db_session:
            return list(self.user.disliked_artists)

    def disliked_genres(self):
        with db_session:
            return list(self.user.disliked_genres)

    async def dislike_artist(self, artist):
        with db_session:
            if isinstance(artist, str):
                # pylint: disable=consider-using-in
                artist = get(a for a in Artist if a.id == artist or a.name == artist)
            elif isinstance(artist, dict):
                artist = Artist.get(artist.id) or (await Artist.from_dict_async(artist))
            self.user.dislike(artist=artist)

    def dislike_genre(self, genre):
        with db_session:
            if isinstance(genre, str):
                genre = genre.lower()
                genre = Genre.get(name=genre) or Genre(name=genre)
            self.user.dislike(genre=genre)
