from apollo_utils.core.constants import ALL
from apollo_utils.core.constants.market import Market
from collections import defaultdict
from typing import Dict, List

from server.utils import sort_limit_dict


def calc_artist_markets_top(
    streams_data: List[dict], combine: bool, include_worldwide: bool, limit: int
) -> Dict[str, List[dict]]:
    """Calculate artist markets top.

    Args:
        streams_data: Delphi streams data.
        combine: Combine DSP.
        include_worldwide: Include worldwide or not.
        limit: Top items count limit.

    Returns:
        Artist markets top.

    """
    market_metrics = defaultdict(lambda: defaultdict(int))
    for item in streams_data:
        country_code = item["country_code"]
        if include_worldwide or country_code != Market.WORLDWIDE:
            market_metrics[ALL if combine else item["dsp"]][country_code] += item.get("streams", 0)

    return {
        dsp: [{"market": market, "streams": streams} for market, streams in sort_limit_dict(markets_streams, limit)]
        for dsp, markets_streams in market_metrics.items()
    }