/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package org.apache.druid.query.scan;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSortedSet;
import com.google.common.collect.Sets;
import org.apache.druid.java.util.common.DateTimes;
import org.apache.druid.java.util.common.Intervals;
import org.apache.druid.java.util.common.Pair;
import org.apache.druid.java.util.common.concurrent.Execs;
import org.apache.druid.java.util.common.guava.MergeSequence;
import org.apache.druid.java.util.common.guava.Sequence;
import org.apache.druid.java.util.common.guava.Sequences;
import org.apache.druid.query.DefaultGenericQueryMetricsFactory;
import org.apache.druid.query.Druids;
import org.apache.druid.query.QueryPlus;
import org.apache.druid.query.QueryRunner;
import org.apache.druid.query.context.ResponseContext;
import org.apache.druid.query.spec.MultipleIntervalSegmentSpec;
import org.apache.druid.query.spec.MultipleSpecificSegmentSpec;
import org.apache.druid.segment.RowAdapter;
import org.apache.druid.segment.RowBasedSegment;
import org.apache.druid.segment.column.ColumnHolder;
import org.apache.druid.segment.column.ColumnType;
import org.apache.druid.segment.column.RowSignature;
import org.apache.druid.timeline.SegmentId;
import org.joda.time.DateTime;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
import java.util.Set;
import java.util.TreeSet;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

/**
 * Tests the order in which Scan query results come back.
 *
 * Ensures that we have run-to-run stability of result order, which is important for offset-based pagination.
 */
@RunWith(Parameterized.class)
public class ScanQueryResultOrderingTest
{
  private static final String DATASOURCE = "datasource";
  private static final String ID_COLUMN = "id";

  private static final RowAdapter<Object[]> ROW_ADAPTER = columnName -> {
    if (ID_COLUMN.equals(columnName)) {
      return row -> row[1];
    } else if (ColumnHolder.TIME_COLUMN_NAME.equals(columnName)) {
      return row -> ((DateTime) row[0]).getMillis();
    } else {
      return row -> null;
    }
  };

  private static final RowSignature ROW_SIGNATURE = RowSignature.builder()
                                                                .addTimeColumn()
                                                                .add(ID_COLUMN, ColumnType.LONG)
                                                                .build();

  private static final List<RowBasedSegment<Object[]>> SEGMENTS = ImmutableList.of(
      new RowBasedSegment<>(
          SegmentId.of(DATASOURCE, Intervals.of("2000-01-01/P1D"), "1", 0),
          Sequences.simple(
              ImmutableList.of(
                  new Object[]{DateTimes.of("2000T01"), 101},
                  new Object[]{DateTimes.of("2000T01"), 80},
                  new Object[]{DateTimes.of("2000T01"), 232},
                  new Object[]{DateTimes.of("2000T01"), 12},
                  new Object[]{DateTimes.of("2000T02"), 808},
                  new Object[]{DateTimes.of("2000T02"), 411},
                  new Object[]{DateTimes.of("2000T02"), 383},
                  new Object[]{DateTimes.of("2000T05"), 22}
              )
          ),
          ROW_ADAPTER,
          ROW_SIGNATURE
      ),
      new RowBasedSegment<>(
          SegmentId.of(DATASOURCE, Intervals.of("2000-01-01/P1D"), "1", 1),
          Sequences.simple(
              ImmutableList.of(
                  new Object[]{DateTimes.of("2000T01"), 333},
                  new Object[]{DateTimes.of("2000T01"), 222},
                  new Object[]{DateTimes.of("2000T01"), 444},
                  new Object[]{DateTimes.of("2000T01"), 111},
                  new Object[]{DateTimes.of("2000T03"), 555},
                  new Object[]{DateTimes.of("2000T03"), 999},
                  new Object[]{DateTimes.of("2000T03"), 888},
                  new Object[]{DateTimes.of("2000T05"), 777}
              )
          ),
          ROW_ADAPTER,
          ROW_SIGNATURE
      ),
      new RowBasedSegment<>(
          SegmentId.of(DATASOURCE, Intervals.of("2000-01-02/P1D"), "1", 0),
          Sequences.simple(
              ImmutableList.of(
                  new Object[]{DateTimes.of("2000-01-02T00"), 7},
                  new Object[]{DateTimes.of("2000-01-02T02"), 9},
                  new Object[]{DateTimes.of("2000-01-02T03"), 8}
              )
          ),
          ROW_ADAPTER,
          ROW_SIGNATURE
      )
  );

