from typing import Any, Dict, List, Mapping, Union

from analytics.constants import cache
from analytics.constants.streams import (
    STREAMS_SOS_COLUMNS,
    STREAMS_SOS_COLUMNS_STORE_FILTERS,
)
from analytics.queries.format import format_row
from analytics.queries.sound_recording import (
    SoundRecordingDownloadsTimeSeries,
    SoundRecordingStreamsTimeSeries,
    SoundRecordingSummary,
)
from analytics.utils.cache import cache_in_redis

# TODO implement remaining aggregations
from analytics.utils.streams import (
    breakdown_by_sos,
    calc_date_skip_rate,
    check_store_ids_for_sos,
    get_streams_sos_columns,
    get_streams_sos_columns_from_stream_sources,
)

TOTAL_TIMESERIES_TYPE = "TRACK_STREAMS"

TOTAL_SUMMARY_TYPE = "TOTAL"

AGGREGATION_FIELDS_TIMESERIES = {
    "TRACK_STREAMS_BY_STORE": "store_id",
    "TRACK_STREAMS_BY_COUNTRY": "country_code",
    "TRACK_STREAMS_BY_PRODUCT": "product_id",
    "TRACK_DOWNLOADS_BY_STORE": "store_id",
    "TRACK_DOWNLOADS_BY_COUNTRY": "country_code",
    "TRACK_DOWNLOADS_BY_PRODUCT": "product_id",
}

AGGREGATION_FIELDS_SUMMARY = {
    "TOTAL": "isrc",
    "STORE": "store_id",
    "COUNTRY": "country_code",
    "SOS": "isrc",
    "SOS_V2": "isrc",
    "PRODUCT": "product_id",
}

# map of type -> aggregation type -> table
SOUND_RECORDING_TABLES_TIMESERIES = {
    # Track streams
    "TRACK_STREAMS": {
        "total": "V_STREAMS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
        "country": "V_STREAMS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
    },
    "TRACK_STREAMS_BY_PRODUCT": {
        "total": "V_STREAMS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
        "country": "V_STREAMS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
    },
    "TRACK_STREAMS_BY_COUNTRY": {
        "total": "V_STREAMS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
        "country": "V_STREAMS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
    },
    "TRACK_STREAMS_BY_STORE": {
        "total": "V_STREAMS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
        "country": "V_STREAMS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
    },
    "TRACK_STREAMS_BY_SOS": {
        "total": "V_STREAMS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
        "country": "V_STREAMS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
    },
    "TRACK_STREAMS_BY_SOS_V2": {
        "total": "V_STREAMS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
        "country": "V_STREAMS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
    },
    # Track downloads
    "TRACK_DOWNLOADS": {
        "total": "DOWNLOADS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
        "country": "DOWNLOADS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
    },
    "TRACK_DOWNLOADS_BY_PRODUCT": {
        "total": "DOWNLOADS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
        "country": "DOWNLOADS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
    },
    "TRACK_DOWNLOADS_BY_COUNTRY": {
        "total": "DOWNLOADS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
        "country": "DOWNLOADS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
    },
    "TRACK_DOWNLOADS_BY_STORE": {
        "total": "DOWNLOADS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
        "country": "DOWNLOADS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
    },
}

