/*
 * 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.io.iceberg;

import static java.nio.charset.StandardCharsets.UTF_8;
import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.hamcrest.Matchers.containsInAnyOrder;
import static org.hamcrest.Matchers.containsString;
import static org.hamcrest.Matchers.either;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;

import java.io.IOException;
import java.net.URLEncoder;
import java.nio.ByteBuffer;
import java.time.Duration;
import java.time.Instant;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.LocalTime;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.sdk.values.WindowedValue;
import org.apache.beam.sdk.values.WindowedValues;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
import org.apache.commons.lang3.RandomStringUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.iceberg.AppendFiles;
import org.apache.iceberg.DataFile;
import org.apache.iceberg.FileFormat;
import org.apache.iceberg.NullOrder;
import org.apache.iceberg.PartitionField;
import org.apache.iceberg.PartitionKey;
import org.apache.iceberg.PartitionSpec;
import org.apache.iceberg.Snapshot;
import org.apache.iceberg.SortDirection;
import org.apache.iceberg.SortOrder;
import org.apache.iceberg.Table;
import org.apache.iceberg.catalog.Catalog;
import org.apache.iceberg.catalog.Namespace;
import org.apache.iceberg.catalog.TableIdentifier;
import org.apache.iceberg.data.Record;
import org.apache.iceberg.hadoop.HadoopCatalog;
import org.apache.iceberg.io.FileIO;
import org.apache.iceberg.io.InputFile;
import org.apache.iceberg.io.OutputFile;
import org.apache.iceberg.io.PositionOutputStream;
import org.apache.iceberg.transforms.Transform;
import org.apache.iceberg.types.Conversions;
import org.apache.iceberg.types.Type;
import org.apache.iceberg.types.Types;
import org.apache.iceberg.util.DateTimeUtil;
import org.checkerframework.checker.nullness.qual.Nullable;
import org.joda.time.DateTime;
import org.joda.time.DateTimeZone;
import org.joda.time.ReadableDateTime;
import org.junit.Before;
import org.junit.ClassRule;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
import org.junit.rules.TemporaryFolder;
import org.junit.rules.TestName;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.Mockito;

/** Test class for {@link RecordWriterManager}. */
@RunWith(JUnit4.class)
public class RecordWriterManagerTest {
  @ClassRule public static final TemporaryFolder TEMPORARY_FOLDER = new TemporaryFolder();

  @Rule
  public transient TestDataWarehouse warehouse = new TestDataWarehouse(TEMPORARY_FOLDER, "default");

  @Rule public TestName testName = new TestName();
  private static final Schema BEAM_SCHEMA =
      Schema.builder().addInt32Field("id").addStringField("name").addBooleanField("bool").build();
  private static final org.apache.iceberg.Schema ICEBERG_SCHEMA =
      IcebergUtils.beamSchemaToIcebergSchema(BEAM_SCHEMA);
  private static final PartitionSpec PARTITION_SPEC =
      PartitionSpec.builderFor(ICEBERG_SCHEMA).truncate("name", 3).identity("bool").build();

  private WindowedValue<IcebergDestination> windowedDestination;
  private HadoopCatalog catalog;

  @Before
  public void setUp() {
    windowedDestination =
        getWindowedDestination("table_" + testName.getMethodName(), PARTITION_SPEC);
    catalog = new HadoopCatalog(new Configuration(), warehouse.location);
    RecordWriterManager.LAST_REFRESHED_TABLE_CACHE.invalidateAll();
  }

  private WindowedValue<IcebergDestination> getWindowedDestination(
      String tableName, @Nullable PartitionSpec partitionSpec) {
    return getWindowedDestination(tableName, ICEBERG_SCHEMA, partitionSpec);
  }

  private WindowedValue<IcebergDestination> getWindowedDestination(
      String tableName, org.apache.iceberg.Schema schema, @Nullable PartitionSpec partitionSpec) {
    TableIdentifier tableIdentifier = TableIdentifier.of("default", tableName);

    // TODO: remove when we enable dynamic table creation with partition specs
    if (partitionSpec != null) {
      warehouse.createTable(tableIdentifier, schema, partitionSpec);
    }

    IcebergDestination icebergDestination =
        IcebergDestination.builder()
            .setFileFormat(FileFormat.PARQUET)
            .setTableIdentifier(tableIdentifier)
            .build();
    return WindowedValues.valueInGlobalWindow(icebergDestination);
  }

  @Test
  public void testCreateNamespaceAndTable() {
    RecordWriterManager writerManager = new RecordWriterManager(catalog, "test_file_name", 1000, 3);
    Namespace newNamespace = Namespace.of("new_namespace");
    TableIdentifier identifier = TableIdentifier.of(newNamespace, testName.getMethodName());
    WindowedValue<IcebergDestination> dest =
        WindowedValues.valueInGlobalWindow(
            IcebergDestination.builder()
                .setFileFormat(FileFormat.PARQUET)
                .setTableIdentifier(identifier)
                .build());

    Row row = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();

    assertFalse(catalog.namespaceExists(newNamespace));
    boolean writeSuccess = writerManager.write(dest, row);
    assertTrue(writeSuccess);
    assertTrue(catalog.namespaceExists(newNamespace));
  }

  @Test
  public void testCreateTableWithSortOrder() throws IOException {
    RecordWriterManager writerManager = new RecordWriterManager(catalog, "test_file_name", 1000, 3);
    TableIdentifier identifier = TableIdentifier.of("default", testName.getMethodName());
    WindowedValue<IcebergDestination> dest =
        WindowedValues.valueInGlobalWindow(
            IcebergDestination.builder()
                .setFileFormat(FileFormat.PARQUET)
                .setTableIdentifier(identifier)
                .setTableCreateConfig(
                    IcebergTableCreateConfig.builder()
                        .setSchema(BEAM_SCHEMA)
                        .setPartitionFields(null)
                        .setSortFields(java.util.Arrays.asList("id desc", "name asc nulls last"))
                        .build())
                .build());

    Row row = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();
    assertTrue(writerManager.write(dest, row));
    writerManager.close();

    Table created = catalog.loadTable(identifier);
    SortOrder order = created.sortOrder();
    assertEquals(2, order.fields().size());
    assertEquals(SortDirection.DESC, order.fields().get(0).direction());
    assertEquals(NullOrder.NULLS_LAST, order.fields().get(0).nullOrder());
    assertEquals(SortDirection.ASC, order.fields().get(1).direction());
    assertEquals(NullOrder.NULLS_LAST, order.fields().get(1).nullOrder());
  }

