/*
 * 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.formats.avro;

import java.io.Serializable;
import java.lang.reflect.Array;
import java.nio.ByteBuffer;
import java.time.Instant;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.LocalTime;
import java.time.temporal.ChronoField;
import java.util.List;
import java.util.Map;
import org.apache.avro.generic.GenericFixed;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.generic.IndexedRecord;
import org.apache.flink.annotation.Internal;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.data.DecimalData;
import org.apache.flink.table.data.GenericArrayData;
import org.apache.flink.table.data.GenericMapData;
import org.apache.flink.table.data.GenericRowData;
import org.apache.flink.table.data.RowData;
import org.apache.flink.table.data.StringData;
import org.apache.flink.table.data.TimestampData;
import org.apache.flink.table.types.logical.ArrayType;
import org.apache.flink.table.types.logical.DecimalType;
import org.apache.flink.table.types.logical.LogicalType;
import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.table.types.logical.utils.LogicalTypeUtils;
import org.apache.iceberg.flink.formats.avro.typeutils.AvroSchemaConverter;
import org.apache.iceberg.relocated.com.google.common.collect.Maps;

/**
 * Tool class used to convert from Avro {@link GenericRecord} to {@link RowData}.
 *
 * <p>This class is adapted in Iceberg to add support for nanosecond precision timestamps
 * (FLINK-39251). Once that ticket is resolved in Flink, this custom converter may be removed.
 */
@Internal
public class AvroToRowDataConverters {

  private AvroToRowDataConverters() {}

  /**
   * Runtime converter that converts Avro data structures into objects of Flink Table &amp; SQL
   * internal data structures.
   */
  @FunctionalInterface
  public interface AvroToRowDataConverter extends Serializable {
    Object convert(Object object);
  }

  // -------------------------------------------------------------------------------------
  // Runtime Converters
  // -------------------------------------------------------------------------------------

  public static AvroToRowDataConverter createRowConverter(RowType rowType) {
    return createRowConverter(rowType, true);
  }

  public static AvroToRowDataConverter createRowConverter(
      RowType rowType, boolean legacyTimestampMapping) {
    final AvroToRowDataConverter[] fieldConverters =
        rowType.getFields().stream()
            .map(RowType.RowField::getType)
            .map(type -> createNullableConverter(type, legacyTimestampMapping))
            .toArray(AvroToRowDataConverter[]::new);
    final int arity = rowType.getFieldCount();

    return avroObject -> {
      IndexedRecord record = (IndexedRecord) avroObject;
      GenericRowData row = new GenericRowData(arity);
      for (int i = 0; i < arity; ++i) {
        // avro always deserialize successfully even though the type isn't matched
        // so no need to throw exception about which field can't be deserialized
        row.setField(i, fieldConverters[i].convert(record.get(i)));
      }
      return row;
    };
  }

  /** Creates a runtime converter which is null safe. */
  private static AvroToRowDataConverter createNullableConverter(
      LogicalType type, boolean legacyTimestampMapping) {
    final AvroToRowDataConverter converter = createConverter(type, legacyTimestampMapping);
    return avroObject -> {
      if (avroObject == null) {
        return null;
      }
      return converter.convert(avroObject);
    };
  }

  /** Creates a runtime converter which assuming input object is not null. */
  private static AvroToRowDataConverter createConverter(
      LogicalType type, boolean legacyTimestampMapping) {
    switch (type.getTypeRoot()) {
      case NULL:
        return avroObject -> null;
      case TINYINT:
        return avroObject -> ((Integer) avroObject).byteValue();
      case SMALLINT:
        return avroObject -> ((Integer) avroObject).shortValue();
      case BOOLEAN: // boolean
      case INTEGER: // int
      case INTERVAL_YEAR_MONTH: // long
      case BIGINT: // long
      case INTERVAL_DAY_TIME: // long
      case FLOAT: // float
      case DOUBLE: // double
        return avroObject -> avroObject;
      case DATE:
        return AvroToRowDataConverters::convertToDate;
      case TIME_WITHOUT_TIME_ZONE:
        return AvroToRowDataConverters::convertToTime;
      case TIMESTAMP_WITHOUT_TIME_ZONE:
        // Iceberg: Added support for nanoseconds precision (FLINK-39251)
        return avroObject -> convertToTimestamp(avroObject, type);
      case TIMESTAMP_WITH_LOCAL_TIME_ZONE:
        if (legacyTimestampMapping) {
          throw new UnsupportedOperationException("Unsupported type: " + type);
        } else {
          // Iceberg: Added support for nanoseconds precision (FLINK-39251)
          return avroObject -> convertToTimestamp(avroObject, type);
        }
      case CHAR:
      case VARCHAR:
        return avroObject -> StringData.fromString(avroObject.toString());
      case BINARY:
      case VARBINARY:
        return AvroToRowDataConverters::convertToBytes;
      case DECIMAL:
        return createDecimalConverter((DecimalType) type);
      case ARRAY:
        return createArrayConverter((ArrayType) type, legacyTimestampMapping);
      case ROW:
        return createRowConverter((RowType) type);
      case MAP:
      case MULTISET:
        return createMapConverter(type, legacyTimestampMapping);
      case RAW:
      default:
        throw new UnsupportedOperationException("Unsupported type: " + type);
    }
  }

