"""Logic for retrieving streams."""

import datetime
import json
from collections import defaultdict
from typing import Any, Mapping

from ddtrace import tracer
from oto import response as oto_response

from analytics.config import STORES
from analytics.constants import cache
from analytics.constants.parameters import ALL_TIME
from analytics.logic import data_availability
from analytics.logic.streams_helpers import rows_to_dicts
from analytics.queries.format import format_row
from analytics.queries.sound_recording_streams import StreamsAll
from analytics.queries.streams import (
    GlobalParticipantAggregatedStreamsByCountryOrStore,
    GlobalParticipantAggregatedStreamsBySOS,
    ProductAggregatedStreamsByCountryOrStore,
    ProductAggregatedStreamsBySOS,
    SoundRecordingAggregatedStreamsByCountryOrStore,
    SoundRecordingAggregatedStreamsBySOS,
    SoundRecordingAggregatedStreamsBySOSV2,
)
from analytics.schemas.streams import StreamsAllSchema
from analytics.utils import date as date_utils
from analytics.utils import store_availability
from analytics.utils import streams as streams_utils
from analytics.utils.cache import cache_in_redis
from analytics.validation.schema import schema_dump


@cache_in_redis(ttl=cache.ONE_DAY)
@tracer.wrap(name="get_streams_all")
def get_streams_all(
    query_params: Mapping[str, Any],
    permissions: Mapping[str, Any],
):
    """Return aggregate streams timeseries for ISRC.

    Args:
        query_params: Dict with isrc, distributors, country_ids, store_ids,
            start_date, end_date.
        permissions: Dict with permission_* keys.

    Returns:
        oto.response.Response with aggregate streams.

    """
    isrc = query_params["isrc"]
    response_body = {"isrc": isrc, "items": []}

    start_date = query_params.get("start_date")
    end_date = query_params.get("end_date")

    if not (start_date and end_date):
        start_date, end_date = data_availability.get_date_range(
            data_availability.HIGHWATERMARK_DATE, days=28
        )

    all_time = start_date == ALL_TIME

    store_ids = store_availability.get_query_store_ids(
        query_params.get("store_ids", [])
    )
    if not store_ids:
        schema = StreamsAllSchema()
        return oto_response.Response(schema_dump(schema, response_body))

    end_date_str = (
        end_date.strftime("%Y-%m-%d")
        if isinstance(end_date, datetime.date)
        else end_date
    )
    query_input = {
        **permissions,
        "isrc": isrc,
        "store_ids": store_ids,
        "distributors": query_params["distributors"],
        "end_date": end_date_str,
        "country_ids": query_params.get("country_ids", []),
        "all_time": all_time,
        "transfer_product_ownership_enabled": query_params.get(
            "transfer_product_ownership_enabled", False
        ),
    }
    if not all_time:
        start_date_str = (
            start_date.strftime("%Y-%m-%d")
            if isinstance(start_date, datetime.date)
            else start_date
        )
        query_input["start_date"] = start_date_str

    streams = rows_to_dicts(StreamsAll(query_input).execute())
    streams_all = streams_utils.get_streams_all(streams, start_date, end_date)

    if streams_all:
        response_body["items"] = streams_all
    schema = StreamsAllSchema()
    return oto_response.Response(schema_dump(schema, response_body))


@cache_in_redis(ttl=cache.ONE_DAY)
def get_aggregated_streams(
    query_params: Mapping[str, Any], permissions: Mapping[str, Any]
):
    QUERIES = {
        "isrc": {
            "SOS": SoundRecordingAggregatedStreamsBySOS,
            "SOS_V2": SoundRecordingAggregatedStreamsBySOSV2,
            "STORE": SoundRecordingAggregatedStreamsByCountryOrStore,
            "COUNTRY": SoundRecordingAggregatedStreamsByCountryOrStore,
        },
        "product_id": {
            "SOS": ProductAggregatedStreamsBySOS,
            "STORE": ProductAggregatedStreamsByCountryOrStore,
            "COUNTRY": ProductAggregatedStreamsByCountryOrStore,
        },
        "global_participant_id": {
            "SOS": GlobalParticipantAggregatedStreamsBySOS,
            "STORE": GlobalParticipantAggregatedStreamsByCountryOrStore,
            "COUNTRY": GlobalParticipantAggregatedStreamsByCountryOrStore,
        },
    }

    if query_params.get("isrc"):
        type = "isrc"
    elif query_params.get("product_id"):
        type = "product_id"
    elif query_params.get("global_participant_id"):
        type = "global_participant_id"
    else:
        raise ValueError(
            "get_aggregated_streams: Must provide either isrc, product_id "
            "or global_participant_id"
        )

    if query_params["dimension"] == "SOS_V2":
        query_params["streams_sos_columns"] = []
        if query_params.get("store_ids", []):
            query_params["store_ids"] = streams_utils.check_store_ids_for_sos(
                query_params["store_ids"]
            )
            query_params["streams_sos_columns"] = streams_utils.get_streams_sos_columns(
                query_params["store_ids"]
            )
    else:
        if query_params.get("store_ids", []):
            query_params["store_ids"] = store_availability.get_query_store_ids(
                query_params.get("store_ids", []),
            )

    query_params[
        "is_feed_data_available"
    ] = store_availability.is_apple_spotify_data_in_sync()
    query_params["max_available_date"] = str(data_availability.get_max_available_date())

    query_params.setdefault("transfer_product_ownership_enabled", False)

    query = QUERIES[type][query_params["dimension"]]({**query_params, **permissions})

    aggregate = [format_row(data_point) for data_point in query.execute()]
    if aggregate:
        parsed = json.loads(aggregate[0]["result"].lower())
        master_calendar = set(
            date_utils.generate_date_range_starting_from_date(
                start_date=query_params["max_available_date"],
                days_back=query_params["days_back"],
            )
        )

        # Group by id and sort by date + fill in missing dates using master_calendar
        topn_timeseries = defaultdict(list)
        dates_with_data = defaultdict(set)
        for item in parsed["topn_timeseries"]:
            if query_params["dimension"] == "STORE":
                key = STORES[item["id"]]
            else:
                key = item["id"]
            topn_timeseries[key].append({"date": item["date"], "value": item["value"]})
            dates_with_data[key].add(item["date"])

        for key in topn_timeseries:
            dates_with_no_data_to_add = master_calendar - dates_with_data[key]
            topn_timeseries[key] = streams_utils.clean_up_dates_from_lag(
                topn_timeseries[key],
                key,
                dates_with_no_data_to_add,
                query_params,
                parsed["missing_store_dates"],
            )
            topn_timeseries[key].sort(key=lambda x: x["date"])
        parsed["topn_timeseries"] = topn_timeseries

        all_other_timeseries = parsed["all_other_timeseries"]
        dates_with_data = set([item["date"] for item in all_other_timeseries])
        dates_with_no_data_to_add = master_calendar - dates_with_data
        all_other_timeseries = streams_utils.clean_up_dates_from_lag(
            all_other_timeseries,
            None,
            dates_with_no_data_to_add,
            query_params,
            parsed["missing_store_dates"],
        )
        parsed["all_other_timeseries"] = sorted(
            all_other_timeseries, key=lambda x: x["date"]
        )

        for item in parsed["topn_rollup"]:
            if query_params["dimension"] == "STORE":
                item["id"] = STORES[item["id"]]

        # Delete missing store dates key since only used for computation
        del parsed["missing_store_dates"]
    else:
        parsed = {}
    return parsed