/*
 * 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;

import java.io.Serializable;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.AbstractMap;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
import org.apache.iceberg.exceptions.ValidationException;
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
import org.apache.iceberg.relocated.com.google.common.collect.ListMultimap;
import org.apache.iceberg.relocated.com.google.common.collect.Lists;
import org.apache.iceberg.relocated.com.google.common.collect.Maps;
import org.apache.iceberg.relocated.com.google.common.collect.Multimaps;
import org.apache.iceberg.relocated.com.google.common.collect.Sets;
import org.apache.iceberg.transforms.Transform;
import org.apache.iceberg.transforms.Transforms;
import org.apache.iceberg.transforms.UnknownTransform;
import org.apache.iceberg.types.Type;
import org.apache.iceberg.types.TypeUtil;
import org.apache.iceberg.types.Types;
import org.apache.iceberg.types.Types.StructType;

/**
 * Represents how to produce partition data for a table.
 *
 * <p>Partition data is produced by transforming columns in a table. Each column transform is
 * represented by a named {@link PartitionField}.
 */
public class PartitionSpec implements Serializable {
  // IDs for partition fields start at 1000
  private static final int PARTITION_DATA_ID_START = 1000;

  private final Schema schema;

  // this is ordered so that DataFile has a consistent schema
  private final int specId;
  private final PartitionField[] fields;
  private transient volatile ListMultimap<Integer, PartitionField> fieldsBySourceId = null;
  private transient volatile Class<?>[] lazyJavaClasses = null;
  private transient volatile StructType lazyPartitionType = null;
  private transient volatile StructType lazyRawPartitionType = null;
  private transient volatile List<PartitionField> fieldList = null;
  private final int lastAssignedFieldId;

  private PartitionSpec(
      Schema schema, int specId, List<PartitionField> fields, int lastAssignedFieldId) {
    this.schema = schema;
    this.specId = specId;
    this.fields = fields.toArray(new PartitionField[0]);
    this.lastAssignedFieldId = lastAssignedFieldId;
  }

  /** Returns the {@link Schema} for this spec. */
  public Schema schema() {
    return schema;
  }

  /** Returns the ID of this spec. */
  public int specId() {
    return specId;
  }

  /** Returns the list of {@link PartitionField partition fields} for this spec. */
  public List<PartitionField> fields() {
    return lazyFieldList();
  }

  public boolean isPartitioned() {
    return fields.length > 0 && fields().stream().anyMatch(f -> !f.transform().isVoid());
  }

  public boolean isUnpartitioned() {
    return !isPartitioned();
  }

  int lastAssignedFieldId() {
    return lastAssignedFieldId;
  }

  public UnboundPartitionSpec toUnbound() {
    UnboundPartitionSpec.Builder builder = UnboundPartitionSpec.builder().withSpecId(specId);

    for (PartitionField field : fields) {
      builder.addField(
          field.transform().toString(), field.sourceId(), field.fieldId(), field.name());
    }

    return builder.build();
  }

  /**
   * Returns the {@link PartitionField field} that partitions the given source field
   *
   * @param fieldId a field id from the source schema
   * @return the {@link PartitionField field} that partitions the given source field
   */
  public List<PartitionField> getFieldsBySourceId(int fieldId) {
    return lazyFieldsBySourceId().get(fieldId);
  }

  /** Returns a {@link StructType} for partition data defined by this spec. */
  public StructType partitionType() {
    if (lazyPartitionType == null) {
      synchronized (this) {
        if (lazyPartitionType == null) {
          List<Types.NestedField> structFields = Lists.newArrayListWithExpectedSize(fields.length);

          for (PartitionField field : fields) {
            structFields.add(
                Types.NestedField.optional(field.fieldId(), field.name(), resultType(field)));
          }

          this.lazyPartitionType = Types.StructType.of(structFields);
        }
      }
    }

    return lazyPartitionType;
  }

  /**
   * Returns a struct matching partition information as written into manifest files. See {@link
   * #partitionType()} for a struct with field ID's potentially re-assigned to avoid conflict.
   */
  public StructType rawPartitionType() {
    if (schema.idsToOriginal().isEmpty()) {
      // not re-assigned.
      return partitionType();
    }
    if (lazyRawPartitionType == null) {
      synchronized (this) {
        if (lazyRawPartitionType == null) {
          this.lazyRawPartitionType =
              StructType.of(
                  partitionType().fields().stream()
                      .map(f -> f.withFieldId(schema.idsToOriginal().get(f.fieldId())))
                      .collect(Collectors.toList()));
        }
      }
    }

    return lazyRawPartitionType;
  }

