/*
 * 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.runners.dataflow.worker;

import static org.apache.beam.runners.dataflow.worker.ReaderUtils.readAllFromReader;
import static org.apache.beam.runners.dataflow.worker.ReaderUtils.readNItemsFromUnstartedIterator;
import static org.apache.beam.runners.dataflow.worker.ReaderUtils.readRemainingFromIterator;
import static org.apache.beam.runners.dataflow.worker.SourceTranslationUtils.readerProgressToCloudProgress;

import com.google.api.services.dataflow.model.ApproximateReportedProgress;
import java.io.File;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.nio.channels.Channels;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Random;
import org.apache.avro.Schema;
import org.apache.avro.file.DataFileWriter;
import org.apache.avro.generic.GenericDatumWriter;
import org.apache.avro.io.DatumWriter;
import org.apache.beam.runners.dataflow.worker.AvroByteReader.AvroByteFileIterator;
import org.apache.beam.runners.dataflow.worker.util.common.worker.ExecutorTestUtils;
import org.apache.beam.runners.dataflow.worker.util.common.worker.NativeReader.DynamicSplitResult;
import org.apache.beam.sdk.TestUtils;
import org.apache.beam.sdk.coders.BigEndianIntegerCoder;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.io.FileSystems;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.util.CoderUtils;
import org.apache.beam.sdk.util.MimeTypes;
import org.junit.Assert;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/** Tests for AvroByteReader. */
@RunWith(JUnit4.class)
@SuppressWarnings({
  "rawtypes", // TODO(https://github.com/apache/beam/issues/20447)
})
public class AvroByteReaderTest {
  @Rule public TemporaryFolder tmpFolder = new TemporaryFolder();

  /** Class representing information about an Avro file generated from a list of elements. */
  private static class AvroFileInfo<T> {
    String filename;
    List<Integer> elementSizes = new ArrayList<>();
    List<Long> syncPoints = new ArrayList<>();
    long totalElementEncodedSize = 0;
  }

  /** Write input elements to a file and return information about the Avro-encoded file. */
  private <T> AvroFileInfo<T> initInputFile(List<List<T>> elemsList, Coder<T> coder)
      throws Exception {
    File tmpFile = tmpFolder.newFile("file.avro");
    AvroFileInfo<T> fileInfo = new AvroFileInfo<>();
    fileInfo.filename = tmpFile.getPath();

    // Write the data.
    OutputStream outStream =
        Channels.newOutputStream(
            FileSystems.create(
                FileSystems.matchNewResource(fileInfo.filename, false), MimeTypes.BINARY));
    Schema schema = Schema.create(Schema.Type.BYTES);
    DatumWriter<ByteBuffer> datumWriter = new GenericDatumWriter<>(schema);
    try (DataFileWriter<ByteBuffer> fileWriter = new DataFileWriter<>(datumWriter)) {
      fileWriter.create(schema, outStream);
      boolean first = true;
      for (List<T> elems : elemsList) {
        if (first) {
          first = false;
        } else {
          // Ensure a block boundary here.
          long syncPoint = fileWriter.sync();
          fileInfo.syncPoints.add(syncPoint);
        }
        for (T elem : elems) {
          byte[] encodedElement = CoderUtils.encodeToByteArray(coder, elem);
          fileWriter.append(ByteBuffer.wrap(encodedElement));
          fileInfo.elementSizes.add(encodedElement.length);
          fileInfo.totalElementEncodedSize += encodedElement.length;
        }
      }
    }

    return fileInfo;
  }

  /**
   * Reads from a file generated from a collection of elements and verifies that the elements read
   * are the same as the elements written.
   *
   * @param elemsList a list of blocks of elements, each of which is as a list of elements.
   * @param coder the coder used to encode the elements
   * @param requireExactMatch if true, each block must match exactly
   * @throws Exception
   */
  private <T> void runTestRead(List<List<T>> elemsList, Coder<T> coder, boolean requireExactMatch)
      throws Exception {
    AvroFileInfo<T> fileInfo = initInputFile(elemsList, coder);

    // Test reading the data back.
    List<List<T>> actualElemsList = new ArrayList<>();
    List<Integer> actualSizes = new ArrayList<>();
    long startOffset = 0;
    long endOffset;
    long prevSyncPoint = 0;
    for (long syncPoint : fileInfo.syncPoints) {
      endOffset = (prevSyncPoint + syncPoint) / 2;
      actualElemsList.add(readElems(fileInfo.filename, startOffset, endOffset, coder, actualSizes));
      startOffset = endOffset;
      prevSyncPoint = syncPoint;
    }
    actualElemsList.add(
        readElems(fileInfo.filename, startOffset, Long.MAX_VALUE, coder, actualSizes));

    // Compare the expected and the actual elements.
    if (requireExactMatch) {
      // Require the blocks to match exactly.  (This works only for
      // small block sizes.  Large block sizes, bigger than Avro's
      // internal sizes, lead to different splits.)
      Assert.assertEquals(elemsList, actualElemsList);
    } else {
      // Just require the overall elements to be the same.  (This
      // works for any block size.)
      List<T> expected = new ArrayList<>();
      for (List<T> elems : elemsList) {
        expected.addAll(elems);
      }
      List<T> actual = new ArrayList<>();
      for (List<T> actualElems : actualElemsList) {
        actual.addAll(actualElems);
      }
      Assert.assertEquals(expected, actual);
    }

    long actualTotalSize = 0;
    for (int elemSize : actualSizes) {
      actualTotalSize += elemSize;
    }
    Assert.assertEquals(fileInfo.totalElementEncodedSize, actualTotalSize);
  }

