package io.delphiplatform.api.v3.bigtable;

import com.google.cloud.bigtable.data.v2.models.Query;
import com.google.cloud.bigtable.data.v2.models.Row;
import com.sonymusic.delphi.etl.apps.proto.Amazon.AmazonCountryStats;
import com.sonymusic.delphi.etl.apps.proto.Amazon.AmazonStreamStats;
import com.sonymusic.delphi.etl.apps.proto.Apple.AppleCountryStats;
import com.sonymusic.delphi.etl.apps.proto.Apple.AppleStreamStats;
import com.sonymusic.delphi.etl.apps.proto.Spotify.SpotifyCountryStats;
import com.sonymusic.delphi.etl.apps.proto.Spotify.SpotifyStreamStats;

import org.apache.commons.lang3.tuple.Pair;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;

import java.time.LocalDate;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Locale;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;

import io.delphiplatform.api.integration.util.MockCurrentDateService;
import io.delphiplatform.api.utils.TestUtils;
import io.delphiplatform.api.util.DateUtils;
import io.delphiplatform.api.v3.bigtable.entity.AmazonTrackStream;
import io.delphiplatform.api.v3.bigtable.entity.AppleTrackStream;
import io.delphiplatform.api.v3.bigtable.entity.SpotifyTrackStream;
import io.delphiplatform.api.v3.bigtable.processing.PaginationService;
import io.delphiplatform.api.v3.bigtable.processing.StreamDatesConverter;
import io.delphiplatform.api.v3.bigtable.processing.StreamDatesSortingService;
import io.delphiplatform.api.v3.bigtable.reader.BigtableColumns;
import io.delphiplatform.api.v3.constant.ApplicationConstants;
import io.delphiplatform.api.v3.constant.CountryCodeConstant;
import io.delphiplatform.api.v3.constant.DspConstants;
import io.delphiplatform.api.v3.model.StreamDate;
import io.delphiplatform.api.v3.view.util.Params;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

class BigtableStreamDatesServiceTest {

    private BigtableStreamDatesService datesService;

    private final BigtableQueryExecutor queryExecutor = mock(BigtableQueryExecutor.class);
    private final BigtableRowReader rowReader = mock(BigtableRowReader.class);
    private final StreamDatesSortingService sortingService = spy(new StreamDatesSortingService());
    private final PaginationService paginationService = spy(new PaginationService());
    private final StreamDatesConverter datesConverter = spy(new StreamDatesConverter());

    @BeforeEach
    public void setUp() {
        datesService = new BigtableStreamDatesService(queryExecutor, rowReader, sortingService, paginationService,
            datesConverter);
    }

