package io.delphiplatform.api.v3.bigtable;

import com.google.common.collect.Streams;

import org.apache.commons.lang3.tuple.Pair;
import org.springframework.stereotype.Service;

import java.time.LocalDate;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import javax.annotation.Nullable;

import io.delphiplatform.api.util.CollectionUtils;
import io.delphiplatform.api.util.FunctionUtils;
import io.delphiplatform.api.util.model.ModelUtils;
import io.delphiplatform.api.v3.constant.CountryCodeConstant;
import io.delphiplatform.api.v3.constant.DspConstants;
import io.delphiplatform.api.v3.model.AggBy;
import io.delphiplatform.api.v3.model.GroupByField;
import io.delphiplatform.api.v3.model.IncludeTracklistPlaylist;
import io.delphiplatform.api.v3.model.PlaylistStreamingForPeriodsInfo;
import io.delphiplatform.api.v3.model.PlaylistStreamingInfo;
import io.delphiplatform.api.v3.model.PublicPlaylistableModel;
import io.delphiplatform.api.v3.model.StreamModel;
import io.delphiplatform.api.v3.model.StreamingInfoAwareModel;
import io.delphiplatform.api.v3.model.SubsetParam;
import io.delphiplatform.api.v3.model.TrackAwareModel;
import io.delphiplatform.api.v3.rdb.entity.ChartmetricToAmazonPlaylistIdMappingProjection;
import io.delphiplatform.api.v3.rdb.service.DataCompletenessService;
import io.delphiplatform.api.v3.rdb.service.playlist_public.PlaylistPublicService;
import io.delphiplatform.api.v3.view.util.Params;

import static io.delphiplatform.api.v3.constant.CountryCodeConstant.LOCAL;

@Service
public class BigtableTrackPositionStreamsDataEnrichmentPublicService {

    private static final int TRACK_POS_STREAMS_FOR_PERIOD_BT_REQUEST_OFFSET_DAYS = 14;

    private static final int TRACK_POS_STREAMS_FOR_PERIOD_7_DAYS_FILTERING_OFFSET = 6;
    private static final int TRACK_POS_STREAMS_FOR_PERIOD_14_DAYS_FILTERING_OFFSET = 13;

    public static final String BT_STREAMS_QUERY_EXCEPTION_MSG = "Cannot get future result for BigTable streams query.";

    private final PlaylistPublicService playlistPublicService;
    private final BigtableStreamsService bigtableStreamsService;
    private final DataCompletenessService dataCompletenessService;

    public BigtableTrackPositionStreamsDataEnrichmentPublicService(
        PlaylistPublicService playlistPublicService,
        BigtableStreamsService bigtableStreamsService,
        DataCompletenessService dataCompletenessService
    ) {
        this.playlistPublicService = playlistPublicService;
        this.bigtableStreamsService = bigtableStreamsService;
        this.dataCompletenessService = dataCompletenessService;
    }

    public <T extends PublicPlaylistableModel & TrackAwareModel & StreamingInfoAwareModel> Stream<T> populatePlaylistStreamsForXDaysPeriods(
        Stream<T> models, Params params
    ) {
        return populatePlaylistStreamsForXDaysPeriods(models, params.getEndDate(), params.getStreamsCountryCode(), params.getDsp(),
            params.getInclude(), true
        );
    }