  private final List<Integer> segmentToServerMap;
  private final int limit;
  private final int batchSize;
  private final int maxRowsQueuedForOrdering;

  private ScanQueryRunnerFactory queryRunnerFactory;
  private List<QueryRunner<ScanResultValue>> segmentRunners;

  @Parameterized.Parameters(name = "Segment-to-server map[{0}], limit[{1}], batchSize[{2}], maxRowsQueuedForOrdering[{3}]")
  public static Iterable<Object[]> constructorFeeder()
  {
    // Set number of server equal to number of segments, then try all possible distributions of segments to servers.
    final int numServers = SEGMENTS.size();

    final Set<List<Integer>> segmentToServerMaps = Sets.cartesianProduct(
        IntStream.range(0, SEGMENTS.size())
                 .mapToObj(i -> IntStream.range(0, numServers).boxed().collect(Collectors.toSet()))
                 .collect(Collectors.toList())
    );

    // Try every limit up to one past the total number of rows.
    final Set<Integer> limits = new TreeSet<>();
    final int totalNumRows = SEGMENTS.stream().mapToInt(s -> s.asStorageAdapter().getNumRows()).sum();
    for (int i = 0; i <= totalNumRows + 1; i++) {
      limits.add(i);
    }

    // Try various batch sizes.
    final Set<Integer> batchSizes = ImmutableSortedSet.of(1, 2, 100);
    final Set<Integer> maxRowsQueuedForOrderings = ImmutableSortedSet.of(1, 7, 100000);

    return Sets.cartesianProduct(
        segmentToServerMaps,
        limits,
        batchSizes,
        maxRowsQueuedForOrderings
    ).stream().map(args -> args.toArray(new Object[0])).collect(Collectors.toList());
  }

  public ScanQueryResultOrderingTest(
      final List<Integer> segmentToServerMap,
      final int limit,
      final int batchSize,
      final int maxRowsQueuedForOrdering
  )
  {
    this.segmentToServerMap = segmentToServerMap;
    this.limit = limit;
    this.batchSize = batchSize;
    this.maxRowsQueuedForOrdering = maxRowsQueuedForOrdering;
  }

  @Before
  public void setUp()
  {
    queryRunnerFactory = new ScanQueryRunnerFactory(
        new ScanQueryQueryToolChest(
            new ScanQueryConfig(),
            new DefaultGenericQueryMetricsFactory()
        ),
        new ScanQueryEngine(),
        new ScanQueryConfig()
    );

    segmentRunners = SEGMENTS.stream().map(queryRunnerFactory::createRunner).collect(Collectors.toList());
  }

  @Test
  public void testOrderNone()
  {
    assertResultsEquals(
        Druids.newScanQueryBuilder()
              .dataSource("ds")
              .intervals(new MultipleIntervalSegmentSpec(Collections.singletonList(Intervals.of("2000/P1D"))))
              .columns(ColumnHolder.TIME_COLUMN_NAME, ID_COLUMN)
              .order(ScanQuery.Order.NONE)
              .build(),
        ImmutableList.of(
            101,
            80,
            232,
            12,
            808,
            411,
            383,
            22,
            333,
            222,
            444,
            111,
            555,
            999,
            888,
            777,
            7,
            9,
            8
        )
    );
  }

  @Test
  public void testOrderTimeAscending()
  {
    assertResultsEquals(
        Druids.newScanQueryBuilder()
              .dataSource("ds")
              .intervals(new MultipleIntervalSegmentSpec(Collections.singletonList(Intervals.of("2000/P1D"))))
              .columns(ColumnHolder.TIME_COLUMN_NAME, ID_COLUMN)
              .order(ScanQuery.Order.ASCENDING)
              .build(),
        ImmutableList.of(
            101,
            80,
            232,
            12,
            333,
            222,
            444,
            111,
            808,
            411,
            383,
            555,
            999,
            888,
            22,
            777,
            7,
            9,
            8
        )
    );
  }