# map of type -> aggregation type -> table
SOUND_RECORDING_TABLES_SUMMARY = {
    "TOTAL": {
        "total": {
            "streams_table": "V_STREAMS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
            "downloads_table": "DOWNLOADS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
        },
        "country": {
            "streams_table": "V_STREAMS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
            "downloads_table": "DOWNLOADS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
        },
    },
    "COUNTRY": {
        "total": {
            "streams_table": "V_STREAMS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
            "downloads_table": "DOWNLOADS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
        },
        "country": {
            "streams_table": "V_STREAMS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
            "downloads_table": "DOWNLOADS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
        },
    },
    "STORE": {
        "total": {
            "streams_table": "V_STREAMS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
            "downloads_table": "DOWNLOADS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
        },
        "country": {
            "streams_table": "V_STREAMS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
            "downloads_table": "DOWNLOADS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
        },
    },
    "SOS": {
        "total": {
            "streams_table": "V_STREAMS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
            "downloads_table": "DOWNLOADS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
        },
        "country": {
            "streams_table": "V_STREAMS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
            "downloads_table": "DOWNLOADS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
        },
    },
    "SOS_V2": {
        "total": {
            "streams_table": "V_STREAMS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
            "downloads_table": "DOWNLOADS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
        },
        "country": {
            "streams_table": "V_STREAMS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
            "downloads_table": "DOWNLOADS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",
        },
    },
    "PRODUCT": {
        "total": {
            "streams_table": "V_STREAMS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
            "downloads_table": "DOWNLOADS_BY_TRACK_FEED_DISTRIBUTOR_DAILY",
        },
        "country": {
            "streams_table": "V_STREAMS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",  # noqa
            "downloads_table": "DOWNLOADS_BY_TRACK_COUNTRY_FEED_DISTRIBUTOR_DAILY",  # noqa
        },
    },
}


def _get_sum_expression_of_streams_sos_columns(
    store_ids: List[str], stream_sources
) -> str:
    sos_columns = get_streams_sos_columns_from_stream_sources(store_ids, stream_sources)
    sum_expression = "SUM(" + "+".join(sos_columns) + ")"
    return sum_expression


def _get_streams_sos_columns_from_stream_sources(
    store_ids: List[str], stream_sources
) -> List[str]:
    """Deprecated: use get_streams_sos_columns_from_stream_sources from analytics.utils.streams."""
    return get_streams_sos_columns_from_stream_sources(store_ids, stream_sources)


def get_sound_recording_table_for_timeseries(query_type: str, countries: List) -> str:
    key = "total"
    if len(countries) != 0:
        key = "country"
    table_options = SOUND_RECORDING_TABLES_TIMESERIES.get(query_type)
    if not table_options:
        raise NotImplementedError(f"Query {query_type} not supported")
    table = table_options.get(key)
    if not table:
        raise NotImplementedError(
            f"Query {query_type} not supported with " f"type {key}"
        )
    return table


def get_sound_recording_tables_for_summary(
    query_type: str,
    countries: List,
) -> dict:
    key = "total"
    if len(countries) != 0:
        key = "country"

    table_options = SOUND_RECORDING_TABLES_SUMMARY.get(query_type)
    if not table_options:
        raise NotImplementedError(f"Query {query_type} not supported")
    tables = table_options.get(key)
    if not tables:
        raise NotImplementedError(f"Query {query_type} not supported with type {key}")

    return tables


@cache_in_redis(ttl=cache.ONE_DAY)
def get_sound_recording_timeseries(
    query_params: Mapping[str, Any],
    permissions: Mapping[str, Any],
) -> List[Dict]:
    query_type = query_params.get("type", TOTAL_TIMESERIES_TYPE)

    aggregation_field = AGGREGATION_FIELDS_TIMESERIES.get(query_type)
    if aggregation_field:
        query_params["aggregation_field"] = aggregation_field
    countries = query_params.get("countries", [])

    query_params["track_table"] = get_sound_recording_table_for_timeseries(
        query_type, countries
    )

    if query_params["stream_sources"]:
        query_params[
            "streams_sum_expression"
        ] = _get_sum_expression_of_streams_sos_columns(
            check_store_ids_for_sos(query_params.get("store_ids")),
            query_params["stream_sources"],
        )

    if query_type == "TRACK_STREAMS_BY_SOS_V2":
        store_ids = check_store_ids_for_sos(query_params.get("store_ids"))
        if query_params["stream_sources"]:
            query_params[
                "streams_sos_columns"
            ] = _get_streams_sos_columns_from_stream_sources(
                store_ids, query_params["stream_sources"]
            )
        else:
            query_params["streams_sos_columns"] = get_streams_sos_columns(store_ids)

    if query_type.startswith("TRACK_STREAMS"):
        query = SoundRecordingStreamsTimeSeries({**query_params, **permissions})
    else:
        query = SoundRecordingDownloadsTimeSeries({**query_params, **permissions})

    time_series = query.execute()

    ts = []
    for data_point in time_series:
        data_point = format_row(data_point)
        if "id" in data_point:
            data_point["id"] = str(data_point["id"])
        if (
            data_point.get("streams_with_skips")
            or data_point.get("streams_with_skips") == 0
        ):
            skip_rate = calc_date_skip_rate(data_point)
            del data_point["streams_with_skips"]
            data_point["skip_rate"] = skip_rate
        ts.append(data_point)

    if query_type in ("TRACK_STREAMS_BY_SOS", "TRACK_STREAMS_BY_SOS_V2"):
        ids = query_params.get("ids")

        version = 2 if query_type.endswith("_V2") else 1
        if version == 2:
            ids = query_params["streams_sos_columns"]

        timeseries = []
        for item in ts:
            timeseries.extend(breakdown_by_sos(item, sources=ids, version=version))
        ts = timeseries

    return ts