    public <T extends PublicPlaylistableModel & TrackAwareModel & StreamingInfoAwareModel> Stream<T> populatePlaylistStreamsForXDaysPeriods(
        Stream<T> models, LocalDate targetEndDate, Set<String> streamsCountryCode, List<String> requestDsps, List<String> include, boolean includeIsrcStreams
    ) {
        if (!CollectionUtils.contains(include, IncludeTracklistPlaylist.STREAMS_FOR_PERIOD.getValue())) {
            return models;
        }

        List<T> positions = models.collect(Collectors.toList());

        if (positions.isEmpty()) {
            return Stream.empty();
        }

        /* Prepare playlist ids to query BT */

        Set<String> playlistIdsNonAmazon = positions.stream()
            .filter(position -> !DspConstants.AMAZON.equals(position.getDsp()))
            .map(T::getPlaylistId)
            .collect(Collectors.toSet());

        Set<String> playlistIdsAmazon = new HashSet<>();

        //mapping between Chartmetric Amazon playlist id with country and Amazon playlist id
        Map<String, String> cmAmazonPlaylistIdToAmazonPlaylistId = new HashMap<>();

        Set<String> amazonCmPlaylistIds = positions.stream()
            .filter(position -> DspConstants.AMAZON.equals(position.getDsp()))
            .filter(position -> CollectionUtils.isNotEmpty(position.getCountryCode()))
            .map(T::getPlaylistId)
            .collect(Collectors.toSet());

        if (!amazonCmPlaylistIds.isEmpty()) {
            cmAmazonPlaylistIdToAmazonPlaylistId.putAll(
                playlistPublicService.getChartmetricToAmazonPlaylistIdMapping(amazonCmPlaylistIds).stream()
                    .collect(Collectors.toMap(
                        ChartmetricToAmazonPlaylistIdMappingProjection::getCmAmazonPlaylistIdWithCountry,
                        projection -> ModelUtils.getPlaylistIdWithDspPrefix(DspConstants.AMAZON,
                            projection.getAmazonPlaylistId())
                    ))
            );

            playlistIdsAmazon.addAll(
                cmAmazonPlaylistIdToAmazonPlaylistId.values()
            );
        }

        List<String> allPlaylistIdsToQueryBTForStreams = Stream.concat(playlistIdsNonAmazon.stream(),
                playlistIdsAmazon.stream())
            .collect(Collectors.toList());

        if (allPlaylistIdsToQueryBTForStreams.isEmpty()) {
            return positions.stream();
        }

        /* Prepare other params to query BT */

        Set<String> isrcs = includeIsrcStreams ?
            positions.stream()
                .map(T::getIsrc)
                .filter(Objects::nonNull)
                .collect(Collectors.toSet())
            : Collections.emptySet();

        List<String> dsps = CollectionUtils.isEmpty(requestDsps)
            ? List.of(DspConstants.APPLE, DspConstants.SPOTIFY, DspConstants.AMAZON)
            : requestDsps;

        //find streams for each playlist
        //if 'local' stream county code is requested - for each playlist find streams of its 'local' country (i.e. country of playlist)
        //(we just add all countries to filter params each time but this won't be slower as all countries are always fetched in protobuf anyway)

        Set<String> streamCountryCodesToRequest = streamsCountryCode.stream()
            .filter(countryCode -> !LOCAL.equals(countryCode))
            .collect(Collectors.toSet());

        if (streamsCountryCode.contains(LOCAL)) {
            streamCountryCodesToRequest.addAll(
                positions.stream().map(this::resolvePositionLocalCountryCodeForStreamsLookup)
                    .collect(Collectors.toSet())
            );
        }

        Map<String, LocalDate> dspDataLastUpdatedAt = dataCompletenessService.findLastCompletedUpdateDatePerDsp(dsps, targetEndDate);
        Set<String> dspIdsWithDataAvailable = dspDataLastUpdatedAt.keySet();

        List<Params> streamsGroupedByPlaylistAndIsrcParamsForEachDsp = new ArrayList<>();
        List<Params> streamsGroupedByPlaylistParamsForEachDsp = new ArrayList<>();

        //create separate params for each DSP because dates range may be different
        dspIdsWithDataAvailable.forEach(dsp -> {
            LocalDate dspEndDate = dspDataLastUpdatedAt.get(dsp);
            LocalDate dspStartDate = dspEndDate.minusDays(TRACK_POS_STREAMS_FOR_PERIOD_BT_REQUEST_OFFSET_DAYS);

            Set<String> dspPlaylistIdsToQuery = ModelUtils.filterPlaylistIdsByDsp(allPlaylistIdsToQueryBTForStreams, dsp);

            if (dspPlaylistIdsToQuery.isEmpty()) {
                return;
            }

            //request stream models distinct by date, countryCode, playlistId, isrc
            Params streamsGroupedByPlaylistAndIsrcParams = Params.builder()
                .startDate(dspStartDate)
                .endDate(dspEndDate)
                .isrc(List.copyOf(isrcs))
                .playlistId(List.copyOf(dspPlaylistIdsToQuery))
                .countryCode(streamCountryCodesToRequest)
                .dsp(Collections.singletonList(dsp))
                .aggBy(AggBy.ISRC)
                .subset(SubsetParam.PLAYLISTS)
                .groupByFields(Set.of(GroupByField.DATE, GroupByField.COUNTRY))
                .limit(Integer.MAX_VALUE)
                .build();

            //request stream models distinct by date, countryCode, playlistId
            Params streamsGroupedByPlaylistParams = Params.builder()
                .startDate(dspStartDate)
                .endDate(dspEndDate)
                .playlistId(List.copyOf(dspPlaylistIdsToQuery))
                .countryCode(streamCountryCodesToRequest)
                .dsp(Collections.singletonList(dsp))
                .subset(SubsetParam.PLAYLISTS)
                .groupByFields(Set.of(GroupByField.DATE, GroupByField.COUNTRY))
                .limit(Integer.MAX_VALUE)
                .build();

            streamsGroupedByPlaylistAndIsrcParamsForEachDsp.add(streamsGroupedByPlaylistAndIsrcParams);
            streamsGroupedByPlaylistParamsForEachDsp.add(streamsGroupedByPlaylistParams);
        });

        /* Perform BT query */

        Pair<List<StreamModel>, List<StreamModel>> streamsGroupedByPlaylistAndIsrcSlashStreamsGroupedByPlaylist = queryBTForStreams(
            streamsGroupedByPlaylistAndIsrcParamsForEachDsp,
            streamsGroupedByPlaylistParamsForEachDsp
        );

        List<StreamModel> streamsGroupedByPlaylistAndIsrc = streamsGroupedByPlaylistAndIsrcSlashStreamsGroupedByPlaylist.getLeft();
        List<StreamModel> streamsGroupedByPlaylist = streamsGroupedByPlaylistAndIsrcSlashStreamsGroupedByPlaylist.getRight();

        //if no streaming data found - return
        if (CollectionUtils.isEmpty(streamsGroupedByPlaylistAndIsrc)
            && CollectionUtils.isEmpty(streamsGroupedByPlaylist)) {

            Map<String, PlaylistStreamingForPeriodsInfo> emptyStreamingInfo = new HashMap<>();

            for (String requestedCountryCode : streamsCountryCode) {
                emptyStreamingInfo.put(requestedCountryCode, null);
            }

            return positions.stream()
                .peek(position -> position.setStreamingInfo(emptyStreamingInfo));
        }

        /* Prepare retrieved stream count maps for different periods */

        //maps format: {spotify_123-US0000000001: {us: 42}}
        Map<String, Map<String, Long>> streamsFor1DayGroupedByPlaylistAndIsrc = new HashMap<>();
        Map<String, Map<String, Long>> streamsFor7DaysGroupedByPlaylistAndIsrc = new HashMap<>();
        Map<String, Map<String, Long>> streamsFor14DaysGroupedByPlaylistAndIsrc = new HashMap<>();
        //maps format: {spotify_123: {us: 42}}
        Map<String, Map<String, Long>> streamsFor1DayGroupedByPlaylist = new HashMap<>();
        Map<String, Map<String, Long>> streamsFor7DaysGroupedByPlaylist = new HashMap<>();
        Map<String, Map<String, Long>> streamsFor14DaysGroupedByPlaylist = new HashMap<>();

        //we calculate streams counts for different periods separately for each DSP since period dates may be different
        //but after that we just add counts to common maps because playlist id (that is a part of map key) is DSP-specific anyway
        dspIdsWithDataAvailable.forEach(dsp -> {
            LocalDate dspStartDateForLast1Day = dspDataLastUpdatedAt.get(dsp);
            LocalDate dspStartDateForLast7Days = dspStartDateForLast1Day
                .minusDays(TRACK_POS_STREAMS_FOR_PERIOD_7_DAYS_FILTERING_OFFSET);
            LocalDate dspStartDateForLast14Days = dspStartDateForLast1Day
                .minusDays(TRACK_POS_STREAMS_FOR_PERIOD_14_DAYS_FILTERING_OFFSET);

            List<StreamModel> dspStreamsGroupedByPlaylistAndIsrc = streamsGroupedByPlaylistAndIsrc.stream()
                .filter(sm -> dsp.equals(sm.getDsp()))
                .collect(Collectors.toList());

            streamsFor1DayGroupedByPlaylistAndIsrc.putAll(
                getStreamsSumGroupedByKeyAndCountries(dspStreamsGroupedByPlaylistAndIsrc, dspStartDateForLast1Day,
                    ModelUtils::getPlaylistIdWithIsrc)
            );

            streamsFor7DaysGroupedByPlaylistAndIsrc.putAll(
                getStreamsSumGroupedByKeyAndCountries(dspStreamsGroupedByPlaylistAndIsrc, dspStartDateForLast7Days,
                    ModelUtils::getPlaylistIdWithIsrc)
            );

            streamsFor14DaysGroupedByPlaylistAndIsrc.putAll(
                getStreamsSumGroupedByKeyAndCountries(dspStreamsGroupedByPlaylistAndIsrc, dspStartDateForLast14Days,
                    ModelUtils::getPlaylistIdWithIsrc)
            );

            List<StreamModel> dspStreamsGroupedByPlaylist = streamsGroupedByPlaylist.stream()
                .filter(sm -> dsp.equals(sm.getDsp()))
                .collect(Collectors.toList());

            streamsFor1DayGroupedByPlaylist.putAll(
                getStreamsSumGroupedByKeyAndCountries(dspStreamsGroupedByPlaylist, dspStartDateForLast1Day, StreamModel::getPlaylistId)
            );

            streamsFor7DaysGroupedByPlaylist.putAll(
                getStreamsSumGroupedByKeyAndCountries(dspStreamsGroupedByPlaylist, dspStartDateForLast7Days, StreamModel::getPlaylistId)
            );

            streamsFor14DaysGroupedByPlaylist.putAll(
                getStreamsSumGroupedByKeyAndCountries(dspStreamsGroupedByPlaylist, dspStartDateForLast14Days, StreamModel::getPlaylistId)
            );
        });

        return populateTrackPositionsWithStreamingInfo(
            positions,
            streamsFor1DayGroupedByPlaylistAndIsrc,
            streamsFor7DaysGroupedByPlaylistAndIsrc,
            streamsFor14DaysGroupedByPlaylistAndIsrc,
            streamsFor1DayGroupedByPlaylist,
            streamsFor7DaysGroupedByPlaylist,
            streamsFor14DaysGroupedByPlaylist,
            cmAmazonPlaylistIdToAmazonPlaylistId,
            streamsCountryCode
        );
    }