  public Class<?>[] javaClasses() {
    if (lazyJavaClasses == null) {
      synchronized (this) {
        if (lazyJavaClasses == null) {
          Class<?>[] classes = new Class<?>[fields.length];
          for (int i = 0; i < fields.length; i += 1) {
            PartitionField field = fields[i];
            if (field.transform() instanceof UnknownTransform) {
              classes[i] = Object.class;
            } else {
              classes[i] = resultType(field).typeId().javaClass();
            }
          }

          this.lazyJavaClasses = classes;
        }
      }
    }

    return lazyJavaClasses;
  }

  private Type resultType(PartitionField field) {
    Type sourceType = schema.findType(field.sourceId());
    if (sourceType == null) {
      // When the source field has been dropped, the source type has been lost
      // Transforms with a fixed result type still work, so use unknown for source
      sourceType = Types.UnknownType.get();
    }

    return field.transform().getResultType(sourceType);
  }

  @SuppressWarnings("unchecked")
  private <T> T get(StructLike data, int pos, Class<?> javaClass) {
    return data.get(pos, (Class<T>) javaClass);
  }

  private String escape(String string) {
    return URLEncoder.encode(string, StandardCharsets.UTF_8);
  }

  public String partitionToPath(StructLike data) {
    StringBuilder sb = new StringBuilder();
    Class<?>[] javaClasses = javaClasses();
    List<Types.NestedField> outputFields = partitionType().fields();
    for (int i = 0; i < javaClasses.length; i += 1) {
      PartitionField field = fields[i];
      Type type = outputFields.get(i).type();
      String valueString = field.transform().toHumanString(type, get(data, i, javaClasses[i]));

      if (i > 0) {
        sb.append("/");
      }
      sb.append(escape(field.name())).append("=").append(escape(valueString));
    }
    return sb.toString();
  }

  /**
   * Returns true if this spec is equivalent to the other, with partition field ids ignored. That
   * is, if both specs have the same number of fields, field order, field name, source columns, and
   * transforms.
   *
   * @param other another PartitionSpec
   * @return true if the specs have the same fields, source columns, and transforms.
   */
  public boolean compatibleWith(PartitionSpec other) {
    if (equals(other)) {
      return true;
    }

    if (fields.length != other.fields.length) {
      return false;
    }

    for (int i = 0; i < fields.length; i += 1) {
      PartitionField thisField = fields[i];
      PartitionField thatField = other.fields[i];
      if (thisField.sourceId() != thatField.sourceId()
          || !thisField.transform().toString().equals(thatField.transform().toString())
          || !thisField.name().equals(thatField.name())) {
        return false;
      }
    }

    return true;
  }

  @Override
  public boolean equals(Object other) {
    if (this == other) {
      return true;
    } else if (!(other instanceof PartitionSpec)) {
      return false;
    }

    PartitionSpec that = (PartitionSpec) other;
    if (this.specId != that.specId) {
      return false;
    }
    return Arrays.equals(fields, that.fields);
  }

  @Override
  public int hashCode() {
    return 31 * Integer.hashCode(specId) + Arrays.hashCode(fields);
  }

  private List<PartitionField> lazyFieldList() {
    if (fieldList == null) {
      synchronized (this) {
        if (fieldList == null) {
          this.fieldList = ImmutableList.copyOf(fields);
        }
      }
    }
    return fieldList;
  }

  private ListMultimap<Integer, PartitionField> lazyFieldsBySourceId() {
    if (fieldsBySourceId == null) {
      synchronized (this) {
        if (fieldsBySourceId == null) {
          ListMultimap<Integer, PartitionField> multiMap =
              Multimaps.newListMultimap(
                  Maps.newHashMap(), () -> Lists.newArrayListWithCapacity(fields.length));
          for (PartitionField field : fields) {
            multiMap.put(field.sourceId(), field);
          }
          this.fieldsBySourceId = multiMap;
        }
      }
    }

    return fieldsBySourceId;
  }