  private static AvroToRowDataConverter createDecimalConverter(DecimalType decimalType) {
    final int precision = decimalType.getPrecision();
    final int scale = decimalType.getScale();
    return avroObject -> {
      final byte[] bytes;
      if (avroObject instanceof GenericFixed) {
        bytes = ((GenericFixed) avroObject).bytes();
      } else if (avroObject instanceof ByteBuffer) {
        ByteBuffer byteBuffer = (ByteBuffer) avroObject;
        bytes = new byte[byteBuffer.remaining()];
        byteBuffer.get(bytes);
      } else {
        bytes = (byte[]) avroObject;
      }
      return DecimalData.fromUnscaledBytes(bytes, precision, scale);
    };
  }

  private static AvroToRowDataConverter createArrayConverter(
      ArrayType arrayType, boolean legacyTimestampMapping) {
    final AvroToRowDataConverter elementConverter =
        createNullableConverter(arrayType.getElementType(), legacyTimestampMapping);
    final Class<?> elementClass =
        LogicalTypeUtils.toInternalConversionClass(arrayType.getElementType());

    return avroObject -> {
      final List<?> list = (List<?>) avroObject;
      final int length = list.size();
      final Object[] array = (Object[]) Array.newInstance(elementClass, length);
      for (int i = 0; i < length; ++i) {
        array[i] = elementConverter.convert(list.get(i));
      }
      return new GenericArrayData(array);
    };
  }

  private static AvroToRowDataConverter createMapConverter(
      LogicalType type, boolean legacyTimestampMapping) {
    final AvroToRowDataConverter keyConverter =
        createConverter(DataTypes.STRING().getLogicalType(), legacyTimestampMapping);
    final AvroToRowDataConverter valueConverter =
        createNullableConverter(
            AvroSchemaConverter.extractValueTypeToAvroMap(type), legacyTimestampMapping);

    return avroObject -> {
      final Map<?, ?> map = (Map<?, ?>) avroObject;
      Map<Object, Object> result = Maps.newHashMap();
      for (Map.Entry<?, ?> entry : map.entrySet()) {
        Object key = keyConverter.convert(entry.getKey());
        Object value = valueConverter.convert(entry.getValue());
        result.put(key, value);
      }
      return new GenericMapData(result);
    };
  }

  private static TimestampData convertToTimestamp(Object object, LogicalType type) {
    int precision = 3;
    if (type instanceof org.apache.flink.table.types.logical.TimestampType) {
      precision = ((org.apache.flink.table.types.logical.TimestampType) type).getPrecision();
    } else if (type instanceof org.apache.flink.table.types.logical.LocalZonedTimestampType) {
      precision =
          ((org.apache.flink.table.types.logical.LocalZonedTimestampType) type).getPrecision();
    }

    if (object instanceof Long) {
      long timeLong = (Long) object;
      if (precision <= 3) {
        return TimestampData.fromEpochMillis(timeLong);
      } else if (precision <= 6) {
        return TimestampData.fromEpochMillis(
            Math.floorDiv(timeLong, 1000L), (int) Math.floorMod(timeLong, 1000L) * 1000);
      } else {
        // Iceberg: Added support for nanoseconds precision (FLINK-39251)
        return TimestampData.fromEpochMillis(
            Math.floorDiv(timeLong, 1_000_000L), (int) Math.floorMod(timeLong, 1_000_000L));
      }
    } else if (object instanceof Instant) {
      return TimestampData.fromInstant((Instant) object);
    } else if (object instanceof LocalDateTime) {
      return TimestampData.fromLocalDateTime((LocalDateTime) object);
    } else {
      JodaConverter jodaConverter = JodaConverter.getConverter();
      if (jodaConverter != null) {
        return TimestampData.fromEpochMillis(jodaConverter.convertTimestamp(object));
      } else {
        throw new IllegalArgumentException(
            "Unexpected object type for TIMESTAMP logical type. Received: " + object);
      }
    }
  }

  private static int convertToDate(Object object) {
    if (object instanceof Integer) {
      return (Integer) object;
    } else if (object instanceof LocalDate) {
      return (int) ((LocalDate) object).toEpochDay();
    } else {
      JodaConverter jodaConverter = JodaConverter.getConverter();
      if (jodaConverter != null) {
        return (int) jodaConverter.convertDate(object);
      } else {
        throw new IllegalArgumentException(
            "Unexpected object type for DATE logical type. Received: " + object);
      }
    }
  }

  private static int convertToTime(Object object) {
    final int millis;
    if (object instanceof Integer) {
      millis = (Integer) object;
    } else if (object instanceof LocalTime) {
      millis = ((LocalTime) object).get(ChronoField.MILLI_OF_DAY);
    } else {
      JodaConverter jodaConverter = JodaConverter.getConverter();
      if (jodaConverter != null) {
        millis = jodaConverter.convertTime(object);
      } else {
        throw new IllegalArgumentException(
            "Unexpected object type for TIME logical type. Received: " + object);
      }
    }
    return millis;
  }

  private static byte[] convertToBytes(Object object) {
    if (object instanceof GenericFixed) {
      return ((GenericFixed) object).bytes();
    } else if (object instanceof ByteBuffer) {
      ByteBuffer byteBuffer = (ByteBuffer) object;
      byte[] bytes = new byte[byteBuffer.remaining()];
      byteBuffer.get(bytes);
      return bytes;
    } else {
      return (byte[]) object;
    }
  }
}
