package io.delphiplatform.api.v3.bigtable;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;

import org.apache.commons.lang3.tuple.Pair;
import org.apache.commons.text.StringSubstitutor;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.time.LocalDate;
import java.util.Collections;
import java.util.List;
import java.util.Set;

import javax.annotation.concurrent.NotThreadSafe;
import javax.transaction.NotSupportedException;

import io.delphiplatform.api.util.DateUtils;
import io.delphiplatform.api.v3.bigtable.config.RequestRowKeyHandle;
import io.delphiplatform.api.v3.bigtable.config.RowKeyAggPeriod;
import io.delphiplatform.api.v3.bigtable.config.RowKeyGroup;
import io.delphiplatform.api.v3.bigtable.config.RowKeyRequestType;
import io.delphiplatform.api.v3.model.AggBy;
import io.delphiplatform.api.v3.model.tiktok.TikTokContentType;
import io.delphiplatform.api.v3.model.video.YoutubeContentType;
import io.delphiplatform.api.v3.view.util.Params;

import static io.delphiplatform.api.v3.bigtable.config.RowKeyPrefix.ARTIST;
import static io.delphiplatform.api.v3.bigtable.config.RowKeyPrefix.ISRC_CONTENT_DATE;
import static io.delphiplatform.api.v3.bigtable.config.RowKeyPrefix.ISRC_DATE;
import static io.delphiplatform.api.v3.bigtable.config.RowKeyPrefix.ISRC_DATE_PLAYLIST;
import static io.delphiplatform.api.v3.bigtable.config.RowKeyPrefix.ISRC_TIKTOK_CONTENT_TYPE_DATE;
import static io.delphiplatform.api.v3.bigtable.config.RowKeyPrefix.PLAYLIST_DATE_ISRC;
import static io.delphiplatform.api.v3.bigtable.config.RowKeyPrefix.PLAYLIST_ISRC_DATE;
import static io.delphiplatform.api.v3.bigtable.config.RowKeyPrefix.VIDEO_DATE;
import static io.delphiplatform.api.v3.model.video.YoutubeContentType.PARTNER_UPLOADED;
import static io.delphiplatform.api.v3.model.video.YoutubeContentType.PREMIUM_UGC;
import static org.junit.jupiter.api.Assertions.assertEquals;

@NotThreadSafe
class RowKeyGeneratorTest {

    // These should match what is in the test data seeds
    static final ImmutableList<String> TEST_ARTIST_IDS = ImmutableList
        .of("GRAS_7513");
    static final ImmutableList<String> TEST_VIDEO_IDS = ImmutableList
        .of("youtube_-qap5aO4i9A");
    static final ImmutableList<String> TEST_CHART_IDS = ImmutableList
        .of("apple_top100_daily_us", "spotify_top200_daily_us");
    static final ImmutableList<String> TEST_DSPS = ImmutableList
        .of("apple", "spotify");
    static final ImmutableList<String> TEST_ISRCS = ImmutableList
        .of("USSM11914699", "USSM10010774");
    static final ImmutableList<String> TEST_PLAYLIST_IDS = ImmutableList
        .of("apple_pl.0001620d82ce4ebb9570ceb57123af1d", "spotify_02C8uKtc6CoBpD0IdqOZdT");
    static final ImmutableList<String> TEST_PRODUCT_IDS = ImmutableList
        .of("GRAS_G0100042483587", "GRAS_G010000912970T");
    static final ImmutableList<String> TEST_TRACK_IDS = ImmutableList
        .of("GRAS_A10328E00092855030", "GRAS_A10328E0001321390L");
    static final ImmutableList<YoutubeContentType> TEST_VIDEO_CONTENT_TYPES = ImmutableList
        .of(PARTNER_UPLOADED, PREMIUM_UGC);
    static final ImmutableList<TikTokContentType> TEST_TIKTOK_CONTENT_TYPES = ImmutableList
        .of(TikTokContentType.PGC, TikTokContentType.UGC);

    static final LocalDate TEST_START_DATE = DateUtils.parse("2020-01-01");
    static final LocalDate TEST_END_DATE = DateUtils.parse("2020-01-07");

