from collections import defaultdict
from typing import List, Optional, Tuple

from server.client import services
from server.client.utils import prepare_list_arg
from server.config import VENDOR_IS_SONY_DATA_TTL
from server.constants import DSP


@prepare_list_arg("track_id_list")
async def check_tracks_is_sony(
    track_id_list: List[str],
    market: str,
    album_list: Optional[List[dict]] = None,
    with_data: bool = False,
    with_upc: bool = False,
) -> List[str] or Tuple[List[str], List[dict], List[dict]]:
    """Check tracks are Sony or not.

    Args:
        track_id_list: Track ID list.
        market: Market code.
        album_list: Album list.
        with_data: Return tracks and albums metadata.
        with_upc: Return track ids and id-to-upc mapping.

    Returns:
        Sony tracks ID list.

    """
    track_list = await services.vendor.get_tracks(
        DSP.SPOTIFY, ids=track_id_list, market=market, data_ttl=VENDOR_IS_SONY_DATA_TTL
    )

    album_mapping = defaultdict(list)
    for track in track_list:
        album_id = track.get("album", {}).get("id")
        if album_id:
            album_mapping[album_id].append(track["id"])

    if not album_list:
        album_list = await services.vendor.get_albums(
            DSP.SPOTIFY, ids=list(album_mapping.keys()), market=market, data_ttl=VENDOR_IS_SONY_DATA_TTL
        )

    upc_mapping = {a.get("external_ids", {}).get("upc"): a["id"] for a in album_list}

    sony_upc_list = await services.apollo.get_albums_is_sony(upc=list(upc_mapping.keys()), market=market)

    sony_spotify_ids = []
    for upc in sony_upc_list:
        sony_spotify_ids.extend(album_mapping[upc_mapping[upc]])

    if with_data:
        return sony_spotify_ids, track_list, album_list
    if with_upc:
        track_id_to_upc = {}
        for upc in sony_upc_list:
            for track_id in album_mapping[upc_mapping[upc]]:
                track_id_to_upc[track_id] = upc
        return sony_spotify_ids, track_id_to_upc
    return sony_spotify_ids