  @Test
  public void testOrderTimeDescending()
  {
    assertResultsEquals(
        Druids.newScanQueryBuilder()
              .dataSource("ds")
              .intervals(new MultipleIntervalSegmentSpec(Collections.singletonList(Intervals.of("2000/P1D"))))
              .columns(ColumnHolder.TIME_COLUMN_NAME, ID_COLUMN)
              .order(ScanQuery.Order.DESCENDING)
              .build(),
        ImmutableList.of(
            8,
            9,
            7,
            777,
            22,
            888,
            999,
            555,
            383,
            411,
            808,
            111,
            444,
            222,
            333,
            12,
            232,
            80,
            101
        )
    );
  }

  private void assertResultsEquals(final ScanQuery query, final List<Integer> expectedResults)
  {
    final List<List<Pair<SegmentId, QueryRunner<ScanResultValue>>>> serverRunners = new ArrayList<>();
    for (int i = 0; i <= segmentToServerMap.stream().max(Comparator.naturalOrder()).orElse(0); i++) {
      serverRunners.add(new ArrayList<>());
    }

    for (int segmentNumber = 0; segmentNumber < segmentToServerMap.size(); segmentNumber++) {
      final SegmentId segmentId = SEGMENTS.get(segmentNumber).getId();
      final int serverNumber = segmentToServerMap.get(segmentNumber);

      serverRunners.get(serverNumber).add(Pair.of(segmentId, segmentRunners.get(segmentNumber)));
    }

    // Simulates what the Historical servers would do.
    final List<QueryRunner<ScanResultValue>> mergedServerRunners =
        serverRunners.stream()
                     .filter(runners -> !runners.isEmpty())
                     .map(
                         runners ->
                             queryRunnerFactory.getToolchest().mergeResults(
                                 new QueryRunner<ScanResultValue>()
                                 {
                                   @Override
                                   public Sequence<ScanResultValue> run(
                                       final QueryPlus<ScanResultValue> queryPlus,
                                       final ResponseContext responseContext
                                   )
                                   {
                                     return queryRunnerFactory.mergeRunners(
                                         Execs.directExecutor(),
                                         runners.stream().map(p -> p.rhs).collect(Collectors.toList())
                                     ).run(
                                         queryPlus.withQuery(
                                             queryPlus.getQuery()
                                                      .withQuerySegmentSpec(
                                                          new MultipleSpecificSegmentSpec(
                                                              runners.stream()
                                                                     .map(p -> p.lhs.toDescriptor())
                                                                     .collect(Collectors.toList())
                                                          )
                                                      )
                                         ),
                                         responseContext
                                     );
                                   }
                                 }
                             )
                     )
                     .collect(Collectors.toList());

    // Simulates what the Broker would do.
    final QueryRunner<ScanResultValue> brokerRunner = queryRunnerFactory.getToolchest().mergeResults(
        (queryPlus, responseContext) -> {
          final List<Sequence<ScanResultValue>> sequences =
              mergedServerRunners.stream()
                                 .map(runner -> runner.run(queryPlus.withoutThreadUnsafeState()))
                                 .collect(Collectors.toList());

          return new MergeSequence<>(
              queryPlus.getQuery().getResultOrdering(),
              Sequences.simple(sequences)
          );
        }
    );

    // Finally: run the query.
    final List<Integer> results = runQuery(
        (ScanQuery) Druids.ScanQueryBuilder.copy(query)
                                           .limit(limit)
                                           .batchSize(batchSize)
                                           .build()
                                           .withOverriddenContext(
                                               ImmutableMap.of(
                                                   ScanQueryConfig.CTX_KEY_MAX_ROWS_QUEUED_FOR_ORDERING,
                                                   maxRowsQueuedForOrdering
                                               )
                                           ),
        brokerRunner
    );

    Assert.assertEquals(
        expectedResults.stream().limit(limit == 0 ? Long.MAX_VALUE : limit).collect(Collectors.toList()),
        results
    );
  }

  private List<Integer> runQuery(final ScanQuery query, final QueryRunner<ScanResultValue> brokerRunner)
  {
    final List<Object[]> results = queryRunnerFactory.getToolchest().resultsAsArrays(
        query,
        brokerRunner.run(QueryPlus.wrap(query))
    ).toList();

    return results.stream().mapToInt(row -> (int) row[1]).boxed().collect(Collectors.toList());
  }
}