    // Expected values to be in the key substitution
    static final String ARTIST_ID_1 = TEST_ARTIST_IDS.get(0);
    static final String VIDEO_ID_1 = TEST_VIDEO_IDS.get(0);
    static final String CHART_ID_1 = TEST_CHART_IDS.get(0);
    static final String CHART_ID_2 = TEST_CHART_IDS.get(1);
    static final String DSP_1 = TEST_DSPS.get(0);
    static final String DSP_2 = TEST_DSPS.get(1);
    static final String END_DATE = "2020-01-08";
    static final String PLAYLIST_ID_1 = TEST_PLAYLIST_IDS.get(0);
    static final String PLAYLIST_ID_2 = TEST_PLAYLIST_IDS.get(1);
    static final String REV_ISRC_1 = "99641911MSSU";
    static final String REV_ISRC_2 = "47701001MSSU";
    static final String START_DATE = "2020-01-01";
    static final String CONTENT_TYPE_1 = "PARTNER_UPLOADED";
    static final String CONTENT_TYPE_2 = "PREMIUM_UGC";

    static final ImmutableMap<String, String> valuesMap = ImmutableMap.<String, String>builder()
        .put("START_DATE", START_DATE)
        .put("END_DATE", END_DATE)
        .put("REV_ISRC_1", REV_ISRC_1)
        .put("REV_ISRC_2", REV_ISRC_2)
        .put("PLAYLIST_ID_1", PLAYLIST_ID_1)
        .put("PLAYLIST_ID_2", PLAYLIST_ID_2)
        .put("CHART_ID_1", CHART_ID_1)
        .put("CHART_ID_2", CHART_ID_2)
        .put("DSP_1", DSP_1)
        .put("DSP_2", DSP_2)
        .put("ARTIST_ID_1", ARTIST_ID_1)
        .put("VIDEO_ID_1", VIDEO_ID_1)
        .put("CONTENT_TYPE_1", CONTENT_TYPE_1)
        .put("CONTENT_TYPE_2", CONTENT_TYPE_2)
        .build();
    static final StringSubstitutor substitutor = new StringSubstitutor(valuesMap);

    @BeforeEach
    void setUp() {
    }

    /**
     * streams - isrcs
     */
    @Test
    public void getRangeKeySets_isrcStreamData() throws NotSupportedException {
        RowKeyGroup rowKeyGroup = RowKeyGroup.STREAM;
        RowKeyAggPeriod rowKeyAggPeriod = RowKeyAggPeriod.DAY;
        Params params = Params.builder().startDate(TEST_START_DATE).endDate(TEST_END_DATE)
            .isrc(TEST_ISRCS)
            .build();

        List<RequestRowKeyHandle> expected = List.of(
            RequestRowKeyHandle.builder()
                .rowKeyPrefix(ISRC_DATE)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("isrc_date_day~${REV_ISRC_1}~${START_DATE}"),
                    substitutor.replace("isrc_date_day~${REV_ISRC_1}~${END_DATE}")))
                .build(),