    /**
     * Filters list of stream models to include only latest reports (since `sinceDate`), then groups streams by key (e.g. by playlistId or by
     * playlistId+isrc, depending on `getStreamModelKeyFunc`) then further groups streams by country and sums stream counts from different dates
     *
     * @return map of positions (playlistId+isrc) with sum of streams per country e.g. {'spotifyPl1': {'us': 25, 'ca': 105}}
     * {'spotifyPl1-US00000001': {'us': 25, 'ca': 105}}
     */
    private Map<String, Map<String, Long>> getStreamsSumGroupedByKeyAndCountries(
        List<StreamModel> streamModels,
        LocalDate sinceDate,
        Function<StreamModel, String> getStreamModelKeyFunc
    ) {
        return streamModels
            .stream()
            .filter(streamModel -> streamModel.getDate().equals(sinceDate) || streamModel.getDate().isAfter(sinceDate))
            .collect(Collectors.groupingBy(getStreamModelKeyFunc))
            .entrySet()
            .stream()
            .collect(Collectors.toMap(
                Entry::getKey,
                entry -> entry.getValue().stream().collect(Collectors.toMap(
                    StreamModel::getCountryCode,
                    StreamModel::getStreams,
                    Long::sum
                ))
            ));
    }

    private String resolvePositionLocalCountryCodeForStreamsLookup(PublicPlaylistableModel position) {
        String playlistLocalCountryCode = position.getCountryCode();

        if (CountryCodeConstant.GLOBAL.equals(playlistLocalCountryCode)) {
            return CountryCodeConstant.WORLDWIDE;
        } else {
            return playlistLocalCountryCode;
        }
    }