  private <T> List<T> readElems(
      String filename, long startOffset, long endOffset, Coder<T> coder, List<Integer> actualSizes)
      throws Exception {
    AvroByteReader<T> avroReader =
        new AvroByteReader<>(
            filename, startOffset, endOffset, coder, PipelineOptionsFactory.create());
    new ExecutorTestUtils.TestReaderObserver(avroReader, actualSizes);
    return readAllFromReader(avroReader);
  }

  @Test
  public void testRead() throws Exception {
    runTestRead(
        Collections.singletonList(TestUtils.INTS),
        BigEndianIntegerCoder.of(),
        true /* require exact match */);
  }

  @Test
  public void testReadEmpty() throws Exception {
    runTestRead(
        Collections.singletonList(TestUtils.NO_INTS),
        BigEndianIntegerCoder.of(),
        true /* require exact match */);
  }

  private List<List<String>> generateInputBlocks(
      int numBlocks, int blockSizeBytes, int averageLineSizeBytes) {
    Random random = new Random(0);
    List<List<String>> blocks = new ArrayList<>(numBlocks);
    for (int blockNum = 0; blockNum < numBlocks; blockNum++) {
      int numLines = blockSizeBytes / averageLineSizeBytes;
      List<String> lines = new ArrayList<>(numLines);
      for (int lineNum = 0; lineNum < numLines; lineNum++) {
        int numChars = random.nextInt(averageLineSizeBytes * 2);
        StringBuilder sb = new StringBuilder();
        for (int charNum = 0; charNum < numChars; charNum++) {
          sb.appendCodePoint(random.nextInt('z' - 'a' + 1) + 'a');
        }
        lines.add(sb.toString());
      }
      blocks.add(lines);
    }
    return blocks;
  }

  @Test
  public void testReadSmallRanges() throws Exception {
    runTestRead(
        generateInputBlocks(3, 50, 5), StringUtf8Coder.of(), true /* require exact match */);
  }

  @Test
  public void testReadBigRanges() throws Exception {
    runTestRead(
        generateInputBlocks(10, 128 * 1024, 100),
        StringUtf8Coder.of(),
        false /* don't require exact match */);
  }

  // Verification behavior for split requests. Used for testRequestDynamicSplitInternal.
  private static enum SplitVerificationBehavior {
    VERIFY_SUCCESS, // Split request must succeed.
    VERIFY_FAILURE, // Split request must fail.
    DO_NOT_VERIFY; // Perform no verification.
  }

  private <T> void testRequestDynamicSplitInternal(
      AvroByteReader<T> reader,
      float splitAtFraction,
      int readBeforeSplit,
      SplitVerificationBehavior splitVerificationBehavior)
      throws Exception {
    // Read all elements from the reader
    Long endOffset = reader.endPosition;
    List<T> expectedElements = readAllFromReader(reader);

    List<T> primaryElements;
    List<T> residualElements = new ArrayList<>();
    try (AvroByteReader<T>.AvroByteFileIterator iterator = reader.iterator()) {
      // Read n elements from the reader
      primaryElements = readNItemsFromUnstartedIterator(iterator, readBeforeSplit);

      // Request a split at the specified position
      DynamicSplitResult splitResult =
          iterator.requestDynamicSplit(ReaderTestUtils.splitRequestAtFraction(splitAtFraction));

      switch (splitVerificationBehavior) {
        case VERIFY_SUCCESS:
          Assert.assertNotNull(splitResult);
          break;
        case VERIFY_FAILURE:
          Assert.assertNull(splitResult);
          break;
        case DO_NOT_VERIFY:
      }

      // Finish reading from the original reader.
      primaryElements.addAll(readRemainingFromIterator(iterator, readBeforeSplit > 0));

      if (splitResult != null) {
        Long splitPosition = ReaderTestUtils.positionFromSplitResult(splitResult).getByteOffset();
        AvroByteReader<T> residualReader =
            new AvroByteReader<T>(
                reader.avroSource.getFileOrPatternSpec(),
                splitPosition,
                endOffset,
                reader.coder,
                reader.options);
        // Read from the residual until it is complete.
        residualElements = readAllFromReader(residualReader);
      }
    }

    primaryElements.addAll(residualElements);
    Assert.assertEquals(expectedElements, primaryElements);
    if (splitVerificationBehavior == SplitVerificationBehavior.VERIFY_SUCCESS) {
      Assert.assertNotEquals(0, residualElements.size());
    }
  }

