from apollo_utils.core.constants.dsp import DSP
from apollo_utils.core.constants.market import Market
from typing import Optional

from server.client import services


async def get_markets_ranks(dsp: Optional[DSP] = None) -> dict[str, str] or dict[str, dict[str, str]]:
    """Get market to rank mapping.

    Args:
        dsp: DSP.

    Returns:
        Market to rank mapping or DSP to market to rank.

    """
    response = await services.dsp.get_streams_country_code_rank()

    result = {
        dsp_code: {item["country_code"]: item["rank"] for item in market_rank_list if item.get("rank") is not None}
        for dsp_code, market_rank_list in (((dsp.value, response.get(dsp.value, {})),) if dsp else response.items())
    }

    if dsp is None:
        return result

    return result.get(dsp.value, {})


def set_missing_ranks(market_rank_map: dict[str, int], market_list: list[str]) -> dict[str, int]:
    """Set missing markets ranks.

    Args:
        market_rank_map: Existing market to rank mapping.
        market_list: All markets list.

    Returns:
        Market to rank mapping.

    """
    market_set = set(Market.convert_global(i, Market.WORLDWIDE) for i in market_list)
    missing_markets_set = market_set - set(market_rank_map.keys())
    rank_value = max(market_rank_map.values()) if market_rank_map else 0
    for market in sorted(missing_markets_set):
        rank_value += 1
        market_rank_map[market] = rank_value
    return market_rank_map