from apollo_utils.core.constants.dsp import DSP, DSP_SPOTIFY_APPLE_AMAZON
from collections import defaultdict
from datetime import date, timedelta
from typing import Optional

from server.client import services
from server.client.utils import str_to_date


async def get_two_weeks_and_last_day_streams(
    isrc: str, dsp: Optional[DSP], country_code_list: list[str], latest_date: date
) -> dict[str or None, dict[str, dict[str, int]]]:
    """Calculate last day, previous and last week streams for a track.

    Args:
        isrc: Track ISRC.
        dsp: DSP code.
        country_code_list: Country code list.
        latest_date: Latest streams date.

    Returns:
        DSP or all to country code to streams dict mapping.

    """
    if not country_code_list:
        return {}

    week_start_date = latest_date - timedelta(days=6)
    streams_data = await services.dsp.get_streams(
        isrc=isrc,
        start_date=week_start_date - timedelta(days=7),
        end_date=latest_date,
        country_code=country_code_list,
        dsp=dsp.value if dsp else DSP_SPOTIFY_APPLE_AMAZON.values(),
        group_by="date",
    )

    streams_map = defaultdict(
        lambda: defaultdict(lambda: {"streams_last_day": 0, "streams_last_week": 0, "streams_previous_week": 0})
    )
    for item in streams_data:
        streams_date = str_to_date(item["date"])
        streams_count = item.get("streams", 0)
        map_item = streams_map[item["dsp"] if dsp else dsp][item["country_code"]]
        if streams_date == latest_date:
            map_item["streams_last_day"] += streams_count
        map_item["streams_previous_week" if streams_date < week_start_date else "streams_last_week"] += streams_count

    return streams_map