/*
 * 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.beam.sdk.extensions.ordered.combiner;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.atomic.AtomicLong;
import org.apache.beam.sdk.extensions.ordered.ContiguousSequenceRange;
import org.joda.time.Instant;
import org.junit.Assert;
import org.junit.Test;

public class SequenceRangeAccumulatorTest {

  // Atomic just in case tests are run in parallel
  private static final AtomicLong currentTicker = new AtomicLong();

  static Instant nextTimestamp() {
    return Instant.ofEpochMilli(currentTicker.getAndIncrement());
  }

  static Instant eventTimestamp(Event[] events, long eventSequence) {
    for (Event e : events) {
      if (e.sequence == eventSequence) {
        return e.timestamp;
      }
    }
    throw new IllegalStateException("Unable to find event with sequence " + eventSequence);
  }

  static class Event {

    long sequence;
    Instant timestamp;
    boolean initialEvent;

    Event(long sequence, Instant ts) {
      this(sequence, ts, false);
    }

    Event(long sequence, Instant ts, boolean initialEvent) {
      this.sequence = sequence;
      this.timestamp = ts;
      this.initialEvent = initialEvent;
    }
  }

  @Test
  public void testSimpleAccumulation() {
    Event[] events =
        new Event[] {
          new Event(1, nextTimestamp(), true),
          new Event(2, nextTimestamp()),
          new Event(3, nextTimestamp())
        };

    doTestAccumulation(events, ContiguousSequenceRange.of(1, 4, eventTimestamp(events, 3)), 1);
  }

  @Test
  public void testReverseArrivalHandling() {
    Event[] events =
        new Event[] {
          new Event(3, nextTimestamp()),
          new Event(2, nextTimestamp()),
          new Event(1, nextTimestamp(), true)
        };

    Instant timestampOfEventNumber1 = eventTimestamp(events, 1);
    doTestAccumulation(events, ContiguousSequenceRange.of(1, 4, timestampOfEventNumber1), 1);
  }

  @Test
  public void testPartialRangeAccumulation() {
    Event[] events =
        new Event[] {
          new Event(1, nextTimestamp(), true),
          new Event(2, nextTimestamp()),
          new Event(3, nextTimestamp()),
          new Event(5, nextTimestamp()),
          new Event(7, nextTimestamp()),
        };

    doTestAccumulation(events, ContiguousSequenceRange.of(1, 4, eventTimestamp(events, 3)), 3);
  }

  @Test
  public void testMergingRangeAccumulation() {
    Event[] events =
        new Event[] {
          new Event(1, nextTimestamp(), true),
          new Event(2, nextTimestamp()),
          new Event(3, nextTimestamp()),
          new Event(5, nextTimestamp()),
          new Event(7, nextTimestamp()),
          new Event(6, nextTimestamp()),
        };

    doTestAccumulation(events, ContiguousSequenceRange.of(1, 4, eventTimestamp(events, 3)), 2);
  }

  @Test
  public void testNoStartEvent() {
    Event[] events =
        new Event[] {
          new Event(2, nextTimestamp()),
          new Event(3, nextTimestamp()),
          new Event(1, nextTimestamp()),
          new Event(5, nextTimestamp()),
        };

    doTestAccumulation(events, ContiguousSequenceRange.EMPTY, 2);
  }

  @Test
  public void testNoEventsAccumulation() {
    Event[] events = new Event[] {};

    doTestAccumulation(events, ContiguousSequenceRange.EMPTY, 0);
  }

  @Test
  public void testRemovingRangesBelowInitialSequenceDuringAccumulation() {
    Event[] events =
        new Event[] {
          // First range
          new Event(2, nextTimestamp()),
          new Event(3, nextTimestamp()),
          new Event(1, nextTimestamp()),

          // Second range
          new Event(5, nextTimestamp()),
          new Event(6, nextTimestamp()),

          // This event should prune everything below
          new Event(7, nextTimestamp(), true),
        };

    doTestAccumulation(events, ContiguousSequenceRange.of(7, 8, eventTimestamp(events, 7)), 1);
  }

  @Test
  public void testRemovingElementsBelowInitialSequenceDuringAccumulation() {

    Event[] events =
        new Event[] {
          // First range
          new Event(2, nextTimestamp()),
          new Event(3, nextTimestamp()),
          new Event(1, nextTimestamp()),

          // Second range
          new Event(5, nextTimestamp()),
          new Event(6, nextTimestamp()),
          new Event(7, nextTimestamp()),
          new Event(8, nextTimestamp()),

          // This event should reduce the range.
          new Event(7, nextTimestamp(), true),
        };

    Instant timestampOfTheLastEvent = events[events.length - 1].timestamp;
    doTestAccumulation(events, ContiguousSequenceRange.of(7, 9, timestampOfTheLastEvent), 1);
  }

  private static void doTestAccumulation(
      Event[] events, ContiguousSequenceRange expectedResult, int expectedNumberOfRanges) {
    SequenceRangeAccumulator accumulator = new SequenceRangeAccumulator();
    Arrays.stream(events).forEach(e -> accumulator.add(e.sequence, e.timestamp, e.initialEvent));

    Assert.assertEquals(
        "Accumulated results", expectedResult, accumulator.largestContinuousRange());

    Assert.assertEquals("Number of ranges", expectedNumberOfRanges, accumulator.numberOfRanges());
  }

  @Test
  public void testEmptyMerge() {
    Event[] set1 = new Event[] {};
    Event[] set2 = new Event[] {};

    ContiguousSequenceRange expectedResult = ContiguousSequenceRange.EMPTY;
    int expectedNumberOfRanges = 0;

    doTestMerging(set1, set2, expectedResult, expectedNumberOfRanges);
  }

  @Test
  public void testMergingNonEmptyWithEmpty() {
    Event[] set1 =
        new Event[] {
          new Event(3, nextTimestamp()),
          new Event(2, nextTimestamp()),
          new Event(1, nextTimestamp(), true)
        };
    Event[] set2 = new Event[] {};

    ContiguousSequenceRange expectedResult =
        ContiguousSequenceRange.of(1, 4, eventTimestamp(set1, 1L));
    int expectedNumberOfRanges = 1;

    doTestMerging(set1, set2, expectedResult, expectedNumberOfRanges);
  }

  @Test
  public void testMergingWithLowerNonAdjacentRange() {
    Event[] set1 =
        new Event[] {
          new Event(1, nextTimestamp(), true), new Event(2, nextTimestamp()),
        };
    Event[] set2 =
        new Event[] {
          new Event(4, nextTimestamp()),
          new Event(5, nextTimestamp()),
          new Event(6, nextTimestamp())
        };

    ContiguousSequenceRange expectedResult =
        ContiguousSequenceRange.of(1, 3, eventTimestamp(set1, 2L));
    int expectedNumberOfRanges = 2;

    doTestMerging(set1, set2, expectedResult, expectedNumberOfRanges);
  }

  @Test
  public void testMergingWithoutAnyInitialEvents() {
    Event[] set1 =
        new Event[] {
          new Event(1, nextTimestamp()), new Event(2, nextTimestamp()),
        };
    Event[] set2 =
        new Event[] {
          new Event(4, nextTimestamp()),
          new Event(5, nextTimestamp()),
          new Event(6, nextTimestamp())
        };

    ContiguousSequenceRange expectedResult = ContiguousSequenceRange.EMPTY;
    int expectedNumberOfRanges = 2;

    doTestMerging(set1, set2, expectedResult, expectedNumberOfRanges);
  }

  @Test
  public void testMergingAdjacentRanges() {
    Event[] set1 =
        new Event[] {
          new Event(1, nextTimestamp(), true), new Event(2, nextTimestamp()),
        };
    Event[] set2 =
        new Event[] {
          new Event(3, nextTimestamp()),
          new Event(4, nextTimestamp()),
          new Event(5, nextTimestamp()),
          new Event(6, nextTimestamp())
        };

    ContiguousSequenceRange expectedResult =
        ContiguousSequenceRange.of(1, 7, eventTimestamp(set2, 6L));
    int expectedNumberOfRanges = 1;

    doTestMerging(set1, set2, expectedResult, expectedNumberOfRanges);
  }

  @Test
  public void testMergingAdjacentRangesWithSingleValue() {
    Event[] set1 =
        new Event[] {
          new Event(4, nextTimestamp()),
        };
    Event[] set2 =
        new Event[] {
          new Event(1, nextTimestamp(), true),
          new Event(3, nextTimestamp()),
          new Event(2, nextTimestamp())
        };

    ContiguousSequenceRange expectedResult =
        ContiguousSequenceRange.of(1, 5, eventTimestamp(set2, 2L));
    int expectedNumberOfRanges = 1;

    doTestMerging(set1, set2, expectedResult, expectedNumberOfRanges);
  }

  @Test
  public void testLargeVolumeMerging() {
    int numberOfSets = 3000;
    ArrayList<Set<Event>> eventSets = new ArrayList<>(numberOfSets);

    for (int i = 0; i < numberOfSets; i++) {
      eventSets.add(new HashSet<>());
    }

    int eventCount = 100000;

    Random random = new Random();
    Instant lastTimestamp = null;
    for (int i = 0; i < eventCount; i++) {
      int setNumber = random.nextInt(numberOfSets);
      Set<Event> currentSet = eventSets.get(setNumber);
      boolean initialEvent = i == 0;
      lastTimestamp = nextTimestamp();
      currentSet.add(new Event(i, lastTimestamp, initialEvent));
    }

    final SequenceRangeAccumulator initialAccumulator = new SequenceRangeAccumulator();
    eventSets.get(0).forEach(e -> initialAccumulator.add(e.sequence, e.timestamp, e.initialEvent));

    for (int i = 1; i < numberOfSets; i++) {
      SequenceRangeAccumulator nextAccumulator = new SequenceRangeAccumulator();
      eventSets.get(i).forEach(e -> nextAccumulator.add(e.sequence, e.timestamp, e.initialEvent));
      initialAccumulator.merge(nextAccumulator);
    }

    ContiguousSequenceRange expectedResult =
        ContiguousSequenceRange.of(0, eventCount, lastTimestamp);
    Assert.assertEquals(
        "Accumulated results", expectedResult, initialAccumulator.largestContinuousRange());

    Assert.assertEquals("Number of ranges", 1, initialAccumulator.numberOfRanges());
  }

  @Test
  public void testPruningSequencesBelowInitial() {
    Event[] set1 =
        new Event[] {
          new Event(1, nextTimestamp()), new Event(2, nextTimestamp()),
        };
    Event[] set2 =
        new Event[] {
          new Event(3, nextTimestamp(), true),
          new Event(4, nextTimestamp()),
          new Event(5, nextTimestamp()),
          new Event(6, nextTimestamp())
        };

    ContiguousSequenceRange expectedResult =
        ContiguousSequenceRange.of(3, 7, eventTimestamp(set2, 6L));
    int expectedNumberOfRanges = 1;

    doTestMerging(set1, set2, expectedResult, expectedNumberOfRanges);
  }

  @Test
  public void testDuplicateHandling() {
    Event[] set1 =
        new Event[] {
          new Event(1, nextTimestamp(), true),
          new Event(2, nextTimestamp()),
          new Event(3, nextTimestamp()),
          new Event(5, nextTimestamp()),
        };
    Event[] set2 =
        new Event[] {
          new Event(3, nextTimestamp()),
          new Event(4, nextTimestamp()),
          new Event(5, nextTimestamp()),
          new Event(6, nextTimestamp())
        };

    ContiguousSequenceRange expectedResult =
        ContiguousSequenceRange.of(1, 7, eventTimestamp(set2, 6L));
    int expectedNumberOfRanges = 1;

    doTestMerging(set1, set2, expectedResult, expectedNumberOfRanges);
  }

  @Test
  public void testExceptionThrownIfThereAreDifferentInitialSequences() {
    Event[] set1 =
        new Event[] {
          new Event(1, nextTimestamp(), true), new Event(2, nextTimestamp()),
        };
    Event[] set2 =
        new Event[] {
          new Event(3, nextTimestamp(), true),
          new Event(4, nextTimestamp()),
          new Event(5, nextTimestamp()),
          new Event(6, nextTimestamp())
        };

    try {
      doTestMerging(set1, set2, ContiguousSequenceRange.EMPTY, 0);
      Assert.fail("Expected to throw an exception");
    } catch (IllegalStateException e) {
      Assert.assertEquals(
          "Exception message",
          "Two accumulators contain different initial sequences: 1 and 3",
          e.getMessage());
    }
  }

  @Test
  public void testSelectingHighestTimestampWhenMerging() {
    Event[] set1 =
        new Event[] {
          new Event(1, nextTimestamp(), true),
          new Event(2, Instant.ofEpochMilli(currentTicker.get() + 10000)),
        };
    Event[] set2 =
        new Event[] {
          new Event(3, nextTimestamp()),
          new Event(4, nextTimestamp()),
          new Event(5, nextTimestamp()),
          new Event(6, nextTimestamp())
        };

    ContiguousSequenceRange expectedResult =
        ContiguousSequenceRange.of(1, 7, eventTimestamp(set1, 2L));
    int expectedNumberOfRanges = 1;
    doTestMerging(set1, set2, expectedResult, expectedNumberOfRanges);
  }

  private static void doTestMerging(
      Event[] set1,
      Event[] set2,
      ContiguousSequenceRange expectedResult,
      int expectedNumberOfRanges) {
    // Try to merge both set2 to set1 and set1 to set2 - both must return the same results
    mergeAndTest(set1, set2, expectedResult, expectedNumberOfRanges, "set1");
    mergeAndTest(set2, set1, expectedResult, expectedNumberOfRanges, "set2");
  }

  private static void mergeAndTest(
      Event[] set1,
      Event[] set2,
      ContiguousSequenceRange expectedResult,
      int expectedNumberOfRanges,
      String firstSetName) {
    final SequenceRangeAccumulator a1 = new SequenceRangeAccumulator();
    Arrays.stream(set1).forEach(e -> a1.add(e.sequence, e.timestamp, e.initialEvent));

    final SequenceRangeAccumulator a2 = new SequenceRangeAccumulator();
    Arrays.stream(set2).forEach(e -> a2.add(e.sequence, e.timestamp, e.initialEvent));

    a1.merge(a2);

    Assert.assertEquals(
        "Accumulated results - " + firstSetName, expectedResult, a1.largestContinuousRange());

    Assert.assertEquals(
        "Number of ranges - " + firstSetName, expectedNumberOfRanges, a1.numberOfRanges());
  }
}