    @Test
    void getStreamDates_oneIsrc_amazon() throws ExecutionException, InterruptedException {
        MockCurrentDateService.setMockDateTime(LocalDateTime.parse("2023-02-13T00:00:00"));

        String isrc1 = "isrc1";

        Row amazonUnlimitedFirst = mock(Row.class);
        Row amazonAdSupportedFirst = mock(Row.class);
        Row amazonAdSupportedSecond = mock(Row.class);
        when(queryExecutor.executeQuery(any(), any()))
            .thenReturn(CompletableFuture.completedFuture(
                List.of(amazonUnlimitedFirst, amazonAdSupportedFirst, amazonAdSupportedSecond)));
        AmazonTrackStream unlimitedTrackStreamFirst = createAmazonTrackStream(DspConstants.UNLIMITED,
            LocalDate.of(2020, 1, 1),
            isrc1,
            10);
        AmazonTrackStream adSupportedTrackStreamFirst = createAmazonTrackStream(DspConstants.ADSUPPORTED,
            LocalDate.of(2020, 1, 1),
            isrc1,
            20);
        AmazonTrackStream adSupportedTrackStreamSecond = createAmazonTrackStream(DspConstants.ADSUPPORTED,
            LocalDate.of(2020, 1, 2),
            isrc1,
            100);
        when(rowReader.readRow(AmazonTrackStream.class, amazonUnlimitedFirst)).thenReturn(unlimitedTrackStreamFirst);
        when(rowReader.readRow(AmazonTrackStream.class, amazonAdSupportedFirst))
            .thenReturn(adSupportedTrackStreamFirst);
        when(rowReader.readRow(AmazonTrackStream.class, amazonAdSupportedSecond))
            .thenReturn(adSupportedTrackStreamSecond);

        List<StreamDate> streamDates = datesService.getStreamDates(Params.builder()
            .isrc(List.of(isrc1))
            .dsp(List.of(DspConstants.AMAZON))
            .countryCode(Set.of(CountryCodeConstant.WORLDWIDE))
            .startDate(ApplicationConstants.MIN_STREAMS_DATE)
            .endDate(DateUtils.getCurrentDate())
            .build()).get();

        ArgumentCaptor<Query> queryArgumentCaptor = ArgumentCaptor.forClass(Query.class);
        verify(queryExecutor)
            .executeQuery(queryArgumentCaptor.capture(), eq(Collections.singletonList(BigtableColumns.DEMOGRAPHICS)));
        Query sentQuery = queryArgumentCaptor.getValue();
        List<Pair<String, String>> sentRanges = TestUtils.getRanges(sentQuery);
        long sentLimit = TestUtils.getLimit(sentQuery);

        assertEquals(10 * DspConstants.AMAZON_SUB_DSPS.size(), sentLimit);
        assertEquals(1, sentRanges.size());
        assertEquals("isrc_date_day~1crsi~2014-01-01", sentRanges.get(0).getLeft());
        assertEquals("isrc_date_day~1crsi~" + DateUtils.getCurrentDate().plusDays(1), sentRanges.get(0).getRight());

        assertEquals(1, streamDates.size());
        assertEquals(LocalDate.of(2020, 1, 1), streamDates.get(0).getDate());
        assertEquals(CountryCodeConstant.WORLDWIDE, streamDates.get(0).getCountryCode());
        assertEquals(isrc1, streamDates.get(0).getIsrc());
        assertEquals(30, streamDates.get(0).getStreams());
    }

    @Test
    void getStreamDates_oneIsrc_amazon_withTestStreamDate() throws ExecutionException, InterruptedException {
        MockCurrentDateService.setMockDateTime(LocalDateTime.parse("2023-02-13T00:00:00"));

        String isrc1 = "isrc1";

        Row amazonUnlimitedTestDate = mock(Row.class);
        Row amazonAdSupportedTestDate = mock(Row.class);
        Row amazonAdSupportedSecond = mock(Row.class);
        when(queryExecutor.executeQuery(any(), any()))
            .thenReturn(CompletableFuture.completedFuture(
                List.of(amazonUnlimitedTestDate, amazonAdSupportedTestDate, amazonAdSupportedSecond)));
        AmazonTrackStream unlimitedTrackStreamTestDate = createAmazonTrackStream(DspConstants.UNLIMITED,
            LocalDate.of(2020, 1, 1),
            isrc1,
            10);
        AmazonTrackStream adSupportedTrackStreamTestDate = createAmazonTrackStream(DspConstants.ADSUPPORTED,
            LocalDate.of(2020, 1, 1),
            isrc1,
            20);
        AmazonTrackStream adSupportedTrackStreamSecond = createAmazonTrackStream(DspConstants.ADSUPPORTED,
            LocalDate.of(2020, 1, 3),
            isrc1,
            100);
        when(rowReader.readRow(AmazonTrackStream.class, amazonUnlimitedTestDate))
            .thenReturn(unlimitedTrackStreamTestDate);
        when(rowReader.readRow(AmazonTrackStream.class, amazonAdSupportedTestDate))
            .thenReturn(adSupportedTrackStreamTestDate);
        when(rowReader.readRow(AmazonTrackStream.class, amazonAdSupportedSecond))
            .thenReturn(adSupportedTrackStreamSecond);

        List<StreamDate> streamDates = datesService.getStreamDates(Params.builder()
            .isrc(List.of(isrc1))
            .dsp(List.of(DspConstants.AMAZON))
            .countryCode(Set.of(CountryCodeConstant.WORLDWIDE))
            .startDate(ApplicationConstants.MIN_STREAMS_DATE)
            .endDate(DateUtils.getCurrentDate())
            .build()).get();

        assertEquals(1, streamDates.size());
        assertEquals(LocalDate.of(2020, 1, 3), streamDates.get(0).getDate());
        assertEquals(CountryCodeConstant.WORLDWIDE, streamDates.get(0).getCountryCode());
        assertEquals(isrc1, streamDates.get(0).getIsrc());
        assertEquals(100, streamDates.get(0).getStreams());
    }