def _add_item_id_field(
    items: List[Dict],
    query_type: str,
) -> List[Dict]:
    """Adds relevant ID field used by GraphQL resolver"""
    updated_items = []
    for item in items:
        if query_type == "STORE":
            id_field = {"store": {"store_id": item.get("id")}}
        elif query_type == "COUNTRY":
            id_field = {"country_code": item.get("id")}
        elif query_type in ["SOS", "SOS_V2"]:
            id_field = {"stream_source": item.get("id")}
        elif query_type == "PRODUCT":
            id_field = {"product": {"product_id": item.get("id")}}
        else:
            raise ValueError(f"Invalid query type `{query_type} for Summary`")
        updated_items.append({**item, **id_field})
    return updated_items


@cache_in_redis(ttl=cache.ONE_DAY)
def get_sound_recording_summary(
    query_params: Mapping[str, Any],
    permissions: Mapping[str, Any],
) -> Dict[str, Union[List[Dict], int]]:
    query_type = query_params.get("type", TOTAL_SUMMARY_TYPE)

    aggregation_field = AGGREGATION_FIELDS_SUMMARY.get(query_type)
    if aggregation_field:
        query_params["aggregation_field"] = aggregation_field
    countries = query_params.get("countries", [])
    tables = get_sound_recording_tables_for_summary(query_type, countries)
    query_params.update(tables)

    if query_params["stream_sources"]:
        query_params[
            "streams_sum_expression"
        ] = _get_sum_expression_of_streams_sos_columns(
            check_store_ids_for_sos(query_params.get("store_ids")),
            query_params["stream_sources"],
        )

    if query_type == "SOS_V2":
        store_ids = check_store_ids_for_sos(query_params.get("store_ids"))
        if query_params["stream_sources"]:
            query_params[
                "streams_sos_columns"
            ] = _get_streams_sos_columns_from_stream_sources(
                store_ids, query_params["stream_sources"]
            )
        else:
            query_params["streams_sos_columns"] = get_streams_sos_columns(store_ids)

    query = SoundRecordingSummary({**query_params, **permissions})

    summary = query.execute()

    items = []
    total_count = 0
    for data_point in summary:
        data_point = format_row(data_point)
        total_count = data_point.get("total_count")
        del data_point["total_count"]
        if (
            data_point.get("streams_with_skips")
            or data_point.get("streams_with_skips") == 0
        ):
            skip_rate = calc_date_skip_rate(data_point)
            del data_point["streams_with_skips"]
            data_point["skip_rate"] = skip_rate
        items.append(data_point)

    if items and query_type in ("SOS", "SOS_V2"):
        version = 2 if query_type == "SOS_V2" else 1
        items = breakdown_by_sos(
            items[0],
            key="streams",
            timeseries=False,
            version=version,
            sources=query_params["streams_sos_columns"] if version == 2 else None,
        )
        total_count = len(
            items
        )  # for SOS and SOS_V2 total count is the number of sources
    if items and query_type != "TOTAL":
        items = _add_item_id_field(items, query_type)

    return {"items": items, "total_count": total_count}