/*
 * 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.gravitino.catalog.glue;

import static org.apache.gravitino.catalog.glue.GlueConstants.INPUT_FORMAT_CLASS;
import static org.apache.gravitino.catalog.glue.GlueConstants.LOCATION;
import static org.apache.gravitino.catalog.glue.GlueConstants.OUTPUT_FORMAT;
import static org.apache.gravitino.catalog.glue.GlueConstants.SERDE_LIB;
import static org.apache.gravitino.catalog.glue.GlueConstants.SERDE_NAME;
import static org.apache.gravitino.catalog.glue.GlueConstants.SERDE_PARAMETER_PREFIX;

import com.google.common.base.Preconditions;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import lombok.ToString;
import org.apache.commons.lang3.StringUtils;
import org.apache.gravitino.connector.BaseTable;
import org.apache.gravitino.connector.TableOperations;
import org.apache.gravitino.meta.AuditInfo;
import org.apache.gravitino.rel.Column;
import org.apache.gravitino.rel.SupportsPartitions;
import org.apache.gravitino.rel.expressions.NamedReference;
import org.apache.gravitino.rel.expressions.distributions.Distribution;
import org.apache.gravitino.rel.expressions.distributions.Distributions;
import org.apache.gravitino.rel.expressions.sorts.SortDirection;
import org.apache.gravitino.rel.expressions.sorts.SortOrder;
import org.apache.gravitino.rel.expressions.sorts.SortOrders;
import org.apache.gravitino.rel.expressions.transforms.Transform;
import org.apache.gravitino.rel.expressions.transforms.Transforms;
import software.amazon.awssdk.services.glue.GlueClient;
import software.amazon.awssdk.services.glue.model.StorageDescriptor;
import software.amazon.awssdk.services.glue.model.Table;

/**
 * Represents an AWS Glue {@link Table} as a Gravitino table.
 *
 * <p>All entries in {@code Table.parameters()} pass through intact (including {@code table_type},
 * {@code metadata_location}, etc.), so downstream tools can correctly identify the table format.
 * StorageDescriptor fields (location, formats, SerDe) are surfaced as additional properties.
 */
@ToString
public class GlueTable extends BaseTable {

  /** Glue client injected after construction to support partition operations. */
  private GlueClient glueClient;

  /** Nullable — when null, Glue uses the caller's AWS account ID. */
  private String catalogId;

  /** Database (schema) this table belongs to. */
  private String dbName;

  private GlueTable() {}

  /**
   * Injects the Glue connection context required for partition operations.
   *
   * <p>Called by {@link GlueCatalogOperations} after constructing a table instance.
   */
  void initOpsContext(GlueClient glueClient, String catalogId, String dbName) {
    this.glueClient = glueClient;
    this.catalogId = catalogId;
    this.dbName = dbName;
  }

  @Override
  protected TableOperations newOps() {
    Preconditions.checkState(
        glueClient != null,
        "Partition operations require Glue context; call initOpsContext() first");
    String[] partColNames =
        Arrays.stream(partitioning)
            .filter(t -> t instanceof Transforms.IdentityTransform)
            .map(t -> ((Transforms.IdentityTransform) t).fieldName()[0])
            .toArray(String[]::new);
    return new GlueTableOperations(glueClient, catalogId, dbName, name, partColNames);
  }

  @Override
  public SupportsPartitions supportPartitions() throws UnsupportedOperationException {
    return (SupportsPartitions) ops();
  }