  @Test
  public void testRequestDynamicSplit() throws Exception {
    // Note that exhaustive tests for AvroSource's split behavior exist in {@link AvroSourceTest}.
    List<List<String>> elements = generateInputBlocks(10, 100 * 100, 100);
    Coder<String> coder = StringUtf8Coder.of();
    AvroFileInfo<String> fileInfo = initInputFile(elements, coder);
    AvroByteReader<String> reader =
        new AvroByteReader<String>(
            fileInfo.filename, 0L, Long.MAX_VALUE, coder, PipelineOptionsFactory.create());
    // Read most of the records before the proposed split point.
    testRequestDynamicSplitInternal(reader, 0.5F, 490, SplitVerificationBehavior.VERIFY_SUCCESS);
    // Read a single record.
    testRequestDynamicSplitInternal(reader, 0.5F, 1, SplitVerificationBehavior.VERIFY_SUCCESS);
    // Read zero records.
    testRequestDynamicSplitInternal(reader, 0.5F, 0, SplitVerificationBehavior.VERIFY_FAILURE);
    // Read almost the entire input.
    testRequestDynamicSplitInternal(reader, 0.5F, 900, SplitVerificationBehavior.VERIFY_FAILURE);
    // Read the entire input.
    testRequestDynamicSplitInternal(reader, 0.5F, 2000, SplitVerificationBehavior.VERIFY_FAILURE);
  }

  @Test
  public void testRequestDynamicSplitExhaustive() throws Exception {
    List<List<String>> elements = generateInputBlocks(5, 10 * 10, 10);
    Coder<String> coder = StringUtf8Coder.of();
    AvroFileInfo<String> fileInfo = initInputFile(elements, coder);
    AvroByteReader<String> reader =
        new AvroByteReader<String>(
            fileInfo.filename, 0L, Long.MAX_VALUE, coder, PipelineOptionsFactory.create());
    for (float splitFraction = 0.0F; splitFraction < 1.0F; splitFraction += 0.02F) {
      for (int recordsToRead = 0; recordsToRead <= 500; recordsToRead += 5) {
        testRequestDynamicSplitInternal(
            reader, splitFraction, recordsToRead, SplitVerificationBehavior.DO_NOT_VERIFY);
      }
    }
  }

  @Test
  public void testGetProgress() throws Exception {
    // Ensure that AvroByteReader reports progress from the underlying AvroSource.

    // 4 blocks with 4 split points.
    List<List<String>> elements = generateInputBlocks(4, 10, 10);
    Coder<String> coder = StringUtf8Coder.of();
    AvroFileInfo<String> fileInfo = initInputFile(elements, coder);
    AvroByteReader<String> reader =
        new AvroByteReader<String>(
            fileInfo.filename, 0L, Long.MAX_VALUE, coder, PipelineOptionsFactory.create());

    AvroByteFileIterator iterator = reader.iterator();
    Assert.assertTrue(iterator.start());
    ApproximateReportedProgress progress =
        readerProgressToCloudProgress(reader.iterator().getProgress());
    Assert.assertEquals(0.0, progress.getConsumedParallelism().getValue(), 1e-6);
    Assert.assertEquals(0.0, progress.getFractionConsumed(), 1e-6);
    Assert.assertNull(progress.getRemainingParallelism());

    Assert.assertTrue(iterator.advance());
    progress = readerProgressToCloudProgress(iterator.getProgress());
    Assert.assertEquals(1.0, progress.getConsumedParallelism().getValue(), 1e-6);
    Assert.assertNull(progress.getRemainingParallelism());

    // Advance to the end of last block and check consumed parallelism along the way.
    Assert.assertTrue(iterator.advance());
    progress = readerProgressToCloudProgress(iterator.getProgress());
    Assert.assertEquals(2.0, progress.getConsumedParallelism().getValue(), 1e-6);
    Assert.assertNull(progress.getRemainingParallelism());

    Assert.assertTrue(iterator.advance());
    progress = readerProgressToCloudProgress(iterator.getProgress());
    Assert.assertEquals(3.0, progress.getConsumedParallelism().getValue(), 1e-6);
    Assert.assertEquals(1.0, progress.getRemainingParallelism().getValue(), 1e-6);

    Assert.assertFalse(iterator.advance());
    progress = readerProgressToCloudProgress(iterator.getProgress());
    Assert.assertEquals(4.0, progress.getConsumedParallelism().getValue(), 1e-6);
    Assert.assertEquals(0.0, progress.getRemainingParallelism().getValue(), 1e-6);
    Assert.assertEquals(1.0, progress.getFractionConsumed(), 1e-6);
  }

  // TODO: sharded filenames
  // TODO: reading from GCS
}