            RequestRowKeyHandle.builder()
                .rowKeyPrefix(ISRC_DATE)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("isrc_date_day~${REV_ISRC_2}~${START_DATE}"),
                    substitutor.replace("isrc_date_day~${REV_ISRC_2}~${END_DATE}")))
                .build()
        );
        List<RequestRowKeyHandle> result = new RowKeyGenerator(params, rowKeyGroup, rowKeyAggPeriod,
            RowKeyRequestType.RANGE).getRequestRowKeyHandles();
        assertEquals(expected, result);
    }

    /**
     * streams - isrcs + empty playlistId list
     */
    @Test
    public void getRangeKeySets_isrcStreamDataEmptyPlaylistId() throws NotSupportedException {
        RowKeyGroup rowKeyGroup = RowKeyGroup.STREAM;
        RowKeyAggPeriod rowKeyAggPeriod = RowKeyAggPeriod.DAY;
        Params params = Params.builder().startDate(TEST_START_DATE).endDate(TEST_END_DATE)
            .isrc(TEST_ISRCS)
            .playlistId(Collections.emptyList())
            .build();

        List<RequestRowKeyHandle> expected = List.of(
            RequestRowKeyHandle.builder()
                .rowKeyPrefix(ISRC_DATE)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("isrc_date_day~${REV_ISRC_1}~${START_DATE}"),
                    substitutor.replace("isrc_date_day~${REV_ISRC_1}~${END_DATE}")))
                .build(),

            RequestRowKeyHandle.builder()
                .rowKeyPrefix(ISRC_DATE)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("isrc_date_day~${REV_ISRC_2}~${START_DATE}"),
                    substitutor.replace("isrc_date_day~${REV_ISRC_2}~${END_DATE}")))
                .build()
        );
        List<RequestRowKeyHandle> result = new RowKeyGenerator(params, rowKeyGroup, rowKeyAggPeriod,
            RowKeyRequestType.RANGE).getRequestRowKeyHandles();
        assertEquals(expected, result);
    }

    /**
     * streams - playlist_ids and isrcs
     */
    @Test
    public void getRangeKeySets_playlistStreamData() throws NotSupportedException {
        RowKeyGroup rowKeyGroup = RowKeyGroup.STREAM;
        RowKeyAggPeriod rowKeyAggPeriod = RowKeyAggPeriod.DAY;
        Params params = Params.builder().startDate(TEST_START_DATE).endDate(TEST_END_DATE)
            .isrc(TEST_ISRCS)
            .playlistId(TEST_PLAYLIST_IDS)
            .build();

        List<RequestRowKeyHandle> expected = List.of(
            RequestRowKeyHandle.builder()
                .rowKeyPrefix(PLAYLIST_ISRC_DATE)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("playlist_isrc_date_day~${PLAYLIST_ID_1}~${REV_ISRC_1}~${START_DATE}"),
                    substitutor.replace("playlist_isrc_date_day~${PLAYLIST_ID_1}~${REV_ISRC_1}~${END_DATE}")))
                .build(),

            RequestRowKeyHandle.builder()
                .rowKeyPrefix(PLAYLIST_ISRC_DATE)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("playlist_isrc_date_day~${PLAYLIST_ID_1}~${REV_ISRC_2}~${START_DATE}"),
                    substitutor.replace("playlist_isrc_date_day~${PLAYLIST_ID_1}~${REV_ISRC_2}~${END_DATE}")))
                .build(),

            RequestRowKeyHandle.builder()
                .rowKeyPrefix(PLAYLIST_ISRC_DATE)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("playlist_isrc_date_day~${PLAYLIST_ID_2}~${REV_ISRC_1}~${START_DATE}"),
                    substitutor.replace("playlist_isrc_date_day~${PLAYLIST_ID_2}~${REV_ISRC_1}~${END_DATE}")))
                .build(),

            RequestRowKeyHandle.builder()
                .rowKeyPrefix(PLAYLIST_ISRC_DATE)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("playlist_isrc_date_day~${PLAYLIST_ID_2}~${REV_ISRC_2}~${START_DATE}"),
                    substitutor.replace("playlist_isrc_date_day~${PLAYLIST_ID_2}~${REV_ISRC_2}~${END_DATE}")))
                .build()
        );
        List<RequestRowKeyHandle> result = new RowKeyGenerator(params, rowKeyGroup, rowKeyAggPeriod,
            RowKeyRequestType.RANGE).getRequestRowKeyHandles();
        assertEquals(expected, result);
    }

    /**
     * streams - playlist_ids and agg_by isrc
     */
    @Test
    public void getRangeKeySets_playlistAndAggByIsrcStreamData() throws NotSupportedException {
        RowKeyGroup rowKeyGroup = RowKeyGroup.STREAM;
        RowKeyAggPeriod rowKeyAggPeriod = RowKeyAggPeriod.DAY;
        Params params = Params.builder().startDate(TEST_START_DATE).endDate(TEST_END_DATE)
            .playlistId(TEST_PLAYLIST_IDS)
            .aggBy(AggBy.ISRC)
            .build();

        List<RequestRowKeyHandle> expected = List.of(
            RequestRowKeyHandle.builder()
                .rowKeyPrefix(PLAYLIST_DATE_ISRC)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("playlist_date_isrc_day~${PLAYLIST_ID_1}~${START_DATE}"),
                    substitutor.replace("playlist_date_isrc_day~${PLAYLIST_ID_1}~${END_DATE}")))
                .build(),

            RequestRowKeyHandle.builder()
                .rowKeyPrefix(PLAYLIST_DATE_ISRC)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("playlist_date_isrc_day~${PLAYLIST_ID_2}~${START_DATE}"),
                    substitutor.replace("playlist_date_isrc_day~${PLAYLIST_ID_2}~${END_DATE}")))
                .build()
        );
        List<RequestRowKeyHandle> result = new RowKeyGenerator(params, rowKeyGroup, rowKeyAggPeriod,
            RowKeyRequestType.RANGE).getRequestRowKeyHandles();
        assertEquals(expected, result);
    }

    /**
     * playlist positions - isrcs
     */
    @Test
    public void getRangeKeySets_playlistPositionData() throws NotSupportedException {
        RowKeyGroup rowKeyGroup = RowKeyGroup.PLAYLIST;
        RowKeyAggPeriod rowKeyAggPeriod = RowKeyAggPeriod.DAY;
        Params params = Params.builder().startDate(TEST_START_DATE).endDate(TEST_END_DATE)
            .isrc(TEST_ISRCS)
            .build();

        List<RequestRowKeyHandle> expected = List.of(
            RequestRowKeyHandle.builder()
                .rowKeyPrefix(ISRC_DATE_PLAYLIST)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("isrc_date_playlist_day~${REV_ISRC_1}~${START_DATE}"),
                    substitutor.replace("isrc_date_playlist_day~${REV_ISRC_1}~${END_DATE}")))
                .build(),

            RequestRowKeyHandle.builder()
                .rowKeyPrefix(ISRC_DATE_PLAYLIST)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("isrc_date_playlist_day~${REV_ISRC_2}~${START_DATE}"),
                    substitutor.replace("isrc_date_playlist_day~${REV_ISRC_2}~${END_DATE}")))
                .build()
        );
        List<RequestRowKeyHandle> result = new RowKeyGenerator(params, rowKeyGroup, rowKeyAggPeriod,
            RowKeyRequestType.RANGE).getRequestRowKeyHandles();
        assertEquals(expected, result);
    }

    /**
     * chartmetric - artist_id
     */
    @Test
    public void getRangeKeySets_chartmetricData() throws NotSupportedException {
        RowKeyGroup rowKeyGroup = RowKeyGroup.CHARTMETRIC;
        RowKeyAggPeriod rowKeyAggPeriod = RowKeyAggPeriod.DAILY;
        Params params = Params.builder().startDate(TEST_START_DATE).endDate(TEST_END_DATE)
            .artistId(TEST_ARTIST_IDS.get(0))
            .build();

        List<RequestRowKeyHandle> expected = List.of(
            RequestRowKeyHandle.builder()
                .rowKeyPrefix(ARTIST)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("artist_daily~${ARTIST_ID_1}~${START_DATE}"),
                    substitutor.replace("artist_daily~${ARTIST_ID_1}~${END_DATE}")))
                .build()
        );
        List<RequestRowKeyHandle> result = new RowKeyGenerator(params, rowKeyGroup, rowKeyAggPeriod,
            RowKeyRequestType.RANGE).getRequestRowKeyHandles();
        assertEquals(expected, result);
    }


    /**
     * youtube - video_id
     */
    @Test
    public void getRangeKeySets_youtubeVideoId() throws NotSupportedException {
        RowKeyGroup rowKeyGroup = RowKeyGroup.VIDEO;
        RowKeyAggPeriod rowKeyAggPeriod = RowKeyAggPeriod.DAY;
        Params params = Params.builder().startDate(TEST_START_DATE).endDate(TEST_END_DATE)
            .videoIds(Set.of(TEST_VIDEO_IDS.get(0)))
            .build();

        List<RequestRowKeyHandle> expected = List.of(
            RequestRowKeyHandle.builder()
                .rowKeyPrefix(VIDEO_DATE)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("video_date_day~${VIDEO_ID_1}~${START_DATE}"),
                    substitutor.replace("video_date_day~${VIDEO_ID_1}~${END_DATE}")))
                .build()
        );
        List<RequestRowKeyHandle> result = new RowKeyGenerator(params, rowKeyGroup, rowKeyAggPeriod,
            RowKeyRequestType.RANGE).getRequestRowKeyHandles();
        assertEquals(expected, result);
    }

    /**
     * video - isrc
     */
    @Test
    public void getRangeKeySets_videoIsrc() throws NotSupportedException {
        RowKeyGroup rowKeyGroup = RowKeyGroup.VIDEO;
        RowKeyAggPeriod rowKeyAggPeriod = RowKeyAggPeriod.DAY;
        Params params = Params.builder().startDate(TEST_START_DATE).endDate(TEST_END_DATE)
            .isrc(List.of(TEST_ISRCS.get(0)))
            .build();

        List<RequestRowKeyHandle> expected = List.of(
            RequestRowKeyHandle.builder()
                .rowKeyPrefix(ISRC_DATE)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("isrc_date_day~${REV_ISRC_1}~${START_DATE}"),
                    substitutor.replace("isrc_date_day~${REV_ISRC_1}~${END_DATE}")))
                .build()
        );
        List<RequestRowKeyHandle> result = new RowKeyGenerator(params, rowKeyGroup, rowKeyAggPeriod,
            RowKeyRequestType.RANGE).getRequestRowKeyHandles();
        assertEquals(expected, result);
    }

    /**
     * video - isrc, content_types
     */
    @Test
    public void getRangeKeySets_videoIsrcContentTypes() throws NotSupportedException {
        RowKeyGroup rowKeyGroup = RowKeyGroup.VIDEO;
        RowKeyAggPeriod rowKeyAggPeriod = RowKeyAggPeriod.DAY;
        Params params = Params.builder().startDate(TEST_START_DATE).endDate(TEST_END_DATE)
            .isrc(List.of(TEST_ISRCS.get(0)))
            .youTubeContentTypes(Set.copyOf(TEST_VIDEO_CONTENT_TYPES))
            .build();

        List<RequestRowKeyHandle> expected = List.of(
            RequestRowKeyHandle.builder()
                .rowKeyPrefix(ISRC_CONTENT_DATE)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("isrc_content_date_day~${REV_ISRC_1}~${CONTENT_TYPE_1}~${START_DATE}"),
                    substitutor.replace("isrc_content_date_day~${REV_ISRC_1}~${CONTENT_TYPE_1}~${END_DATE}")))
                .build(),

            RequestRowKeyHandle.builder()
                .rowKeyPrefix(ISRC_CONTENT_DATE)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("isrc_content_date_day~${REV_ISRC_1}~${CONTENT_TYPE_2}~${START_DATE}"),
                    substitutor.replace("isrc_content_date_day~${REV_ISRC_1}~${CONTENT_TYPE_2}~${END_DATE}")))
                .build()
        );
        List<RequestRowKeyHandle> result = new RowKeyGenerator(params, rowKeyGroup, rowKeyAggPeriod,
            RowKeyRequestType.RANGE).getRequestRowKeyHandles();
        assertEquals(expected, result);
    }

    /**
     * isrc - tiktok-content_type
     */
    @Test
    public void getRangeKeySets_tikTokContentTypeIsrc() throws NotSupportedException {
        RowKeyGroup rowKeyGroup = RowKeyGroup.TIKTOK;
        RowKeyAggPeriod rowKeyAggPeriod = RowKeyAggPeriod.EMPTY;
        Params params = Params.builder().startDate(TEST_START_DATE).endDate(TEST_END_DATE)
            .isrc(List.of(TEST_ISRCS.get(0)))
            .tikTokContentTypes(Set.of(TikTokContentType.UGC))
            .build();

        List<RequestRowKeyHandle> expected = List.of(
            RequestRowKeyHandle.builder()
                .rowKeyPrefix(ISRC_TIKTOK_CONTENT_TYPE_DATE)
                .rowKeyRequestType(RowKeyRequestType.RANGE)
                .rowKey(Pair.of(
                    substitutor.replace("isrc_content_type_date~${REV_ISRC_1}~UGC~${START_DATE}"),
                    substitutor.replace("isrc_content_type_date~${REV_ISRC_1}~UGC~${END_DATE}")))
                .build()
        );
        List<RequestRowKeyHandle> result = new RowKeyGenerator(params, rowKeyGroup, rowKeyAggPeriod,
            RowKeyRequestType.RANGE).getRequestRowKeyHandles();
        assertEquals(expected, result);
    }

}