/*
 * 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.iceberg.flink.source.reader;

import static org.assertj.core.api.Assertions.assertThat;

import java.io.File;
import java.io.IOException;
import java.nio.file.Path;
import java.util.Arrays;
import java.util.List;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.connector.base.source.reader.RecordsWithSplitIds;
import org.apache.flink.connector.base.source.reader.SourceReaderOptions;
import org.apache.flink.table.data.RowData;
import org.apache.iceberg.BaseCombinedScanTask;
import org.apache.iceberg.CombinedScanTask;
import org.apache.iceberg.FileFormat;
import org.apache.iceberg.FileScanTask;
import org.apache.iceberg.data.GenericRecord;
import org.apache.iceberg.data.RandomGenericData;
import org.apache.iceberg.data.Record;
import org.apache.iceberg.expressions.Expression;
import org.apache.iceberg.expressions.Expressions;
import org.apache.iceberg.expressions.ResidualEvaluator;
import org.apache.iceberg.flink.FlinkConfigOptions;
import org.apache.iceberg.flink.TestFixtures;
import org.apache.iceberg.flink.TestHelpers;
import org.apache.iceberg.flink.source.DataIterator;
import org.apache.iceberg.io.CloseableIterator;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;

public class TestArrayPoolDataIteratorBatcherRowData {

  @TempDir protected Path temporaryFolder;
  private static final FileFormat FILE_FORMAT = FileFormat.PARQUET;
  private final Configuration config =
      new Configuration()
          .set(SourceReaderOptions.ELEMENT_QUEUE_CAPACITY, 1)
          .set(FlinkConfigOptions.SOURCE_READER_FETCH_BATCH_RECORD_COUNT, 2);

  private final DataIteratorBatcher<RowData> batcher =
      new ArrayPoolDataIteratorBatcher<>(config, new RowDataRecordFactory(TestFixtures.ROW_TYPE));

  /** Read a CombinedScanTask that contains a single file with less than a full batch of records */
  @Test
  public void testSingleFileLessThanOneFullBatch() throws Exception {
    List<Record> records = RandomGenericData.generate(TestFixtures.SCHEMA, 1, 1);
    FileScanTask fileTask =
        ReaderUtil.createFileTask(
            records,
            File.createTempFile("junit", null, temporaryFolder.toFile()),
            FILE_FORMAT,
            TestFixtures.SCHEMA);
    CombinedScanTask combinedTask = new BaseCombinedScanTask(fileTask);
    DataIterator<RowData> dataIterator = ReaderUtil.createDataIterator(combinedTask);
    String splitId = "someSplitId";
    CloseableIterator<RecordsWithSplitIds<RecordAndPosition<RowData>>> recordBatchIterator =
        batcher.batch(splitId, dataIterator);

    ArrayBatchRecords<RowData> batch = (ArrayBatchRecords<RowData>) recordBatchIterator.next();
    assertThat(batch.finishedSplits()).isEmpty();
    assertThat(batch.nextSplit()).isEqualTo(splitId);
    assertThat(batch.records()).hasSize(2);
    assertThat(batch.numberOfRecords()).isEqualTo(1);

    RecordAndPosition<RowData> recordAndPosition = batch.nextRecordFromSplit();

    ///////////////////////////////
    // assert first record

    assertThat(recordAndPosition.fileOffset()).isEqualTo(0);
    assertThat(recordAndPosition.recordOffset())
        .as("The position points to where the reader should resume after this record is processed.")
        .isEqualTo(1);
    TestHelpers.assertRowData(TestFixtures.SCHEMA, records.get(0), recordAndPosition.record());

    assertThat(batch.nextRecordFromSplit()).isNull();
    assertThat(batch.nextSplit()).isNull();
    batch.recycle();

    assertThat(recordBatchIterator).isExhausted();
  }

  /**
   * Read a CombinedScanTask that contains a single file with multiple batches.
   *
   * <p>Insert 5 records in a single file that should result in 3 batches
   */
  @Test
  public void testSingleFileWithMultipleBatches() throws Exception {
    List<Record> records = RandomGenericData.generate(TestFixtures.SCHEMA, 5, 1);
    FileScanTask fileTask =
        ReaderUtil.createFileTask(
            records,
            File.createTempFile("junit", null, temporaryFolder.toFile()),
            FILE_FORMAT,
            TestFixtures.SCHEMA);
    CombinedScanTask combinedTask = new BaseCombinedScanTask(fileTask);
    DataIterator<RowData> dataIterator = ReaderUtil.createDataIterator(combinedTask);
    String splitId = "someSplitId";
    CloseableIterator<RecordsWithSplitIds<RecordAndPosition<RowData>>> recordBatchIterator =
        batcher.batch(splitId, dataIterator);

    ///////////////////////////////
    // assert first batch with full batch of 2 records

    ArrayBatchRecords<RowData> batch0 = (ArrayBatchRecords<RowData>) recordBatchIterator.next();
    assertThat(batch0.finishedSplits()).isEmpty();
    assertThat(batch0.nextSplit()).isEqualTo(splitId);
    assertThat(batch0.records()).hasSize(2);
    assertThat(batch0.numberOfRecords()).isEqualTo(2);

    RecordAndPosition<RowData> recordAndPosition;

    // assert first record
    recordAndPosition = batch0.nextRecordFromSplit();
    assertThat(recordAndPosition.fileOffset()).isEqualTo(0);
    assertThat(recordAndPosition.recordOffset())
        .as("The position points to where the reader should resume after this record is processed.")
        .isEqualTo(1);
    TestHelpers.assertRowData(TestFixtures.SCHEMA, records.get(0), recordAndPosition.record());

    // assert second record
    recordAndPosition = batch0.nextRecordFromSplit();
    assertThat(recordAndPosition.fileOffset()).isEqualTo(0);
    assertThat(recordAndPosition.recordOffset())
        .as("The position points to where the reader should resume after this record is processed.")
        .isEqualTo(2);
    TestHelpers.assertRowData(TestFixtures.SCHEMA, records.get(1), recordAndPosition.record());

    assertThat(batch0.nextRecordFromSplit()).isNull();
    assertThat(batch0.nextSplit()).isNull();
    batch0.recycle();

    ///////////////////////////////
    // assert second batch with full batch of 2 records

    ArrayBatchRecords<RowData> batch1 = (ArrayBatchRecords<RowData>) recordBatchIterator.next();
    assertThat(batch1.records()).containsExactlyInAnyOrder(batch0.records());
    assertThat(batch1.finishedSplits()).isEmpty();
    assertThat(batch1.nextSplit()).isEqualTo(splitId);
    assertThat(batch1.records()).hasSize(2);
    assertThat(batch1.numberOfRecords()).isEqualTo(2);

    // assert third record
    recordAndPosition = batch1.nextRecordFromSplit();
    assertThat(recordAndPosition.fileOffset()).isEqualTo(0);
    assertThat(recordAndPosition.recordOffset())
        .as("The position points to where the reader should resume after this record is processed.")
        .isEqualTo(3);
    TestHelpers.assertRowData(TestFixtures.SCHEMA, records.get(2), recordAndPosition.record());

    // assert fourth record
    recordAndPosition = batch1.nextRecordFromSplit();
    assertThat(recordAndPosition.fileOffset()).isEqualTo(0);
    assertThat(recordAndPosition.recordOffset())
        .as("The position points to where the reader should resume after this record is processed.")
        .isEqualTo(4);
    TestHelpers.assertRowData(TestFixtures.SCHEMA, records.get(3), recordAndPosition.record());

    assertThat(batch1.nextRecordFromSplit()).isNull();
    assertThat(batch1.nextSplit()).isNull();
    batch1.recycle();

    ///////////////////////////////
    // assert third batch with partial batch of 1 record

    ArrayBatchRecords<RowData> batch2 = (ArrayBatchRecords<RowData>) recordBatchIterator.next();
    assertThat(batch2.records()).containsExactlyInAnyOrder(batch0.records());
    assertThat(batch2.finishedSplits()).isEmpty();
    assertThat(batch2.nextSplit()).isEqualTo(splitId);
    assertThat(batch2.records()).hasSize(2);
    assertThat(batch2.numberOfRecords()).isEqualTo(1);

    // assert fifth record
    recordAndPosition = batch2.nextRecordFromSplit();
    assertThat(recordAndPosition.fileOffset()).isEqualTo(0);
    assertThat(recordAndPosition.recordOffset())
        .as("The position points to where the reader should resume after this record is processed.")
        .isEqualTo(5);
    TestHelpers.assertRowData(TestFixtures.SCHEMA, records.get(4), recordAndPosition.record());

    assertThat(batch2.nextRecordFromSplit()).isNull();
    assertThat(batch2.nextSplit()).isNull();
    batch2.recycle();

    assertThat(recordBatchIterator).isExhausted();
  }

  /**
   * Read a CombinedScanTask that contains with multiple files.
   *
   * <p>In this test, we also seek the iterator to starting position (1, 1).
   */
  @Test
  public void testMultipleFilesWithSeekPosition() throws Exception {
    List<Record> records0 = RandomGenericData.generate(TestFixtures.SCHEMA, 1, 1);
    FileScanTask fileTask0 =
        ReaderUtil.createFileTask(
            records0,
            File.createTempFile("junit", null, temporaryFolder.toFile()),
            FILE_FORMAT,
            TestFixtures.SCHEMA);
    List<Record> records1 = RandomGenericData.generate(TestFixtures.SCHEMA, 4, 2);
    FileScanTask fileTask1 =
        ReaderUtil.createFileTask(
            records1,
            File.createTempFile("junit", null, temporaryFolder.toFile()),
            FILE_FORMAT,
            TestFixtures.SCHEMA);
    List<Record> records2 = RandomGenericData.generate(TestFixtures.SCHEMA, 3, 3);
    FileScanTask fileTask2 =
        ReaderUtil.createFileTask(
            records2,
            File.createTempFile("junit", null, temporaryFolder.toFile()),
            FILE_FORMAT,
            TestFixtures.SCHEMA);
    CombinedScanTask combinedTask =
        new BaseCombinedScanTask(Arrays.asList(fileTask0, fileTask1, fileTask2));

    DataIterator<RowData> dataIterator = ReaderUtil.createDataIterator(combinedTask);
    dataIterator.seek(1, 1);

    String splitId = "someSplitId";
    CloseableIterator<RecordsWithSplitIds<RecordAndPosition<RowData>>> recordBatchIterator =
        batcher.batch(splitId, dataIterator);

    ///////////////////////////////
    // file0 is skipped by seek

    ///////////////////////////////
    // file1 has 4 records. because the seek position, first record is skipped.
    // we should read 3 remaining records in 2 batches:
    // batch10 with 2 records and batch11 with 1 records.

    // assert first batch from file1 with full batch of 2 records

    // variable naming convention: batch<fileOffset><batchId>
    ArrayBatchRecords<RowData> batch10 = (ArrayBatchRecords<RowData>) recordBatchIterator.next();
    assertThat(batch10.finishedSplits()).isEmpty();
    assertThat(batch10.nextSplit()).isEqualTo(splitId);
    assertThat(batch10.records()).hasSize(2);
    assertThat(batch10.numberOfRecords()).isEqualTo(2);

    RecordAndPosition<RowData> recordAndPosition;

    recordAndPosition = batch10.nextRecordFromSplit();
    assertThat(recordAndPosition.fileOffset()).isEqualTo(1);
    assertThat(recordAndPosition.recordOffset())
        .as("seek should skip the first record in file1. starting from the second record")
        .isEqualTo(2);
    TestHelpers.assertRowData(TestFixtures.SCHEMA, records1.get(1), recordAndPosition.record());

    recordAndPosition = batch10.nextRecordFromSplit();
    assertThat(recordAndPosition.fileOffset()).isEqualTo(1);
    assertThat(recordAndPosition.recordOffset())
        .as("The position points to where the reader should resume after this record is processed.")
        .isEqualTo(3);
    TestHelpers.assertRowData(TestFixtures.SCHEMA, records1.get(2), recordAndPosition.record());

    assertThat(batch10.nextRecordFromSplit()).isNull();
    assertThat(batch10.nextSplit()).isNull();
    batch10.recycle();

    // assert second batch from file1 with partial batch of 1 record

    // variable naming convention: batch_<fileOffset>_<batchId>
    ArrayBatchRecords<RowData> batch11 = (ArrayBatchRecords<RowData>) recordBatchIterator.next();
    assertThat(batch11.records()).containsExactlyInAnyOrder(batch10.records());
    assertThat(batch11.finishedSplits()).isEmpty();
    assertThat(batch11.nextSplit()).isEqualTo(splitId);
    assertThat(batch11.records()).hasSize(2);
    assertThat(batch11.numberOfRecords()).isEqualTo(1);

    recordAndPosition = batch11.nextRecordFromSplit();
    assertThat(recordAndPosition.fileOffset()).isEqualTo(1);
    assertThat(recordAndPosition.recordOffset())
        .as("The position points to where the reader should resume after this record is processed.")
        .isEqualTo(4);
    TestHelpers.assertRowData(TestFixtures.SCHEMA, records1.get(3), recordAndPosition.record());

    assertThat(batch11.nextRecordFromSplit()).isNull();
    assertThat(batch11.nextSplit()).isNull();
    batch11.recycle();

    ///////////////////////////////
    // file2 has 3 records.
    // we should read 3 records in 2 batches:
    // batch20 with 2 records and batch21 with 1 records

    // assert first batch from file2 with full batch of 2 records

    // variable naming convention: batch_<fileOffset>_<batchId>
    ArrayBatchRecords<RowData> batch20 = (ArrayBatchRecords<RowData>) recordBatchIterator.next();
    assertThat(batch20.records()).containsExactlyInAnyOrder(batch10.records());
    assertThat(batch20.finishedSplits()).isEmpty();
    assertThat(batch20.nextSplit()).isEqualTo(splitId);
    assertThat(batch20.records()).hasSize(2);
    assertThat(batch20.numberOfRecords()).isEqualTo(2);

    recordAndPosition = batch20.nextRecordFromSplit();
    assertThat(recordAndPosition.fileOffset()).isEqualTo(2);
    assertThat(recordAndPosition.recordOffset())
        .as("The position points to where the reader should resume after this record is processed.")
        .isEqualTo(1);
    TestHelpers.assertRowData(TestFixtures.SCHEMA, records2.get(0), recordAndPosition.record());

    recordAndPosition = batch20.nextRecordFromSplit();
    assertThat(recordAndPosition.fileOffset()).isEqualTo(2);
    assertThat(recordAndPosition.recordOffset())
        .as("The position points to where the reader should resume after this record is processed.")
        .isEqualTo(2);
    TestHelpers.assertRowData(TestFixtures.SCHEMA, records2.get(1), recordAndPosition.record());

    assertThat(batch20.nextRecordFromSplit()).isNull();
    assertThat(batch20.nextSplit()).isNull();
    batch20.recycle();

    ///////////////////////////////
    // assert second batch from file2 with partial batch of 1 record

    // variable naming convention: batch_<fileOffset>_<batchId>
    ArrayBatchRecords<RowData> batch21 = (ArrayBatchRecords<RowData>) recordBatchIterator.next();
    assertThat(batch21.records()).containsExactlyInAnyOrder(batch10.records());
    assertThat(batch21.finishedSplits()).isEmpty();
    assertThat(batch21.nextSplit()).isEqualTo(splitId);
    assertThat(batch21.records()).hasSize(2);
    assertThat(batch21.numberOfRecords()).isEqualTo(1);

    recordAndPosition = batch21.nextRecordFromSplit();
    assertThat(recordAndPosition.fileOffset()).isEqualTo(2);

    assertThat(recordAndPosition.recordOffset())
        .as("The position points to where the reader should resume after this record is processed.")
        .isEqualTo(3);
    TestHelpers.assertRowData(TestFixtures.SCHEMA, records2.get(2), recordAndPosition.record());

    assertThat(batch21.nextRecordFromSplit()).isNull();
    assertThat(batch21.nextSplit()).isNull();
    batch21.recycle();

    assertThat(recordBatchIterator).isExhausted();
  }

  @Test
  public void testDataIteratorWithResidualFilter() throws IOException {
    GenericRecord record0 = GenericRecord.create(TestFixtures.SCHEMA);
    record0.setField("data", "file-1");
    record0.setField("id", 0L);
    record0.setField("dt", "-");

    GenericRecord record1 = GenericRecord.create(TestFixtures.SCHEMA);
    record1.setField("data", "file-2");
    record1.setField("id", 1L);
    record1.setField("dt", "-");

    List<Record> fileRecords0 = ImmutableList.of(record0);
    List<Record> fileRecords1 = ImmutableList.of(record1);

    validateFileOffsetWithResidualFilter(fileRecords0, fileRecords1, Expressions.alwaysTrue());

    validateFileOffsetWithResidualFilter(
        fileRecords0, fileRecords1, Expressions.greaterThan("id", 0));
  }

  private void validateFileOffsetWithResidualFilter(
      List<Record> fileRecords0, List<Record> fileRecords1, Expression residualFilter)
      throws IOException {
    ResidualEvaluator residualEvaluator = ResidualEvaluator.unpartitioned(residualFilter);

    FileScanTask fileTask0 =
        ReaderUtil.createFileTask(
            fileRecords0,
            File.createTempFile("junit", null, temporaryFolder.toFile()),
            FILE_FORMAT,
            TestFixtures.SCHEMA,
            residualEvaluator);

    FileScanTask fileTask1 =
        ReaderUtil.createFileTask(
            fileRecords1,
            File.createTempFile("junit", null, temporaryFolder.toFile()),
            FILE_FORMAT,
            TestFixtures.SCHEMA,
            residualEvaluator);

    CombinedScanTask combinedTask = new BaseCombinedScanTask(Arrays.asList(fileTask0, fileTask1));

    DataIterator<RowData> dataIterator = ReaderUtil.createDataIterator(combinedTask);

    dataIterator.seek(0, 0);

    while (dataIterator.hasNext()) {
      assertThat(dataIterator.fileOffset()).isEqualTo(dataIterator.next().getLong(1));
    }
  }

  @Test
  public void testInitializationWithHeadFilesSkipped() throws IOException {
    GenericRecord record0 = GenericRecord.create(TestFixtures.SCHEMA);
    record0.setField("data", "a");
    record0.setField("id", 1L);
    record0.setField("dt", "-");

    GenericRecord record1 = GenericRecord.create(TestFixtures.SCHEMA);
    record1.setField("data", "a");
    record1.setField("id", 10L);
    record1.setField("dt", "-");

    ResidualEvaluator residuals = ResidualEvaluator.unpartitioned(Expressions.greaterThan("id", 5));

    FileScanTask fileTask0 =
        ReaderUtil.createFileTask(
            ImmutableList.of(record0),
            File.createTempFile("junit", null, temporaryFolder.toFile()),
            FILE_FORMAT,
            TestFixtures.SCHEMA,
            residuals);

    FileScanTask fileTask1 =
        ReaderUtil.createFileTask(
            ImmutableList.of(record1),
            File.createTempFile("junit", null, temporaryFolder.toFile()),
            FILE_FORMAT,
            TestFixtures.SCHEMA,
            residuals);

    CombinedScanTask combinedTask = new BaseCombinedScanTask(Arrays.asList(fileTask0, fileTask1));

    DataIterator<RowData> dataIterator = ReaderUtil.createDataIterator(combinedTask);

    dataIterator.seek(0, 0);

    assertThat(dataIterator.fileOffset())
        .as("File offset should be 1 because file 0 should be skipped")
        .isEqualTo(1);

    TestHelpers.assertRowData(TestFixtures.SCHEMA, record1, dataIterator.next());
  }

  @Test
  void testSeekResumesCorrectlyAfterHeadFilesSkipped() throws IOException {
    // Setup: [F0(filtered), F1(1 record), F2(2 records)]
    // F1 has only 1 record so that after reading F1, the checkpoint position (fileOffset=1,
    // recordOffset=1) falls at a file boundary. Without the fix, fileOffset is incorrectly
    // reported as 0 instead of 1 after seek(0,0) skips F0. This causes F2's first record to
    // be reported at (fileOffset=1, recordOffset=1). On restore with seek(1, 1), the iterator
    // skips only F0 + 1 record from F1, landing at the start of F2 instead of after F2 record0,
    // resulting in F2 record0 being read again (duplicate).
    GenericRecord filteredRecord = GenericRecord.create(TestFixtures.SCHEMA);
    filteredRecord.setField("data", "a");
    filteredRecord.setField("id", 1L);
    filteredRecord.setField("dt", "-");

    GenericRecord f1Record0 = GenericRecord.create(TestFixtures.SCHEMA);
    f1Record0.setField("data", "b");
    f1Record0.setField("id", 10L);
    f1Record0.setField("dt", "-");

    GenericRecord f2Record0 = GenericRecord.create(TestFixtures.SCHEMA);
    f2Record0.setField("data", "d");
    f2Record0.setField("id", 20L);
    f2Record0.setField("dt", "-");

    GenericRecord f2Record1 = GenericRecord.create(TestFixtures.SCHEMA);
    f2Record1.setField("data", "e");
    f2Record1.setField("id", 21L);
    f2Record1.setField("dt", "-");

    ResidualEvaluator residuals = ResidualEvaluator.unpartitioned(Expressions.greaterThan("id", 5));

    FileScanTask fileTask0 =
        ReaderUtil.createFileTask(
            ImmutableList.of(filteredRecord),
            File.createTempFile("junit", null, temporaryFolder.toFile()),
            FILE_FORMAT,
            TestFixtures.SCHEMA,
            residuals);

    FileScanTask fileTask1 =
        ReaderUtil.createFileTask(
            ImmutableList.of(f1Record0),
            File.createTempFile("junit", null, temporaryFolder.toFile()),
            FILE_FORMAT,
            TestFixtures.SCHEMA,
            residuals);

    FileScanTask fileTask2 =
        ReaderUtil.createFileTask(
            ImmutableList.of(f2Record0, f2Record1),
            File.createTempFile("junit", null, temporaryFolder.toFile()),
            FILE_FORMAT,
            TestFixtures.SCHEMA,
            residuals);

    CombinedScanTask combinedTask =
        new BaseCombinedScanTask(Arrays.asList(fileTask0, fileTask1, fileTask2));

    // First read: consume file 1's record and file 2's first record
    DataIterator<RowData> dataIterator = ReaderUtil.createDataIterator(combinedTask);
    dataIterator.seek(0, 0);

    TestHelpers.assertRowData(TestFixtures.SCHEMA, f1Record0, dataIterator.next());
    // F1 exhausted, updateCurrentIterator moves to F2
    TestHelpers.assertRowData(TestFixtures.SCHEMA, f2Record0, dataIterator.next());

    // Capture checkpoint position after reading F2 record0.
    // With the fix: fileOffset=2, recordOffset=1 (correctly points into F2)
    // Without fix: fileOffset=1, recordOffset=1 (incorrectly points to "F1 record 1")
    int checkpointFileOffset = dataIterator.fileOffset();
    long checkpointRecordOffset = dataIterator.recordOffset();

    // Simulate restore: seek to the checkpoint position
    DataIterator<RowData> restoredIterator = ReaderUtil.createDataIterator(combinedTask);
    restoredIterator.seek(checkpointFileOffset, checkpointRecordOffset);

    // After restore, should read only F2 record1 (the unread remainder)
    assertThat(restoredIterator.hasNext()).isTrue();
    TestHelpers.assertRowData(TestFixtures.SCHEMA, f2Record1, restoredIterator.next());
    assertThat(restoredIterator.hasNext()).isFalse();
  }
}
