import asyncio
from apollo_utils.core.constants.market import Market
from collections import defaultdict
from datetime import date
from typing import List, Optional

from server import config
from server.cache.utils import cached
from server.client import services
from server.constants import SPOTIFY_ARTIST_URI_PREFIX, UNIX_EPOCH_START_DATE
from server.constants.common import SortOrder


@cached(ttl=config.ARTIST_GRAS_ID_TLL)
async def get_artist_gras_id(spotify_artist_id: str) -> Optional[str]:
    """Get artist GRAS ID by Spotify ID.

    Args:
        spotify_artist_id: Spotify artist ID,

    Returns:
        GRAS ID or None.

    """
    artists_data = await services.dsp.get_search(
        index="artist", query=f'spotify_artist_id:"{SPOTIFY_ARTIST_URI_PREFIX}{spotify_artist_id}"'
    )

    if not artists_data:
        return

    return artists_data[0]["artist_id"]


async def get_artist_streams(
    gras_artist_id: str, dsp_list: List[str], country_code_list: List[str], start_date: date, end_date: date
) -> List[dict]:
    """Get artist streams.

    Args:
        gras_artist_id: GRAS artist ID.
        dsp_list: DSP code list.
        country_code_list: Country code list.
        start_date: Date range from.
        end_date: Date range to.

    Returns:
        Country code to date streams list.

    """
    streams_list = await services.dsp.get_streams(
        artist_id=gras_artist_id,
        dsp=dsp_list,
        country_code=country_code_list,
        start_date=start_date,
        end_date=end_date,
        group_by="date",
        include="sources",
        sort_by="date",
        sort_order=SortOrder.ASC.value,
    )

    country_code_map = defaultdict(list)

    for item in streams_list:
        country_code_map[item["country_code"]].append(item)

    return [{"country_code": country_code, "data": items_list} for country_code, items_list in country_code_map.items()]


async def calc_artist_streams_totals(
    gras_artist_id: str, dsp_list: List[str], country_code: Optional[str], end_date: date
) -> dict:
    """Calculate artist streams totals.

    Args:
        gras_artist_id: GRAS artist ID.
        dsp_list: DSP code list.
        country_code: Country code.
        end_date: End (latest streams) date.

    Returns:
        Artist totals.

    """
    task_list = [
        services.dsp.get_streams(
            artist_id=gras_artist_id,
            dsp=dsp_list,
            start_date=UNIX_EPOCH_START_DATE,
            end_date=end_date,
            country_code=[Market.WORLDWIDE] + ([country_code] if country_code else []),
        )
    ]

    for dsp in dsp_list:
        task_list.append(
            services.dsp.get_streams(
                artist_id=gras_artist_id,
                dsp=dsp,
                start_date=UNIX_EPOCH_START_DATE,
                end_date=end_date,
                group_by="date",
                limit=1,
                sort_by="date",
                sort_order=SortOrder.ASC.value,
            )
        )

    response_list = await asyncio.gather(*task_list)

    streams_data = response_list[0]
    country_streams_map = defaultdict(int)
    for item in streams_data:
        country_streams_map[item["country_code"]] += item.get("streams", 0)

    dsp_date_map = {}
    for index, dsp in enumerate(dsp_list):
        date_data = response_list[index + 1]
        dsp_date_map[dsp] = date_data[0]["date"] if date_data else None

    return {
        "first_stream_date": dsp_date_map,
        Market.WORLDWIDE: country_streams_map[Market.WORLDWIDE],
        "country": country_streams_map[country_code],
    }