    @Test
    void getStreamDates_oneIsrc_apple_withTestStreamDate() throws ExecutionException, InterruptedException {
        MockCurrentDateService.setMockDateTime(LocalDateTime.parse("2023-02-13T00:00:00"));

        String isrc1 = "isrc1";

        Row appleTestDate1 = mock(Row.class);
        Row appleTestDate2 = mock(Row.class);
        Row appleTestDate3 = mock(Row.class);
        when(queryExecutor.executeQuery(any(), any()))
            .thenReturn(CompletableFuture.completedFuture(
                List.of(appleTestDate1, appleTestDate2, appleTestDate3)));
        AppleTrackStream appleTrackStreamTestDate1 = createAppleTrackStream(
            LocalDate.of(2021, 1, 1),
            isrc1,
            10,
            Locale.US.getCountry());
        AppleTrackStream appleTrackStreamTestDate2 = createAppleTrackStream(
            LocalDate.of(2021, 1, 6),
            isrc1,
            20,
            CountryCodeConstant.WORLDWIDE);
        AppleTrackStream appleTrackStreamTestDate3 = createAppleTrackStream(
            LocalDate.of(2021, 1, 3),
            isrc1,
            100,
            CountryCodeConstant.WORLDWIDE);
        when(rowReader.readRow(AppleTrackStream.class, appleTestDate1))
            .thenReturn(appleTrackStreamTestDate1);
        when(rowReader.readRow(AppleTrackStream.class, appleTestDate2))
            .thenReturn(appleTrackStreamTestDate2);
        when(rowReader.readRow(AppleTrackStream.class, appleTestDate3))
            .thenReturn(appleTrackStreamTestDate3);

        List<StreamDate> streamDates = datesService.getStreamDates(Params.builder()
            .isrc(List.of(isrc1))
            .dsp(List.of(DspConstants.APPLE))
            .countryCode(Set.of(CountryCodeConstant.WORLDWIDE))
            .startDate(ApplicationConstants.MIN_STREAMS_DATE)
            .endDate(DateUtils.getCurrentDate())
            .build()).get();

        assertEquals(1, streamDates.size());
        assertEquals(LocalDate.of(2021, 1, 3), streamDates.get(0).getDate());
        assertEquals(CountryCodeConstant.WORLDWIDE, streamDates.get(0).getCountryCode());
        assertEquals(isrc1, streamDates.get(0).getIsrc());
        assertEquals(100, streamDates.get(0).getStreams());
    }