  /**
   * Returns the source field ids for identity partitions.
   *
   * @return a set of source ids for the identity partitions.
   */
  public Set<Integer> identitySourceIds() {
    Set<Integer> sourceIds = Sets.newHashSet();
    for (PartitionField field : fields()) {
      if (field.transform().isIdentity()) {
        sourceIds.add(field.sourceId());
      }
    }

    return sourceIds;
  }

  @Override
  public String toString() {
    StringBuilder sb = new StringBuilder();
    sb.append("[");
    for (PartitionField field : fields) {
      sb.append("\n");
      sb.append("  ").append(field);
    }
    if (fields.length > 0) {
      sb.append("\n");
    }
    sb.append("]");
    return sb.toString();
  }

  private static final PartitionSpec UNPARTITIONED_SPEC =
      new PartitionSpec(new Schema(), 0, ImmutableList.of(), unpartitionedLastAssignedId());

  /**
   * Returns a spec for unpartitioned tables.
   *
   * @return a partition spec with no partitions
   */
  public static PartitionSpec unpartitioned() {
    return UNPARTITIONED_SPEC;
  }

  private static int unpartitionedLastAssignedId() {
    return PARTITION_DATA_ID_START - 1;
  }

  /**
   * Creates a new {@link Builder partition spec builder} for the given {@link Schema}.
   *
   * @param schema a schema
   * @return a partition spec builder for the given schema
   */
  public static Builder builderFor(Schema schema) {
    return new Builder(schema);
  }

  /**
   * Used to create valid {@link PartitionSpec partition specs}.
   *
   * <p>Call {@link #builderFor(Schema)} to create a new builder.
   */
  public static class Builder {
    private final Schema schema;
    private final List<PartitionField> fields = Lists.newArrayList();
    private final Set<String> partitionNames = Sets.newHashSet();
    private final Map<Map.Entry<Integer, String>, PartitionField> dedupFields = Maps.newHashMap();
    private int specId = 0;
    private final AtomicInteger lastAssignedFieldId =
        new AtomicInteger(unpartitionedLastAssignedId());
    // check if there are conflicts between partition and schema field name
    private boolean checkConflicts = true;
    private boolean caseSensitive = true;

    private Builder(Schema schema) {
      this.schema = schema;
    }

    private int nextFieldId() {
      return lastAssignedFieldId.incrementAndGet();
    }

    private void checkAndAddPartitionName(String name) {
      checkAndAddPartitionName(name, null);
    }

    Builder checkConflicts(boolean check) {
      checkConflicts = check;
      return this;
    }

    private void checkAndAddPartitionName(String name, Integer sourceColumnId) {
      Types.NestedField schemaField =
          this.caseSensitive ? schema.findField(name) : schema.caseInsensitiveFindField(name);
      if (checkConflicts) {
        if (sourceColumnId == null) {
          Preconditions.checkArgument(
              schemaField == null,
              "Cannot create partition from name that exists in schema: %s",
              name);
        } else {
          boolean sourceFieldExists = schema.findField(sourceColumnId) != null;
          // For identity transforms, require the partition name to match the source column when it
          // still exists in the schema. When the source was dropped, the spec may be historical;
          // skip the identity name check in that case.
          if (sourceFieldExists) {
            Preconditions.checkArgument(
                schemaField == null || schemaField.fieldId() == sourceColumnId,
                "Cannot create identity partition sourced from different field in schema: %s",
                name);
          }
        }
      }
      Preconditions.checkArgument(!name.isEmpty(), "Cannot use empty partition name: %s", name);
      Preconditions.checkArgument(
          !partitionNames.contains(name), "Cannot use partition name more than once: %s", name);
      partitionNames.add(name);
    }

    private void checkForRedundantPartitions(PartitionField field) {
      Map.Entry<Integer, String> dedupKey =
          new AbstractMap.SimpleEntry<>(field.sourceId(), field.transform().dedupName());
      PartitionField partitionField = dedupFields.get(dedupKey);
      Preconditions.checkArgument(
          partitionField == null,
          "Cannot add redundant partition: %s conflicts with %s",
          partitionField,
          field);
      dedupFields.put(dedupKey, field);
    }

    public Builder caseSensitive(boolean sensitive) {
      this.caseSensitive = sensitive;
      return this;
    }