    private <T extends PublicPlaylistableModel & TrackAwareModel & StreamingInfoAwareModel> Stream<T> populateTrackPositionsWithStreamingInfo(
        List<T> positions,
        Map<String, Map<String, Long>> streamsFor1DayGroupedByPlaylistAndIsrc,
        Map<String, Map<String, Long>> streamsFor7DaysGroupedByPlaylistAndIsrc,
        Map<String, Map<String, Long>> streamsFor14DaysGroupedByPlaylistAndIsrc,
        Map<String, Map<String, Long>> streamsFor1DayGroupedByPlaylist,
        Map<String, Map<String, Long>> streamsFor7DaysGroupedByPlaylist,
        Map<String, Map<String, Long>> streamsFor14DaysGroupedByPlaylist,
        Map<String, String> cmAmazonPlaylistIdToAmazonPlaylistId,
        Set<String> requestedStreamsCountryCodes
    ) {
        return positions.stream()
            .peek(position -> {
                String playlistIdToQueryBTForStreams;

                if (DspConstants.AMAZON.equals(position.getDsp())) {
                    //to query BT amazon_music table - resolve 'amazon playlist id' by 'chartmetric amazon playlist id' which is used in position class
                    playlistIdToQueryBTForStreams =
                        cmAmazonPlaylistIdToAmazonPlaylistId.get(position.getPlaylistId());
                } else {
                    playlistIdToQueryBTForStreams = position.getPlaylistId();
                }

                if (playlistIdToQueryBTForStreams != null) {
                    String playlistIdAndIsrcKey = ModelUtils.getPlaylistIdWithIsrc(playlistIdToQueryBTForStreams,
                        position.getIsrc());

                    Map<String, Long> isrcInPlaylistStreamsByCountryFor1Day =
                        streamsFor1DayGroupedByPlaylistAndIsrc.getOrDefault(playlistIdAndIsrcKey, Collections.emptyMap());
                    Map<String, Long> playlistStreamsByCountryFor1Day =
                        streamsFor1DayGroupedByPlaylist.getOrDefault(playlistIdToQueryBTForStreams, Collections.emptyMap());

                    Map<String, Long> isrcInPlaylistStreamsByCountryFor7Days =
                        streamsFor7DaysGroupedByPlaylistAndIsrc.getOrDefault(playlistIdAndIsrcKey, Collections.emptyMap());
                    Map<String, Long> playlistStreamsByCountryFor7Days =
                        streamsFor7DaysGroupedByPlaylist.getOrDefault(playlistIdToQueryBTForStreams, Collections.emptyMap());

                    Map<String, Long> isrcInPlaylistStreamsByCountryFor14Days =
                        streamsFor14DaysGroupedByPlaylistAndIsrc.getOrDefault(playlistIdAndIsrcKey, Collections.emptyMap());
                    Map<String, Long> playlistStreamsByCountryFor14Days =
                        streamsFor14DaysGroupedByPlaylist.getOrDefault(playlistIdToQueryBTForStreams, Collections.emptyMap());

                    Map<String, PlaylistStreamingForPeriodsInfo> streamingInfo = new HashMap<>();

                    for (String requestedCountryCode : requestedStreamsCountryCodes) {
                        String countryCode;

                        if (!LOCAL.equals(requestedCountryCode)) {
                            countryCode = requestedCountryCode;
                        } else {
                            countryCode = resolvePositionLocalCountryCodeForStreamsLookup(position);
                        }

                        @Nullable
                        PlaylistStreamingForPeriodsInfo streamingForPeriodsInfo = null;

                        Long isrcInPlaylistStreams1Day = isrcInPlaylistStreamsByCountryFor1Day.get(countryCode);
                        Long playlistStreams1Day = playlistStreamsByCountryFor1Day.get(countryCode);

                        Long isrcInPlaylistStreams7Days = isrcInPlaylistStreamsByCountryFor7Days.get(countryCode);
                        Long playlistStreams7Days = playlistStreamsByCountryFor7Days.get(countryCode);

                        Long isrcInPlaylistStreams14Days = isrcInPlaylistStreamsByCountryFor14Days.get(countryCode);
                        Long playlistStreams14Days = playlistStreamsByCountryFor14Days.get(countryCode);

                        if (isrcInPlaylistStreams1Day != null || playlistStreams1Day != null) {
                            if (streamingForPeriodsInfo == null) {
                                streamingForPeriodsInfo = new PlaylistStreamingForPeriodsInfo();
                            }

                            streamingForPeriodsInfo.setPeriodFor1Day(
                                PlaylistStreamingInfo.builder()
                                    .isrcInPlaylist(isrcInPlaylistStreams1Day)
                                    .playlist(playlistStreams1Day)
                                    .build()
                            );
                        }

                        if (isrcInPlaylistStreams7Days != null || playlistStreams7Days != null) {
                            if (streamingForPeriodsInfo == null) {
                                streamingForPeriodsInfo = new PlaylistStreamingForPeriodsInfo();
                            }

                            streamingForPeriodsInfo.setPeriodFor7Days(
                                PlaylistStreamingInfo.builder()
                                    .isrcInPlaylist(isrcInPlaylistStreams7Days)
                                    .playlist(playlistStreams7Days)
                                    .build()
                            );
                        }

                        if (isrcInPlaylistStreams14Days != null || playlistStreams14Days != null) {
                            if (streamingForPeriodsInfo == null) {
                                streamingForPeriodsInfo = new PlaylistStreamingForPeriodsInfo();
                            }

                            streamingForPeriodsInfo.setPeriodFor14Days(
                                PlaylistStreamingInfo.builder()
                                    .isrcInPlaylist(isrcInPlaylistStreams14Days)
                                    .playlist(playlistStreams14Days)
                                    .build()
                            );
                        }

                        //requestedCountryCode used here so if this is 'local' - we set this exact name w/o translation to real countryCode
                        streamingInfo.put(requestedCountryCode, streamingForPeriodsInfo);
                    }

                    position.setStreamingInfo(streamingInfo);
                }
            });
    }