    @Test
    void getStreamDates_oneIsrc_spotify_withRetry() throws ExecutionException, InterruptedException {
        MockCurrentDateService.setMockDateTime(LocalDateTime.parse("2023-02-13T00:00:00"));

        String isrc1 = "isrc1";
        int numberOfEmptyStreams = 10;

        Row noStreamsRow = mock(Row.class);
        Row streamsRow = mock(Row.class);

        List<Row> noStreamsRowsForFirstQuery = Collections.nCopies(numberOfEmptyStreams, noStreamsRow);

        when(queryExecutor.executeQuery(any(), any()))
            .thenReturn(CompletableFuture.completedFuture(noStreamsRowsForFirstQuery));

        when(queryExecutor.executeQuerySync(any(), any()))
            .thenReturn(CompletableFuture.completedFuture(List.of(streamsRow)));

        List<SpotifyTrackStream> emptyTrackStreams = new ArrayList<>();
        LocalDate currentDate = LocalDate.of(2020, 1, 1);
        for (int i = 0; i < numberOfEmptyStreams; i++) {
            emptyTrackStreams.add(createSpotifyTrackStream(
                currentDate,
                isrc1,
                null));
            currentDate = currentDate.plusDays(1);
        }

        when(rowReader.readRow(SpotifyTrackStream.class, noStreamsRow))
            .thenReturn(emptyTrackStreams.get(0),
                emptyTrackStreams.subList(1, emptyTrackStreams.size()).toArray(new SpotifyTrackStream[0]));

        when(rowReader.readRow(SpotifyTrackStream.class, streamsRow))
            .thenReturn(createSpotifyTrackStream(LocalDate.of(2020, 1, 11), isrc1, 100));

        List<StreamDate> streamDates = datesService.getStreamDates(Params.builder()
            .isrc(List.of(isrc1))
            .dsp(List.of(DspConstants.SPOTIFY))
            .countryCode(Set.of(CountryCodeConstant.WORLDWIDE))
            .startDate(ApplicationConstants.MIN_STREAMS_DATE)
            .endDate(DateUtils.getCurrentDate())
            .build()).get();

        ArgumentCaptor<Query> queryArgumentCaptor = ArgumentCaptor.forClass(Query.class);
        verify(queryExecutor, times(1))
            .executeQuery(queryArgumentCaptor.capture(), eq(Collections.singletonList(BigtableColumns.DEMOGRAPHICS)));
        verify(queryExecutor, times(1))
            .executeQuerySync(queryArgumentCaptor.capture(), eq(Collections.singletonList(BigtableColumns.DEMOGRAPHICS)));
        assertEquals(2, queryArgumentCaptor.getAllValues().size());

        Query firstQuery = queryArgumentCaptor.getAllValues().get(0);
        List<Pair<String, String>> firstSentRanges = TestUtils.getRanges(firstQuery);
        long firstSentLimit = TestUtils.getLimit(firstQuery);

        Query secondQuery = queryArgumentCaptor.getAllValues().get(1);
        List<Pair<String, String>> secondSentRanges = TestUtils.getRanges(secondQuery);
        long secondSentLimit = TestUtils.getLimit(secondQuery);

        assertEquals(10, firstSentLimit);
        assertEquals(1, firstSentRanges.size());
        assertEquals("isrc_date_day~1crsi~2014-01-01", firstSentRanges.get(0).getLeft());
        assertEquals("isrc_date_day~1crsi~" + DateUtils.getCurrentDate().plusDays(1), firstSentRanges.get(0).getRight());

        assertEquals(20, secondSentLimit);
        assertEquals(1, secondSentRanges.size());
        assertEquals("isrc_date_day~1crsi~2020-01-10", secondSentRanges.get(0).getLeft());
        assertEquals("isrc_date_day~1crsi~" + DateUtils.getCurrentDate().plusDays(1), secondSentRanges.get(0).getRight());

        assertEquals(1, streamDates.size());
        assertEquals(LocalDate.of(2020, 1, 11), streamDates.get(0).getDate());
        assertEquals(CountryCodeConstant.WORLDWIDE, streamDates.get(0).getCountryCode());
        assertEquals(isrc1, streamDates.get(0).getIsrc());
        assertEquals(100, streamDates.get(0).getStreams());
    }

    private AmazonTrackStream createAmazonTrackStream(String fullDsp, LocalDate date, String isrc, Integer streams) {
        AmazonTrackStream stream = new AmazonTrackStream();
        stream.setFullDsp(fullDsp);
        stream.setDate(date);
        stream.setIsrc(isrc);
        if (stream != null) {
            stream.setStreams(AmazonCountryStats.newBuilder()
                .putCountries(CountryCodeConstant.WORLDWIDE, AmazonStreamStats.newBuilder().setStreams(streams).build())
                .build());
        }

        return stream;
    }

    private SpotifyTrackStream createSpotifyTrackStream(LocalDate date, String isrc, Integer streams) {
        SpotifyTrackStream stream = new SpotifyTrackStream();
        stream.setDate(date);
        stream.setIsrc(isrc);
        if (streams != null) {
            stream.setStreams(SpotifyCountryStats.newBuilder()
                .putCountries(CountryCodeConstant.WORLDWIDE,
                    SpotifyStreamStats.newBuilder().setStreams(streams).build())
                .build());
        }

        return stream;
    }

    private AppleTrackStream createAppleTrackStream(LocalDate date, String isrc, Integer streams, String countryCode) {
        AppleTrackStream stream = new AppleTrackStream();
        stream.setDate(date);
        stream.setIsrc(isrc);
        stream.setDsp(DspConstants.APPLE);
        if (streams != null) {
            stream.setStreams(AppleCountryStats.newBuilder()
                .putCountries(countryCode,
                    AppleStreamStats.newBuilder().setStreams(streams).build())
                .build());
        }

        return stream;
    }
}