    public Builder withSpecId(int newSpecId) {
      this.specId = newSpecId;
      return this;
    }

    private Types.NestedField findSourceColumn(String sourceName) {
      Types.NestedField sourceColumn =
          this.caseSensitive
              ? schema.findField(sourceName)
              : schema.caseInsensitiveFindField(sourceName);
      Preconditions.checkArgument(
          sourceColumn != null, "Cannot find source column: %s", sourceName);
      return sourceColumn;
    }

    public Builder identity(String sourceName, String targetName) {
      return identity(findSourceColumn(sourceName), targetName);
    }

    private Builder identity(Types.NestedField sourceColumn, String targetName) {
      checkAndAddPartitionName(targetName, sourceColumn.fieldId());
      PartitionField field =
          new PartitionField(
              sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.identity());
      checkForRedundantPartitions(field);
      fields.add(field);
      return this;
    }

    public Builder identity(String sourceName) {
      Types.NestedField sourceColumn = findSourceColumn(sourceName);
      return identity(sourceColumn, schema.findColumnName(sourceColumn.fieldId()));
    }

    public Builder year(String sourceName, String targetName) {
      return year(findSourceColumn(sourceName), targetName);
    }

    private Builder year(Types.NestedField sourceColumn, String targetName) {
      checkAndAddPartitionName(targetName);
      PartitionField field =
          new PartitionField(sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.year());
      checkForRedundantPartitions(field);
      fields.add(field);
      return this;
    }

    public Builder year(String sourceName) {
      Types.NestedField sourceColumn = findSourceColumn(sourceName);
      String columnName = schema.findColumnName(sourceColumn.fieldId());
      return year(sourceColumn, columnName + "_year");
    }

    public Builder month(String sourceName, String targetName) {
      return month(findSourceColumn(sourceName), targetName);
    }

    private Builder month(Types.NestedField sourceColumn, String targetName) {
      checkAndAddPartitionName(targetName);
      PartitionField field =
          new PartitionField(sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.month());
      checkForRedundantPartitions(field);
      fields.add(field);
      return this;
    }

    public Builder month(String sourceName) {
      Types.NestedField sourceColumn = findSourceColumn(sourceName);
      String columnName = schema.findColumnName(sourceColumn.fieldId());
      return month(sourceColumn, columnName + "_month");
    }

    public Builder day(String sourceName, String targetName) {
      return day(findSourceColumn(sourceName), targetName);
    }

    private Builder day(Types.NestedField sourceColumn, String targetName) {
      checkAndAddPartitionName(targetName);
      PartitionField field =
          new PartitionField(sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.day());
      checkForRedundantPartitions(field);
      fields.add(field);
      return this;
    }

    public Builder day(String sourceName) {
      Types.NestedField sourceColumn = findSourceColumn(sourceName);
      String columnName = schema.findColumnName(sourceColumn.fieldId());
      return day(sourceColumn, columnName + "_day");
    }

    public Builder hour(String sourceName, String targetName) {
      return hour(findSourceColumn(sourceName), targetName);
    }

    private Builder hour(Types.NestedField sourceColumn, String targetName) {
      checkAndAddPartitionName(targetName);
      PartitionField field =
          new PartitionField(sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.hour());
      checkForRedundantPartitions(field);
      fields.add(field);
      return this;
    }

    public Builder hour(String sourceName) {
      Types.NestedField sourceColumn = findSourceColumn(sourceName);
      String columnName = schema.findColumnName(sourceColumn.fieldId());
      return hour(sourceColumn, columnName + "_hour");
    }

    public Builder bucket(String sourceName, int numBuckets, String targetName) {
      return bucket(findSourceColumn(sourceName), numBuckets, targetName);
    }

    private Builder bucket(Types.NestedField sourceColumn, int numBuckets, String targetName) {
      checkAndAddPartitionName(targetName);
      fields.add(
          new PartitionField(
              sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.bucket(numBuckets)));
      return this;
    }

    public Builder bucket(String sourceName, int numBuckets) {
      Types.NestedField sourceColumn = findSourceColumn(sourceName);
      String columnName = schema.findColumnName(sourceColumn.fieldId());
      return bucket(sourceColumn, numBuckets, columnName + "_bucket");
    }