  @Test
  public void testCreateNewWriterForEachDestination() throws IOException {
    // Writer manager with a maximum limit of 3 writers
    RecordWriterManager writerManager = new RecordWriterManager(catalog, "test_file_name", 1000, 3);
    assertEquals(0, writerManager.openWriters);

    boolean writeSuccess;

    WindowedValue<IcebergDestination> dest1 = getWindowedDestination("dest1", null);
    WindowedValue<IcebergDestination> dest2 = getWindowedDestination("dest2", null);
    WindowedValue<IcebergDestination> dest3 = getWindowedDestination("dest3", PARTITION_SPEC);
    WindowedValue<IcebergDestination> dest4 = getWindowedDestination("dest4", null);

    // dest1
    // This is a new destination so a new writer will be created.
    Row row = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();
    writeSuccess = writerManager.write(dest1, row);
    assertTrue(writeSuccess);
    assertEquals(1, writerManager.openWriters);

    // dest2
    // This is a new destination so a new writer will be created.
    row = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();
    writeSuccess = writerManager.write(dest2, row);
    assertTrue(writeSuccess);
    assertEquals(2, writerManager.openWriters);

    // dest3, partition: [aaa, true]
    // This is a new destination so a new writer will be created.
    row = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();
    writeSuccess = writerManager.write(dest3, row);
    assertTrue(writeSuccess);
    assertEquals(3, writerManager.openWriters);

    // dest4
    // This is a new destination, but the writer manager is saturated with 3
    // writers. reject the
    // record
    row = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();
    writeSuccess = writerManager.write(dest4, row);
    assertFalse(writeSuccess);
    assertEquals(3, writerManager.openWriters);

    // dest3, partition: [aaa, false]
    // new partition, but the writer manager is saturated with 3 writers. reject the
    // record
    row = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", false).build();
    writeSuccess = writerManager.write(dest3, row);
    assertFalse(writeSuccess);
    assertEquals(3, writerManager.openWriters);

    // Closing PartitionRecordWriter will close all writers.
    writerManager.close();
    assertEquals(0, writerManager.openWriters);

    // We should only have 3 data files (one for each destination we wrote to)
    assertEquals(3, writerManager.getSerializableDataFiles().keySet().size());
    assertThat(
        writerManager.getSerializableDataFiles().keySet(), containsInAnyOrder(dest1, dest2, dest3));
  }

  @Test
  public void testCreateNewWriterForEachPartition() throws IOException {
    // Writer manager with a maximum limit of 3 writers
    RecordWriterManager writerManager = new RecordWriterManager(catalog, "test_file_name", 1000, 3);
    assertEquals(0, writerManager.openWriters);

    boolean writeSuccess;

    // partition: [aaa, true].
    // This is a new partition so a new writer will be created.
    Row row = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();
    writeSuccess = writerManager.write(windowedDestination, row);
    assertTrue(writeSuccess);
    assertEquals(1, writerManager.openWriters);

    // partition: [bbb, false].
    // This is a new partition so a new writer will be created.
    row = Row.withSchema(BEAM_SCHEMA).addValues(2, "bbb", false).build();
    writeSuccess = writerManager.write(windowedDestination, row);
    assertTrue(writeSuccess);
    assertEquals(2, writerManager.openWriters);

    // partition: [bbb, false].
    // A writer already exists for this partition, so no new writers are created.
    row = Row.withSchema(BEAM_SCHEMA).addValues(3, "bbbaaa", false).build();
    writeSuccess = writerManager.write(windowedDestination, row);
    assertTrue(writeSuccess);
    assertEquals(2, writerManager.openWriters);

    // partition: [bbb, true].
    // This is a new partition so a new writer will be created.
    row = Row.withSchema(BEAM_SCHEMA).addValues(4, "bbb123", true).build();
    writeSuccess = writerManager.write(windowedDestination, row);
    assertTrue(writeSuccess);
    assertEquals(3, writerManager.openWriters);

    // partition: [aaa, false].
    // The writerManager is already saturated with three writers. This record is
    // rejected.
    row = Row.withSchema(BEAM_SCHEMA).addValues(5, "aaa123", false).build();
    writeSuccess = writerManager.write(windowedDestination, row);
    assertFalse(writeSuccess);
    assertEquals(3, writerManager.openWriters);

    // Closing RecordWriterManager will close all writers.
    writerManager.close();
    assertEquals(0, writerManager.openWriters);

    // We should have only one destination
    assertEquals(1, writerManager.getSerializableDataFiles().size());
    assertTrue(writerManager.getSerializableDataFiles().containsKey(windowedDestination));
    // We should have 3 data files (one for each partition we wrote to)
    assertEquals(3, writerManager.getSerializableDataFiles().get(windowedDestination).size());
    long totalRows = 0;
    for (SerializableDataFile dataFile :
        writerManager.getSerializableDataFiles().get(windowedDestination)) {
      totalRows += dataFile.getRecordCount();
    }
    assertEquals(4L, totalRows);
  }

  @Test
  public void testRespectMaxFileSize() throws IOException {
    // Writer manager with a maximum file size of 100 bytes
    RecordWriterManager writerManager = new RecordWriterManager(catalog, "test_file_name", 100, 2);
    assertEquals(0, writerManager.openWriters);
    boolean writeSuccess;

    PartitionKey partitionKey = new PartitionKey(PARTITION_SPEC, ICEBERG_SCHEMA);
    // row partition:: [aaa, true].
    // This is a new partition so a new writer will be created.
    Row row = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();
    writeSuccess = writerManager.write(windowedDestination, row);
    assertTrue(writeSuccess);
    assertEquals(1, writerManager.openWriters);

    partitionKey.partition(IcebergUtils.beamRowToIcebergRecord(ICEBERG_SCHEMA, row));
    Map<PartitionKey, Integer> writerCounts =
        writerManager.destinations.get(windowedDestination).writerCounts;
    // this is our first writer
    assertEquals(1, writerCounts.get(partitionKey).intValue());

    // row partition:: [aaa, true].
    // existing partition. use existing writer
    row =
        Row.withSchema(BEAM_SCHEMA)
            .addValues(2, "aaa" + RandomStringUtils.randomAlphanumeric(1000), true)
            .build();
    writeSuccess = writerManager.write(windowedDestination, row);
    assertTrue(writeSuccess);
    assertEquals(1, writerManager.openWriters);
    // check that we still use our first writer
    assertEquals(1, writerCounts.get(partitionKey).intValue());

    // row partition:: [aaa, true].
    // writer has reached max file size. create a new writer
    row = Row.withSchema(BEAM_SCHEMA).addValues(2, "aaabb", true).build();
    writeSuccess = writerManager.write(windowedDestination, row);
    assertTrue(writeSuccess);
    // check that we have opened and are using a second writer
    assertEquals(2, writerCounts.get(partitionKey).intValue());
    // check that only one writer is open (we have closed the first writer)
    assertEquals(1, writerManager.openWriters);

    writerManager.close();
    assertEquals(0, writerManager.openWriters);
  }