  /**
   * Converts an AWS Glue {@link Table} to a {@link GlueTable}.
   *
   * <p>Column assembly:
   *
   * <ol>
   *   <li>Data columns from {@code storageDescriptor.columns()} (Hive-format tables).
   *   <li>Partition columns from {@code table.partitionKeys()} appended after data columns.
   * </ol>
   *
   * <p>For Iceberg-format tables the StorageDescriptor columns are typically empty; all metadata
   * (including {@code table_type=ICEBERG} and {@code metadata_location}) is in {@code
   * table.parameters()} and passes through as-is.
   *
   * @param glueTable the Glue Table returned by the AWS SDK
   * @param typeConverter the type converter to use for column type mapping
   * @return a populated {@link GlueTable}
   */
  public static GlueTable fromGlueTable(Table glueTable, GlueTypeConverter typeConverter) {
    StorageDescriptor sd = glueTable.storageDescriptor();

    // --- Columns ---
    List<Column> columns = new ArrayList<>();
    if (sd != null && sd.hasColumns()) {
      for (software.amazon.awssdk.services.glue.model.Column c : sd.columns()) {
        columns.add(GlueColumn.fromGlueColumn(c, typeConverter));
      }
    }
    List<String> partitionColNames = new ArrayList<>();
    if (glueTable.hasPartitionKeys()) {
      for (software.amazon.awssdk.services.glue.model.Column pk : glueTable.partitionKeys()) {
        columns.add(GlueColumn.fromGlueColumn(pk, typeConverter));
        partitionColNames.add(pk.name());
      }
    }

    // --- Partitioning ---
    Transform[] partitioning =
        partitionColNames.stream().map(Transforms::identity).toArray(Transform[]::new);

    // --- Distribution (bucket) ---
    Distribution distribution = Distributions.NONE;
    Integer numBuckets = sd != null ? sd.numberOfBuckets() : null;
    if (sd != null && sd.hasBucketColumns() && numBuckets != null && numBuckets > 0) {
      distribution =
          Distributions.hash(
              numBuckets,
              sd.bucketColumns().stream()
                  .map(NamedReference::field)
                  .toArray(org.apache.gravitino.rel.expressions.Expression[]::new));
    }

    // --- Sort orders ---
    SortOrder[] sortOrders = SortOrders.NONE;
    if (sd != null && sd.hasSortColumns()) {
      sortOrders =
          sd.sortColumns().stream()
              .map(
                  o -> {
                    Integer sortOrder = o.sortOrder();
                    return SortOrders.of(
                        NamedReference.field(o.column()),
                        sortOrder != null && sortOrder == 1
                            ? SortDirection.ASCENDING
                            : SortDirection.DESCENDING);
                  })
              .toArray(SortOrder[]::new);
    }

    // --- Properties ---
    Map<String, String> properties = new HashMap<>();
    if (glueTable.hasParameters()) {
      properties.putAll(glueTable.parameters());
    }

    if (sd != null) {
      putIfNotBlank(properties, LOCATION, sd.location());
      putIfNotBlank(properties, INPUT_FORMAT_CLASS, sd.inputFormat());
      putIfNotBlank(properties, OUTPUT_FORMAT, sd.outputFormat());
      if (sd.serdeInfo() != null) {
        putIfNotBlank(properties, SERDE_LIB, sd.serdeInfo().serializationLibrary());
        putIfNotBlank(properties, SERDE_NAME, sd.serdeInfo().name());
        if (sd.serdeInfo().parameters() != null) {
          sd.serdeInfo()
              .parameters()
              .forEach((k, v) -> properties.put(SERDE_PARAMETER_PREFIX + k, v));
        }
      }
    }

    // --- AuditInfo ---
    AuditInfo auditInfo =
        AuditInfo.builder()
            .withCreateTime(glueTable.createTime())
            .withLastModifiedTime(glueTable.updateTime())
            .build();

    return GlueTable.builder()
        .withName(glueTable.name())
        .withComment(glueTable.description())
        .withColumns(columns.toArray(new Column[0]))
        .withProperties(properties)
        .withPartitioning(partitioning)
        .withDistribution(distribution)
        .withSortOrders(sortOrders)
        .withAuditInfo(auditInfo)
        .build();
  }

  private static void putIfNotBlank(Map<String, String> map, String key, String value) {
    if (StringUtils.isNotBlank(value)) {
      map.put(key, value);
    }
  }

  /** Builder for {@link GlueTable}. */
  public static class Builder extends BaseTableBuilder<Builder, GlueTable> {

    private Builder() {}

    @Override
    protected GlueTable internalBuild() {
      GlueTable table = new GlueTable();
      table.name = name;
      table.comment = comment;
      table.columns = columns;
      table.properties = properties;
      table.partitioning = partitioning;
      table.sortOrders = sortOrders;
      table.distribution = distribution;
      table.indexes = indexes;
      table.auditInfo = auditInfo;
      table.proxyPlugin = Optional.empty();
      return table;
    }
  }

  /**
   * Creates a new {@link Builder}.
   *
   * @return a new builder instance
   */
  public static Builder builder() {
    return new Builder();
  }

  /**
   * Sets the partitioning of the table. Package-private: only {@link GlueIcebergTableHelper} should
   * call this to enrich a table loaded from Glue with Iceberg-specific partition metadata.
   *
   * @param partitioning the partitioning transforms
   */
  void setPartitioning(Transform[] partitioning) {
    this.partitioning = partitioning;
  }

  /**
   * Sets the sort orders of the table. Package-private: only {@link GlueIcebergTableHelper} should
   * call this to enrich a table loaded from Glue with Iceberg-specific sort order metadata.
   *
   * @param sortOrders the sort orders
   */
  void setSortOrders(SortOrder[] sortOrders) {
    this.sortOrders = sortOrders;
  }

  /**
   * Sets the properties of the table. Package-private: only {@link GlueIcebergTableHelper} should
   * call this to merge Iceberg metadata properties into a table loaded from Glue.
   *
   * @param properties the properties map
   */
  void setProperties(Map<String, String> properties) {
    this.properties = properties;
  }

  /**
   * Replaces the columns with types sourced from the Iceberg schema, which is authoritative for
   * Iceberg tables. Glue stores some types (e.g. TIME) as "string", so this corrects the loss.
   *
   * @param columns the columns derived from the Iceberg table schema
   */
  void setColumns(Column[] columns) {
    this.columns = columns;
  }
}