    public Builder truncate(String sourceName, int width, String targetName) {
      return truncate(findSourceColumn(sourceName), width, targetName);
    }

    private Builder truncate(Types.NestedField sourceColumn, int width, String targetName) {
      checkAndAddPartitionName(targetName);
      fields.add(
          new PartitionField(
              sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.truncate(width)));
      return this;
    }

    public Builder truncate(String sourceName, int width) {
      Types.NestedField sourceColumn = findSourceColumn(sourceName);
      String columnName = schema.findColumnName(sourceColumn.fieldId());
      return truncate(sourceColumn, width, columnName + "_trunc");
    }

    public Builder alwaysNull(String sourceName, String targetName) {
      return alwaysNull(findSourceColumn(sourceName), targetName);
    }

    private Builder alwaysNull(Types.NestedField sourceColumn, String targetName) {
      checkAndAddPartitionName(
          targetName, sourceColumn.fieldId()); // can duplicate a source column name
      fields.add(
          new PartitionField(
              sourceColumn.fieldId(), nextFieldId(), targetName, Transforms.alwaysNull()));
      return this;
    }

    public Builder alwaysNull(String sourceName) {
      Types.NestedField sourceColumn = findSourceColumn(sourceName);
      String columnName = schema.findColumnName(sourceColumn.fieldId());
      return alwaysNull(sourceColumn, columnName + "_null");
    }

    // add a partition field with an auto-increment partition field id starting from
    // PARTITION_DATA_ID_START
    Builder add(int sourceId, String name, Transform<?, ?> transform) {
      return add(sourceId, nextFieldId(), name, transform);
    }

    Builder add(int sourceId, int fieldId, String name, Transform<?, ?> transform) {
      checkAndAddPartitionName(name, sourceId);
      fields.add(new PartitionField(sourceId, fieldId, name, transform));
      lastAssignedFieldId.getAndAccumulate(fieldId, Math::max);
      return this;
    }

    public PartitionSpec build() {
      return build(false);
    }

    public PartitionSpec build(boolean allowMissingFields) {
      PartitionSpec spec = buildUnchecked();
      checkCompatibility(spec, schema, allowMissingFields);
      return spec;
    }

    PartitionSpec buildUnchecked() {
      return new PartitionSpec(schema, specId, fields, lastAssignedFieldId.get());
    }
  }

  static void checkCompatibility(PartitionSpec spec, Schema schema) {
    checkCompatibility(spec, schema, false);
  }

  static void checkCompatibility(PartitionSpec spec, Schema schema, boolean allowMissingFields) {
    final Map<Integer, Integer> parents = TypeUtil.indexParents(schema.asStruct());
    for (PartitionField field : spec.fields) {
      Type sourceType = schema.findType(field.sourceId());
      Transform<?, ?> transform = field.transform();
      // In the case the underlying field is dropped, we cannot check if they are compatible
      if (allowMissingFields && sourceType == null) {
        continue;
      }
      // In the case of a Version 1 partition-spec field gets deleted,
      // it is replaced with a void transform, see:
      // https://iceberg.apache.org/spec/#partition-transforms
      // We don't care about the source type since a VoidTransform is always compatible and skip the
      // checks
      if (!transform.equals(Transforms.alwaysNull())) {
        ValidationException.check(
            sourceType != null, "Cannot find source column for partition field: %s", field);
        ValidationException.check(
            sourceType.isPrimitiveType(),
            "Cannot partition by non-primitive source field: %s",
            sourceType);
        ValidationException.check(
            transform.canTransform(sourceType),
            "Invalid source type %s for transform: %s",
            sourceType,
            transform);
        // The only valid parent types for a PartitionField are StructTypes. This must be checked
        // recursively.
        Integer parentId = parents.get(field.sourceId());
        while (parentId != null) {
          Type parentType = schema.findType(parentId);
          ValidationException.check(
              parentType.isStructType(), "Invalid partition field parent: %s", parentType);
          parentId = parents.get(parentId);
        }
      }
    }
  }

  static boolean hasSequentialIds(PartitionSpec spec) {
    for (int i = 0; i < spec.fields.length; i += 1) {
      if (spec.fields[i].fieldId() != PARTITION_DATA_ID_START + i) {
        return false;
      }
    }
    return true;
  }
}