  @Test
  public void testRequireClosingBeforeFetchingDataFiles() {
    RecordWriterManager writerManager = new RecordWriterManager(catalog, "test_file_name", 100, 2);
    Row row = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();
    writerManager.write(windowedDestination, row);
    assertEquals(1, writerManager.openWriters);

    assertThrows(IllegalStateException.class, writerManager::getSerializableDataFiles);
  }

  /** DataFile doesn't implement a .equals() method. Check equality manually. */
  private static void checkDataFileEquality(DataFile d1, DataFile d2) {
    assertEquals(d1.path().toString(), d2.path().toString());
    assertEquals(d1.format(), d2.format());
    assertEquals(d1.recordCount(), d2.recordCount());
    assertEquals(d1.partition(), d2.partition());
    assertEquals(d1.specId(), d2.specId());
    assertEquals(d1.keyMetadata(), d2.keyMetadata());
    assertEquals(d1.splitOffsets(), d2.splitOffsets());
    assertEquals(d1.columnSizes(), d2.columnSizes());
    assertEquals(d1.valueCounts(), d2.valueCounts());
    assertEquals(d1.nullValueCounts(), d2.nullValueCounts());
    assertEquals(d1.nanValueCounts(), d2.nanValueCounts());
    assertEquals(d1.equalityFieldIds(), d2.equalityFieldIds());
    assertEquals(d1.fileSequenceNumber(), d2.fileSequenceNumber());
    assertEquals(d1.dataSequenceNumber(), d2.dataSequenceNumber());
    assertEquals(d1.pos(), d2.pos());
  }

  @Test
  public void testSerializableDataFileRoundTripEquality() throws IOException {
    PartitionKey partitionKey = new PartitionKey(PARTITION_SPEC, ICEBERG_SCHEMA);

    Row row = Row.withSchema(BEAM_SCHEMA).addValues(1, "abcdef", true).build();
    Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2, "abcxyz", true).build();
    // same partition for both records (name_trunc=abc, bool=true)
    partitionKey.partition(IcebergUtils.beamRowToIcebergRecord(ICEBERG_SCHEMA, row));

    RecordWriter writer =
        new RecordWriter(catalog, windowedDestination.getValue(), "test_file_name", partitionKey);
    writer.write(IcebergUtils.beamRowToIcebergRecord(ICEBERG_SCHEMA, row));
    writer.write(IcebergUtils.beamRowToIcebergRecord(ICEBERG_SCHEMA, row2));

    writer.close();
    DataFile datafile = writer.getDataFile();
    assertEquals(2L, datafile.recordCount());

    Map<String, PartitionField> partitionFieldMap = new HashMap<>();
    for (PartitionField partitionField : PARTITION_SPEC.fields()) {
      partitionFieldMap.put(partitionField.name(), partitionField);
    }

    String partitionPath =
        RecordWriterManager.getPartitionDataPath(partitionKey.toPath(), partitionFieldMap);
    DataFile roundTripDataFile =
        SerializableDataFile.from(datafile, partitionPath)
            .createDataFile(ImmutableMap.of(PARTITION_SPEC.specId(), PARTITION_SPEC));