    private Pair<List<StreamModel>, List<StreamModel>> queryBTForStreams(
        List<Params> streamsGroupedByPlaylistAndIsrcParamsForEachDsp,
        List<Params> streamsGroupedByPlaylistParamsForEachDsp
    ) {
        if (streamsGroupedByPlaylistAndIsrcParamsForEachDsp.isEmpty()
            && streamsGroupedByPlaylistParamsForEachDsp.isEmpty()) {
            return Pair.of(Collections.emptyList(), Collections.emptyList());
        }

        List<CompletableFuture<List<StreamModel>>> streamsGroupedByPlaylistAndIsrcFutures =
            streamsGroupedByPlaylistAndIsrcParamsForEachDsp.stream()
                .map(prms -> FunctionUtils.applyWrappedInRuntimeEx(bigtableStreamsService::getRows, prms, BT_STREAMS_QUERY_EXCEPTION_MSG))
                .collect(Collectors.toList());

        List<CompletableFuture<List<StreamModel>>> streamsGroupedByPlaylistFutures =
            streamsGroupedByPlaylistParamsForEachDsp.stream()
                .map(prms -> FunctionUtils.applyWrappedInRuntimeEx(bigtableStreamsService::getRows, prms, BT_STREAMS_QUERY_EXCEPTION_MSG))
                .collect(Collectors.toList());

        CompletableFuture.allOf(
            Streams.concat(
                streamsGroupedByPlaylistAndIsrcFutures.stream(),
                streamsGroupedByPlaylistFutures.stream()
            ).toArray(CompletableFuture[]::new)
        ).join();

        List<StreamModel> streamsGroupedByPlaylistAndIsrc = streamsGroupedByPlaylistAndIsrcFutures.stream()
            .flatMap(ftr -> FunctionUtils.applyWrappedInRuntimeEx(f -> f.get().stream(), ftr, BT_STREAMS_QUERY_EXCEPTION_MSG))
            .collect(Collectors.toList());

        List<StreamModel> streamsGroupedByPlaylist = streamsGroupedByPlaylistFutures.stream()
            .flatMap(ftr -> FunctionUtils.applyWrappedInRuntimeEx(f -> f.get().stream(), ftr, BT_STREAMS_QUERY_EXCEPTION_MSG))
            .collect(Collectors.toList());

        return Pair.of(streamsGroupedByPlaylistAndIsrc, streamsGroupedByPlaylist);
    }
}