    checkDataFileEquality(datafile, roundTripDataFile);
  }

  /**
   * Users may update the table's spec while a write pipeline is running. Sometimes, this can happen
   * after converting {@link DataFile} to {@link SerializableDataFile}s. When converting back to
   * {@link DataFile} to commit in the {@link AppendFilesToTables} step, we need to make sure to use
   * the same {@link PartitionSpec} it was originally created with.
   *
   * <p>This test checks that we're preserving the right {@link PartitionSpec} when such an update
   * happens.
   */
  @Test
  public void testRecreateSerializableDataAfterUpdatingPartitionSpec() throws IOException {
    PartitionKey partitionKey = new PartitionKey(PARTITION_SPEC, ICEBERG_SCHEMA);

    Row row = Row.withSchema(BEAM_SCHEMA).addValues(1, "abcdef", true).build();
    Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2, "abcxyz", true).build();
    // same partition for both records (name_trunc=abc, bool=true)
    partitionKey.partition(IcebergUtils.beamRowToIcebergRecord(ICEBERG_SCHEMA, row));

    // write some rows
    RecordWriter writer =
        new RecordWriter(catalog, windowedDestination.getValue(), "test_file_name", partitionKey);
    writer.write(IcebergUtils.beamRowToIcebergRecord(ICEBERG_SCHEMA, row));
    writer.write(IcebergUtils.beamRowToIcebergRecord(ICEBERG_SCHEMA, row2));
    writer.close();

    // fetch data file and its serializable version
    Map<String, PartitionField> partitionFieldMap = new HashMap<>();
    for (PartitionField partitionField : PARTITION_SPEC.fields()) {
      partitionFieldMap.put(partitionField.name(), partitionField);
    }
    String partitionPath =
        RecordWriterManager.getPartitionDataPath(partitionKey.toPath(), partitionFieldMap);
    DataFile datafile = writer.getDataFile();
    SerializableDataFile serializableDataFile = SerializableDataFile.from(datafile, partitionPath);

    assertEquals(2L, datafile.recordCount());
    assertEquals(serializableDataFile.getPartitionSpecId(), datafile.specId());

    // update spec
    Table table = catalog.loadTable(windowedDestination.getValue().getTableIdentifier());
    table.updateSpec().addField("id").removeField("bool").commit();

    Map<Integer, PartitionSpec> updatedSpecs = table.specs();
    DataFile roundTripDataFile = serializableDataFile.createDataFile(updatedSpecs);

    checkDataFileEquality(datafile, roundTripDataFile);
  }

  @Test
  public void testWriterKeepsUpWithUpdatingPartitionSpec() throws IOException {
    Table table = catalog.loadTable(windowedDestination.getValue().getTableIdentifier());
    Row row = Row.withSchema(BEAM_SCHEMA).addValues(1, "abcdef", true).build();
    Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2, "abcxyz", true).build();

    // write some rows
    RecordWriterManager writer =
        new RecordWriterManager(catalog, "test_prefix", Long.MAX_VALUE, Integer.MAX_VALUE);
    writer.write(windowedDestination, row);
    writer.write(windowedDestination, row2);
    writer.close();
    DataFile dataFile =
        writer
            .getSerializableDataFiles()
            .get(windowedDestination)
            .get(0)
            .createDataFile(table.specs());

    // check data file path contains the correct partition components
    assertEquals(2L, dataFile.recordCount());
    assertEquals(dataFile.specId(), PARTITION_SPEC.specId());
    assertThat(dataFile.path().toString(), containsString("name_trunc=abc"));
    assertThat(dataFile.path().toString(), containsString("bool=true"));

    // table is cached
    assertEquals(1, RecordWriterManager.LAST_REFRESHED_TABLE_CACHE.size());

    // update spec
    table.updateSpec().addField("id").removeField("bool").commit();
    // Make the cached table stale to force reloading its metadata.
    RecordWriterManager.LAST_REFRESHED_TABLE_CACHE.getIfPresent(
                windowedDestination.getValue().getTableIdentifier())
            .lastRefreshTime =
        Instant.EPOCH;

    // write a second data file
    // should refresh the table and use the new partition spec
    RecordWriterManager writer2 =
        new RecordWriterManager(catalog, "test_prefix_2", Long.MAX_VALUE, Integer.MAX_VALUE);
    writer2.write(windowedDestination, row);
    writer2.write(windowedDestination, row2);
    writer2.close();

    List<SerializableDataFile> serializableDataFiles =
        writer2.getSerializableDataFiles().get(windowedDestination);
    assertEquals(2, serializableDataFiles.size());
    for (SerializableDataFile serializableDataFile : serializableDataFiles) {
      assertEquals(table.spec().specId(), serializableDataFile.getPartitionSpecId());
      dataFile = serializableDataFile.createDataFile(table.specs());
      assertEquals(1L, dataFile.recordCount());
      assertThat(dataFile.path().toString(), containsString("name_trunc=abc"));
      assertThat(
          dataFile.path().toString(), either(containsString("id=1")).or(containsString("id=2")));
    }
  }

  @Test
  public void testIdentityPartitioning() throws IOException {
    Schema primitiveTypeSchema =
        Schema.builder()
            .addBooleanField("bool")
            .addInt32Field("int")
            .addInt64Field("long")
            .addFloatField("float")
            .addDoubleField("double")
            .addStringField("str")
            .addLogicalTypeField("date", SqlTypes.DATE)
            .addLogicalTypeField("time", SqlTypes.TIME)
            .addLogicalTypeField("datetime", SqlTypes.DATETIME)
            .addDateTimeField("datetime_tz")
            .build();
    String timestamp = "2025-01-21T13:18:20.053";
    LocalDateTime localDateTime = LocalDateTime.parse(timestamp);
    Row row =
        Row.withSchema(primitiveTypeSchema)
            .addValues(
                true,
                1,
                1L,
                1.23f,
                4.56,
                "str",
                localDateTime.toLocalDate(),
                localDateTime.toLocalTime(),
                localDateTime,
                DateTime.parse(timestamp))
            .build();
    org.apache.iceberg.Schema icebergSchema =
        IcebergUtils.beamSchemaToIcebergSchema(primitiveTypeSchema);
    PartitionSpec spec =
        PartitionSpec.builderFor(icebergSchema)
            .identity("bool")
            .identity("int")
            .identity("long")
            .identity("float")
            .identity("double")
            .identity("str")
            .identity("date")
            .identity("time")
            .identity("datetime")
            .identity("datetime_tz")
            .build();
    WindowedValue<IcebergDestination> dest =
        getWindowedDestination("identity_partitioning", icebergSchema, spec);

    RecordWriterManager writer =
        new RecordWriterManager(catalog, "test_prefix", Long.MAX_VALUE, Integer.MAX_VALUE);
    writer.write(dest, row);
    writer.close();
    List<SerializableDataFile> files = writer.getSerializableDataFiles().get(dest);
    assertEquals(1, files.size());
    SerializableDataFile dataFile = files.get(0);
    assertEquals(1, dataFile.getRecordCount());
    // build this string: bool=true/int=1/long=1/float=1.0/double=1.0/str=str
    List<String> expectedPartitions = new ArrayList<>();

    for (PartitionField field : spec.fields()) {
      String name = field.name();
      Type type = spec.schema().findType(name);
      Transform<Object, Object> transform = (Transform<Object, Object>) field.transform();
      String val;
      switch (name) {
        case "date":
          LocalDate localDate = checkStateNotNull(row.getValue(name));
          Integer day = Integer.parseInt(String.valueOf(localDate.toEpochDay()));
          val = transform.toHumanString(type, day);
          break;
        case "time":
          LocalTime localTime = checkStateNotNull(row.getValue(name));
          val = transform.toHumanString(type, localTime.toNanoOfDay() / 1000);
          break;
        case "datetime":
          LocalDateTime ldt = checkStateNotNull(row.getValue(name));
          val = transform.toHumanString(type, DateTimeUtil.microsFromTimestamp(ldt));
          break;
        case "datetime_tz":
          ReadableDateTime dt = checkStateNotNull(row.getDateTime(name));
          val = transform.toHumanString(type, dt.getMillis() * 1000);
          break;
        default:
          val = transform.toHumanString(type, checkStateNotNull(row.getValue(name)));
          break;
      }
      expectedPartitions.add(name + "=" + URLEncoder.encode(val, UTF_8.toString()));
    }
    String expectedPartitionPath = String.join("/", expectedPartitions);
    assertEquals(expectedPartitionPath, dataFile.getPartitionPath());
    assertThat(dataFile.getPath(), containsString(expectedPartitionPath));
  }

  @Test
  public void testBucketPartitioning() throws IOException {
    Schema bucketSchema =
        Schema.builder()
            .addInt32Field("int")
            .addInt64Field("long")
            .addStringField("str")
            .addLogicalTypeField("date", SqlTypes.DATE)
            .addLogicalTypeField("time", SqlTypes.TIME)
            .addLogicalTypeField("datetime", SqlTypes.DATETIME)
            .addDateTimeField("datetime_tz")
            .build();

    String timestamp = "2024-10-08T13:18:20.053";
    LocalDateTime localDateTime = LocalDateTime.parse(timestamp);

    Row row =
        Row.withSchema(bucketSchema)
            .addValues(
                1,
                1L,
                "str",
                localDateTime.toLocalDate(),
                localDateTime.toLocalTime(),
                localDateTime,
                DateTime.parse(timestamp))
            .build();
    org.apache.iceberg.Schema icebergSchema = IcebergUtils.beamSchemaToIcebergSchema(bucketSchema);
    PartitionSpec spec =
        PartitionSpec.builderFor(icebergSchema)
            .bucket("int", 2)
            .bucket("long", 2)
            .bucket("str", 2)
            .bucket("date", 2)
            .bucket("time", 2)
            .bucket("datetime", 2)
            .bucket("datetime_tz", 2)
            .build();
    WindowedValue<IcebergDestination> dest =
        getWindowedDestination("bucket_partitioning", icebergSchema, spec);

    RecordWriterManager writer =
        new RecordWriterManager(catalog, "test_prefix", Long.MAX_VALUE, Integer.MAX_VALUE);
    writer.write(dest, row);
    writer.close();
    List<SerializableDataFile> files = writer.getSerializableDataFiles().get(dest);
    assertEquals(1, files.size());
    SerializableDataFile dataFile = files.get(0);
    assertEquals(1, dataFile.getRecordCount());
    for (Schema.Field field : bucketSchema.getFields()) {
      String expectedPartition = field.getName() + "_bucket";
      assertThat(dataFile.getPartitionPath(), containsString(expectedPartition));
      assertThat(dataFile.getPath(), containsString(expectedPartition));
    }
  }

  @Test
  public void testTimePartitioning() throws IOException {
    Schema timePartitioningSchema =
        Schema.builder()
            .addLogicalTypeField("y_date", SqlTypes.DATE)
            .addLogicalTypeField("y_datetime", SqlTypes.DATETIME)
            .addDateTimeField("y_datetime_tz")
            .addLogicalTypeField("m_date", SqlTypes.DATE)
            .addLogicalTypeField("m_datetime", SqlTypes.DATETIME)
            .addDateTimeField("m_datetime_tz")
            .addLogicalTypeField("d_date", SqlTypes.DATE)
            .addLogicalTypeField("d_datetime", SqlTypes.DATETIME)
            .addDateTimeField("d_datetime_tz")
            .addLogicalTypeField("h_datetime", SqlTypes.DATETIME)
            .addDateTimeField("h_datetime_tz")
            .build();
    org.apache.iceberg.Schema icebergSchema =
        IcebergUtils.beamSchemaToIcebergSchema(timePartitioningSchema);
    PartitionSpec spec =
        PartitionSpec.builderFor(icebergSchema)
            .year("y_date")
            .year("y_datetime")
            .year("y_datetime_tz")
            .month("m_date")
            .month("m_datetime")
            .month("m_datetime_tz")
            .day("d_date")
            .day("d_datetime")
            .day("d_datetime_tz")
            .hour("h_datetime")
            .hour("h_datetime_tz")
            .build();

    WindowedValue<IcebergDestination> dest =
        getWindowedDestination("time_partitioning", icebergSchema, spec);

    String timestamp = "2024-10-08T13:18:20.053";
    LocalDateTime localDateTime = LocalDateTime.parse(timestamp);
    LocalDate localDate = localDateTime.toLocalDate();
    String timestamptz = "2024-10-08T13:18:20.053+03:27";
    DateTime dateTime = DateTime.parse(timestamptz);

    Row row =
        Row.withSchema(timePartitioningSchema)
            .addValues(localDate, localDateTime, dateTime) // year
            .addValues(localDate, localDateTime, dateTime) // month
            .addValues(localDate, localDateTime, dateTime) // day
            .addValues(localDateTime, dateTime) // hour
            .build();

    // write some rows
    RecordWriterManager writer =
        new RecordWriterManager(catalog, "test_prefix", Long.MAX_VALUE, Integer.MAX_VALUE);
    writer.write(dest, row);
    writer.close();
    List<SerializableDataFile> files = writer.getSerializableDataFiles().get(dest);
    assertEquals(1, files.size());
    SerializableDataFile serializableDataFile = files.get(0);
    assertEquals(1, serializableDataFile.getRecordCount());

    int year = localDateTime.getYear();
    int month = localDateTime.getMonthValue();
    int day = localDateTime.getDayOfMonth();
    int hour = localDateTime.getHour();
    List<String> expectedPartitions = new ArrayList<>();
    for (Schema.Field field : timePartitioningSchema.getFields()) {
      String name = field.getName();
      String expected = "";
      if (name.startsWith("y_")) {
        expected = String.format("%s_year=%s", name, year);
      } else if (name.startsWith("m_")) {
        expected = String.format("%s_month=%s-%02d", name, year, month);
      } else if (name.startsWith("d_")) {
        expected = String.format("%s_day=%s-%02d-%02d", name, year, month, day);
      } else if (name.startsWith("h_")) {
        if (name.contains("tz")) {
          hour = dateTime.withZone(DateTimeZone.UTC).getHourOfDay();
        }
        expected = String.format("%s_hour=%s-%02d-%02d-%02d", name, year, month, day, hour);
      }
      expectedPartitions.add(expected);
    }
    String expectedPartition = String.join("/", expectedPartitions);
    DataFile dataFile =
        serializableDataFile.createDataFile(
            catalog.loadTable(dest.getValue().getTableIdentifier()).specs());
    assertThat(dataFile.path().toString(), containsString(expectedPartition));
  }

  @Rule public ExpectedException thrown = ExpectedException.none();

  @Test
  public void testWriterExceptionGetsCaught() throws IOException {
    RecordWriterManager writerManager = new RecordWriterManager(catalog, "test_file_name", 100, 2);
    Row row = Row.withSchema(BEAM_SCHEMA).addValues(1, "abcdef", true).build();
    PartitionKey partitionKey = new PartitionKey(PARTITION_SPEC, ICEBERG_SCHEMA);
    partitionKey.partition(IcebergUtils.beamRowToIcebergRecord(ICEBERG_SCHEMA, row));

    writerManager.write(windowedDestination, row);

    RecordWriterManager.DestinationState state =
        writerManager.destinations.get(windowedDestination);
    // replace with a failing record writer
    FailingRecordWriter failingWriter =
        new FailingRecordWriter(
            catalog, windowedDestination.getValue(), "test_failing_writer", partitionKey);
    state.writers.put(partitionKey, failingWriter);
    writerManager.write(windowedDestination, row);

    // this tests that we indeed enter the catch block
    thrown.expect(IllegalStateException.class);
    thrown.expectMessage("Encountered 1 failed writer(s)");
    try {
      writerManager.close();
    } catch (IllegalStateException e) {
      // fetch underlying exceptions and validate
      Throwable[] underlyingExceptions = e.getSuppressed();
      assertEquals(1, underlyingExceptions.length);
      for (Throwable t : underlyingExceptions) {
        assertThat(
            t.getMessage(),
            containsString("Encountered an error when closing data writer for table"));
        assertThat(
            t.getMessage(),
            containsString(windowedDestination.getValue().getTableIdentifier().toString()));
        assertThat(t.getMessage(), containsString(failingWriter.path()));
        Throwable realCause = t.getCause();
        assertEquals("I am failing!", realCause.getMessage());
      }

      throw e;
    }
  }

  static class FailingRecordWriter extends RecordWriter {
    FailingRecordWriter(
        Catalog catalog, IcebergDestination destination, String filename, PartitionKey partitionKey)
        throws IOException {
      super(catalog, destination, filename, partitionKey);
    }

    @Override
    public void close() throws IOException {
      throw new IOException("I am failing!");
    }
  }

  @Test
  public void testColumnSpecificMetricsCollection() throws IOException {
    TableIdentifier tableId = TableIdentifier.of("default", "test_column_metrics");
    Table table = warehouse.createTable(tableId, ICEBERG_SCHEMA); // Unpartitioned table
    table
        .updateProperties()
        .set("write.metadata.metrics.column.id", "counts")
        .set("write.metadata.metrics.column.name", "counts")
        .commit();

    IcebergDestination destination =
        IcebergDestination.builder()
            .setTableIdentifier(tableId)
            .setFileFormat(FileFormat.PARQUET)
            .build();
    WindowedValue<IcebergDestination> singleDestination =
        WindowedValues.valueInGlobalWindow(destination);

    RecordWriterManager writerManager = new RecordWriterManager(catalog, "test_file_name", 1000, 3);
    Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();
    Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2, "bbb", false).build();
    Row row3 = Row.withSchema(BEAM_SCHEMA).addValues(3, "ccc", true).build();
    assertTrue("Write operation should succeed", writerManager.write(singleDestination, row1));
    assertTrue("Write operation should succeed", writerManager.write(singleDestination, row2));
    assertTrue("Write operation should succeed", writerManager.write(singleDestination, row3));
    writerManager.close();

    Map<WindowedValue<IcebergDestination>, List<SerializableDataFile>> dataFiles =
        writerManager.getSerializableDataFiles();
    assertFalse("Data files should not be empty", dataFiles.isEmpty());
    assertEquals(
        "Should have exactly one data file for an unpartitioned table", 1, dataFiles.size());

    for (Map.Entry<WindowedValue<IcebergDestination>, List<SerializableDataFile>> entry :
        dataFiles.entrySet()) {
      AppendFiles appendFiles = table.newAppend();
      for (SerializableDataFile dataFile : entry.getValue()) {
        appendFiles.appendFile(dataFile.createDataFile(table.specs()));
      }
      appendFiles.commit();
    }
    table.refresh();

    Snapshot snapshot = table.currentSnapshot();
    assertNotNull("Table should have a snapshot after writing data", snapshot);
    DataFile dataFile = snapshot.addedDataFiles(table.io()).iterator().next();

    Map<Integer, Long> valueCounts = dataFile.valueCounts();
    assertNotNull("Value counts should not be null", valueCounts);
    assertTrue("Value counts should exist for id (column 1)", valueCounts.containsKey(1));
    assertEquals("Value count for id should be 3 in this file", 3L, valueCounts.get(1).longValue());
    assertTrue("Value counts should exist for name (column 2)", valueCounts.containsKey(2));
    assertEquals(
        "Value count for name should be 3 in this file", 3L, valueCounts.get(2).longValue());

    assertNotNull(dataFile.nullValueCounts());
    assertNotNull(dataFile.columnSizes());
    Map<Integer, Long> nullValueCounts = dataFile.nullValueCounts();
    Map<Integer, Long> columnSizes = dataFile.columnSizes();
    for (int i = 1; i <= ICEBERG_SCHEMA.columns().size(); i++) {
      assertTrue(
          "Null value counts should be collected for column " + i, nullValueCounts.containsKey(i));
      assertTrue("Column sizes should be collected for column " + i, columnSizes.containsKey(i));
    }
  }

  @Test
  public void testDefaultMetrics() throws IOException {
    TableIdentifier tableId = TableIdentifier.of("default", "test_default_metrics");
    Table table = warehouse.createTable(tableId, ICEBERG_SCHEMA);
    table.updateProperties().set("write.metadata.metrics.default", "full").commit();

    IcebergDestination destination =
        IcebergDestination.builder()
            .setTableIdentifier(tableId)
            .setFileFormat(FileFormat.PARQUET)
            .build();
    WindowedValue<IcebergDestination> singleDestination =
        WindowedValues.valueInGlobalWindow(destination);

    RecordWriterManager writerManager = new RecordWriterManager(catalog, "test_file_name", 1000, 3);
    Row row1 = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();
    Row row2 = Row.withSchema(BEAM_SCHEMA).addValues(2, "bbb", false).build();
    Row row3 = Row.withSchema(BEAM_SCHEMA).addValues(3, "ccc", true).build();
    assertTrue("Write operation should succeed", writerManager.write(singleDestination, row1));
    assertTrue("Write operation should succeed", writerManager.write(singleDestination, row2));
    assertTrue("Write operation should succeed", writerManager.write(singleDestination, row3));
    writerManager.close();

    Map<WindowedValue<IcebergDestination>, List<SerializableDataFile>> dataFiles =
        writerManager.getSerializableDataFiles();
    assertFalse("Data files should not be empty", dataFiles.isEmpty());
    assertEquals(
        "Should have exactly one data file for an unpartitioned table", 1, dataFiles.size());

    for (Map.Entry<WindowedValue<IcebergDestination>, List<SerializableDataFile>> entry :
        dataFiles.entrySet()) {
      AppendFiles appendFiles = table.newAppend();
      for (SerializableDataFile dataFile : entry.getValue()) {
        appendFiles.appendFile(dataFile.createDataFile(table.specs()));
      }
      appendFiles.commit();
    }
    table.refresh();

    Snapshot snapshot = table.currentSnapshot();
    assertNotNull("Table should have a snapshot after writing data", snapshot);
    DataFile dataFile = snapshot.addedDataFiles(table.io()).iterator().next();

    assertNotNull(dataFile.valueCounts());
    assertNotNull(dataFile.nullValueCounts());
    assertNotNull(dataFile.columnSizes());
    assertNotNull(dataFile.lowerBounds());
    assertNotNull(dataFile.upperBounds());

    Map<Integer, Long> valueCounts = dataFile.valueCounts();
    Map<Integer, Long> nullValueCounts = dataFile.nullValueCounts();
    Map<Integer, Long> columnSizes = dataFile.columnSizes();
    Map<Integer, ByteBuffer> lowerBounds = dataFile.lowerBounds();
    Map<Integer, ByteBuffer> upperBounds = dataFile.upperBounds();

    for (int i = 1; i <= ICEBERG_SCHEMA.columns().size(); i++) {
      assertTrue("Value counts should be collected for column " + i, valueCounts.containsKey(i));
      long expectedCount = 3L;
      assertEquals(
          "Value count for column " + i + " should match total values in this file",
          expectedCount,
          valueCounts.get(i).longValue());

      assertTrue(
          "Null value counts should be collected for column " + i, nullValueCounts.containsKey(i));
      assertTrue("Column sizes should be collected for column " + i, columnSizes.containsKey(i));
      assertTrue("Lower bounds should be collected for column " + i, lowerBounds.containsKey(i));
      assertTrue("Upper bounds should be collected for column " + i, upperBounds.containsKey(i));

      ByteBuffer lower = lowerBounds.get(i);
      ByteBuffer upper = upperBounds.get(i);
      assertNotNull("Lower bound for column " + i + " should not be null", lower);
      assertNotNull("Upper bound for column " + i + " should not be null", upper);

      switch (i) {
        case 1:
          Integer lowerId = Conversions.fromByteBuffer(Types.IntegerType.get(), lower);
          Integer upperId = Conversions.fromByteBuffer(Types.IntegerType.get(), upper);
          assertEquals("Lower bound for id should be min value", 1, lowerId.intValue());
          assertEquals("Upper bound for id should be max value", 3, upperId.intValue());
          break;
        case 2:
          CharSequence lowerName = Conversions.fromByteBuffer(Types.StringType.get(), lower);
          CharSequence upperName = Conversions.fromByteBuffer(Types.StringType.get(), upper);
          assertEquals("Lower bound for name should be min value", "aaa", lowerName.toString());
          assertEquals("Upper bound for name should be max value", "ccc", upperName.toString());
          break;
        case 3:
          Boolean lowerBool = Conversions.fromByteBuffer(Types.BooleanType.get(), lower);
          Boolean upperBool = Conversions.fromByteBuffer(Types.BooleanType.get(), upper);
          assertEquals("Lower bound for bool should be min value", false, lowerBool);
          assertEquals("Upper bound for bool should be max value", true, upperBool);
          break;
        default:
          throw new IllegalStateException("Unexpected column index: " + i);
      }
    }
  }

  @Test
  public void testRecordWriterDoesNotCloseSharedFileIO() throws IOException {
    TableIdentifier tableId =
        TableIdentifier.of(
            "default",
            "table_"
                + testName.getMethodName()
                + "_"
                + UUID.randomUUID().toString().replace("-", "").substring(0, 6));
    Table table = warehouse.createTable(tableId, ICEBERG_SCHEMA);

    CloseTrackingFileIO trackingFileIO = new CloseTrackingFileIO(table.io());
    Table spyTable = Mockito.spy(table);
    Mockito.doReturn(trackingFileIO).when(spyTable).io();

    PartitionKey partitionKey = new PartitionKey(spyTable.spec(), spyTable.schema());
    RecordWriter writer =
        new RecordWriter(spyTable, FileFormat.PARQUET, "file.parquet", partitionKey);

    Row row = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();

    writer.write(IcebergUtils.beamRowToIcebergRecord(ICEBERG_SCHEMA, row));
    writer.close();

    // RecordWriter must NOT close FileIO — it may be a shared instance.
    assertFalse("RecordWriter.close() must not close the shared FileIO", trackingFileIO.closed);
  }

  /**
   * Verifies that when multiple writers share the same FileIO, closing any writer does not close
   * the shared FileIO — that is the responsibility of RecordWriterManager.close().
   */
  @Test
  public void testMultipleWritersSharingFileIOSurviveBatchClose() throws IOException {
    // Create two tables that share the same FileIO (simulating dynamic destinations
    // backed by the same catalog)
    TableIdentifier tableId1 =
        TableIdentifier.of(
            "default",
            "table_batch_close_a_" + UUID.randomUUID().toString().replace("-", "").substring(0, 6));
    TableIdentifier tableId2 =
        TableIdentifier.of(
            "default",
            "table_batch_close_b_" + UUID.randomUUID().toString().replace("-", "").substring(0, 6));
    Table table1 = warehouse.createTable(tableId1, ICEBERG_SCHEMA);
    Table table2 = warehouse.createTable(tableId2, ICEBERG_SCHEMA);

    // Both tables share the same CloseTrackingFileIO — mirrors how some catalogs
    // return a shared FileIO instance across tables
    CloseTrackingFileIO sharedFileIO = new CloseTrackingFileIO(table1.io());
    Table spyTable1 = Mockito.spy(table1);
    Table spyTable2 = Mockito.spy(table2);
    Mockito.doReturn(sharedFileIO).when(spyTable1).io();
    Mockito.doReturn(sharedFileIO).when(spyTable2).io();

    PartitionKey pk1 = new PartitionKey(spyTable1.spec(), spyTable1.schema());
    PartitionKey pk2 = new PartitionKey(spyTable2.spec(), spyTable2.schema());

    RecordWriter writer1 = new RecordWriter(spyTable1, FileFormat.PARQUET, "file1.parquet", pk1);
    RecordWriter writer2 = new RecordWriter(spyTable2, FileFormat.PARQUET, "file2.parquet", pk2);

    Row row = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();
    Record record = IcebergUtils.beamRowToIcebergRecord(ICEBERG_SCHEMA, row);

    writer1.write(record);
    writer2.write(record);

    writer1.close();
    assertFalse("FileIO must remain open between batch writer closes", sharedFileIO.closed);

    writer2.close();
    assertFalse("FileIO must remain open after all writers close", sharedFileIO.closed);

    // Both writers produced valid data files
    assertNotNull(writer1.getDataFile());
    assertNotNull(writer2.getDataFile());
  }

  /**
   * Verifies that RecordWriterManager.close() flushes data files from multiple destinations and
   * closes the shared FileIO.
   */
  @Test
  public void testRecordWriterManagerDoesNotCloseSharedFileIO() throws IOException {
    String tableName1 =
        "table_mgr_io_a_" + UUID.randomUUID().toString().replace("-", "").substring(0, 6);
    String tableName2 =
        "table_mgr_io_b_" + UUID.randomUUID().toString().replace("-", "").substring(0, 6);
    TableIdentifier tableId1 = TableIdentifier.of("default", tableName1);
    TableIdentifier tableId2 = TableIdentifier.of("default", tableName2);

    Table realTable1 = warehouse.createTable(tableId1, ICEBERG_SCHEMA);
    Table realTable2 = warehouse.createTable(tableId2, ICEBERG_SCHEMA);

    CloseTrackingFileIO sharedTrackingIO = new CloseTrackingFileIO(realTable1.io());
    Table spyTable1 = Mockito.spy(realTable1);
    Table spyTable2 = Mockito.spy(realTable2);
    Mockito.doReturn(sharedTrackingIO).when(spyTable1).io();
    Mockito.doReturn(sharedTrackingIO).when(spyTable2).io();

    Catalog spyCatalog = Mockito.spy(catalog);
    Mockito.doReturn(spyTable1).when(spyCatalog).loadTable(tableId1);
    Mockito.doReturn(spyTable2).when(spyCatalog).loadTable(tableId2);

    WindowedValue<IcebergDestination> dest1 = getWindowedDestination(tableName1, null);
    WindowedValue<IcebergDestination> dest2 = getWindowedDestination(tableName2, null);

    RecordWriterManager writerManager =
        new RecordWriterManager(spyCatalog, "test_file_name", 1000, 3);

    Row row = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();
    assertTrue(writerManager.write(dest1, row));
    assertTrue(writerManager.write(dest2, row));
    assertEquals(2, writerManager.openWriters);

    writerManager.close();

    Map<WindowedValue<IcebergDestination>, List<SerializableDataFile>> dataFiles =
        writerManager.getSerializableDataFiles();
    assertTrue("Should have data files for dest1", dataFiles.containsKey(dest1));
    assertTrue("Should have data files for dest2", dataFiles.containsKey(dest2));
    assertFalse(
        "Shared FileIO should NOT be closed by RecordWriterManager", sharedTrackingIO.closed);
  }

  private static final class CloseTrackingFileIO implements FileIO {
    private final FileIO delegate;
    volatile boolean closed = false;

    CloseTrackingFileIO(FileIO delegate) {
      this.delegate = delegate;
    }

    @Override
    public InputFile newInputFile(String path) {
      return delegate.newInputFile(path);
    }

    @Override
    public OutputFile newOutputFile(String path) {
      OutputFile underlying = delegate.newOutputFile(path);
      return new CloseAwareOutputFile(underlying, this);
    }

    @Override
    public void deleteFile(String path) {
      delegate.deleteFile(path);
    }

    @Override
    public Map<String, String> properties() {
      return delegate.properties();
    }

    @Override
    public void close() {
      closed = true;
      delegate.close();
    }
  }

  private static final class CloseAwareOutputFile implements OutputFile {
    private final OutputFile delegate;
    private final CloseTrackingFileIO io;

    CloseAwareOutputFile(OutputFile delegate, CloseTrackingFileIO io) {
      this.delegate = delegate;
      this.io = io;
    }

    @Override
    public PositionOutputStream create() {
      if (io.closed) {
        throw new IllegalStateException("Connection pool shut down");
      }
      return delegate.create();
    }

    @Override
    public PositionOutputStream createOrOverwrite() {
      if (io.closed) {
        throw new IllegalStateException("Connection pool shut down");
      }
      return delegate.createOrOverwrite();
    }

    @Override
    public String location() {
      return delegate.location();
    }

    @Override
    public InputFile toInputFile() {
      return delegate.toInputFile();
    }
  }

  @Test
  public void testGetOrCreateTable_refreshLogic() {
    Table mockTable = mock(Table.class);
    TableIdentifier identifier = TableIdentifier.of("db", "table");
    IcebergDestination destination =
        IcebergDestination.builder()
            .setTableIdentifier(identifier)
            .setFileFormat(FileFormat.PARQUET)
            .setTableCreateConfig(
                IcebergTableCreateConfig.builder()
                    .setPartitionFields(null)
                    .setSchema(BEAM_SCHEMA)
                    .build())
            .build();
    // The schema is only used if the table is created, so a null is fine for this
    // test.
    Schema beamSchema = null;

    // Instantiate a RecordWriterManager with a dummy catalog.
    RecordWriterManager writer = new RecordWriterManager(null, "p", 1L, 1);

    // Clean up cache before test
    RecordWriterManager.LAST_REFRESHED_TABLE_CACHE.invalidateAll();

    // --- 1. Test the fast path (entry is not stale) ---
    Instant freshTimestamp = Instant.now().minus(Duration.ofMinutes(1));
    RecordWriterManager.LastRefreshedTable freshEntry =
        new RecordWriterManager.LastRefreshedTable(mockTable, freshTimestamp);
    RecordWriterManager.LAST_REFRESHED_TABLE_CACHE.put(identifier, freshEntry);

    // Access the table
    writer.getOrCreateTable(destination, beamSchema);

    // Verify that refresh() was NOT called because the entry is fresh.
    verify(mockTable, never()).refresh();

    // --- 2. Test the stale path (entry is stale) ---
    Instant staleTimestamp = Instant.now().minus(Duration.ofMinutes(5));
    RecordWriterManager.LastRefreshedTable staleEntry =
        new RecordWriterManager.LastRefreshedTable(mockTable, staleTimestamp);
    RecordWriterManager.LAST_REFRESHED_TABLE_CACHE.put(identifier, staleEntry);

    // Access the table again
    writer.getOrCreateTable(destination, beamSchema);

    // Verify that refresh() WAS called exactly once because the entry was stale.
    verify(mockTable, times(1)).refresh();
  }

  /**
   * Verifies that the shared FileIO survives across multiple bundles. This is the core regression
   * test: if RecordWriterManager.close() closed the FileIO, the second bundle would fail with
   * "Connection pool shut down".
   */
  @Test
  public void testFileIOSurvivesAcrossBundles() throws IOException {
    String tableName =
        "table_survive_" + UUID.randomUUID().toString().replace("-", "").substring(0, 6);
    TableIdentifier tableId = TableIdentifier.of("default", tableName);

    Table realTable = warehouse.createTable(tableId, ICEBERG_SCHEMA);

    CloseTrackingFileIO sharedIO = new CloseTrackingFileIO(realTable.io());
    Table spyTable = Mockito.spy(realTable);
    Mockito.doReturn(sharedIO).when(spyTable).io();

    Catalog spyCatalog = Mockito.spy(catalog);
    Mockito.doReturn(spyTable).when(spyCatalog).loadTable(tableId);

    WindowedValue<IcebergDestination> dest = getWindowedDestination(tableName, null);
    Row row = Row.withSchema(BEAM_SCHEMA).addValues(1, "aaa", true).build();

    // Bundle 1: write and close
    RecordWriterManager bundle1 = new RecordWriterManager(spyCatalog, "file_b1", 1000, 3);
    assertTrue(bundle1.write(dest, row));
    bundle1.close();
    assertFalse("FileIO must survive after bundle 1 close", sharedIO.closed);
    assertTrue(
        "Bundle 1 should produce data files", bundle1.getSerializableDataFiles().containsKey(dest));

    // Bundle 2: write and close using the same catalog (simulates DoFn reuse)
    RecordWriterManager bundle2 = new RecordWriterManager(spyCatalog, "file_b2", 1000, 3);
    assertTrue(bundle2.write(dest, row));
    bundle2.close();
    assertFalse("FileIO must survive after bundle 2 close", sharedIO.closed);
    assertTrue(
        "Bundle 2 should produce data files", bundle2.getSerializableDataFiles().containsKey(dest));
  }
}
