/*
 * Licensed 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 io.prestosql.plugin.hive.metastore.glue.converter;

import com.amazonaws.services.glue.model.DatabaseInput;
import com.amazonaws.services.glue.model.PartitionInput;
import com.amazonaws.services.glue.model.SerDeInfo;
import com.amazonaws.services.glue.model.StorageDescriptor;
import com.amazonaws.services.glue.model.TableInput;
import com.google.common.collect.ImmutableMap;
import io.prestosql.plugin.hive.PartitionStatistics;
import io.prestosql.plugin.hive.metastore.Column;
import io.prestosql.plugin.hive.metastore.Database;
import io.prestosql.plugin.hive.metastore.Partition;
import io.prestosql.plugin.hive.metastore.PartitionWithStatistics;
import io.prestosql.plugin.hive.metastore.Storage;
import io.prestosql.plugin.hive.metastore.Table;
import io.prestosql.plugin.hive.metastore.glue.GlueColumnStatisticsProvider;

import java.util.List;

import static com.google.common.collect.ImmutableList.toImmutableList;
import static io.prestosql.plugin.hive.metastore.thrift.ThriftMetastoreUtil.updateStatisticsParameters;

public final class GlueInputConverter
{
    private GlueInputConverter() {}

    public static DatabaseInput convertDatabase(Database database)
    {
        DatabaseInput input = new DatabaseInput();
        input.setName(database.getDatabaseName());
        input.setParameters(database.getParameters());
        database.getComment().ifPresent(input::setDescription);
        database.getLocation().ifPresent(input::setLocationUri);
        return input;
    }

    public static TableInput convertTable(Table table)
    {
        TableInput input = new TableInput();
        input.setName(table.getTableName());
        input.setOwner(table.getOwner());
        input.setTableType(table.getTableType());
        input.setStorageDescriptor(convertStorage(table.getStorage(), table.getDataColumns()));
        input.setPartitionKeys(table.getPartitionColumns().stream().map(GlueInputConverter::convertColumn).collect(toImmutableList()));
        input.setParameters(table.getParameters());
        table.getViewOriginalText().ifPresent(input::setViewOriginalText);
        table.getViewExpandedText().ifPresent(input::setViewExpandedText);
        return input;
    }

    public static PartitionInput convertPartition(PartitionWithStatistics partitionWithStatistics, GlueColumnStatisticsProvider columnStatisticsProvider)
    {
        PartitionInput input = convertPartition(partitionWithStatistics.getPartition());
        PartitionStatistics statistics = partitionWithStatistics.getStatistics();
        columnStatisticsProvider.updatePartitionStatistics(input, statistics.getColumnStatistics());
        input.setParameters(updateStatisticsParameters(input.getParameters(), statistics.getBasicStatistics()));
        return input;
    }

    public static PartitionInput convertPartition(Partition partition)
    {
        PartitionInput input = new PartitionInput();
        input.setValues(partition.getValues());
        input.setStorageDescriptor(convertStorage(partition.getStorage(), partition.getColumns()));
        input.setParameters(partition.getParameters());
        return input;
    }

    private static StorageDescriptor convertStorage(Storage storage, List<Column> columns)
    {
        if (storage.isSkewed()) {
            throw new IllegalArgumentException("Writing to skewed table/partition is not supported");
        }
        SerDeInfo serdeInfo = new SerDeInfo()
                .withSerializationLibrary(storage.getStorageFormat().getSerDeNullable())
                .withParameters(storage.getSerdeParameters());

        StorageDescriptor sd = new StorageDescriptor();
        sd.setLocation(storage.getLocation());
        sd.setColumns(columns.stream().map(GlueInputConverter::convertColumn).collect(toImmutableList()));
        sd.setSerdeInfo(serdeInfo);
        sd.setInputFormat(storage.getStorageFormat().getInputFormatNullable());
        sd.setOutputFormat(storage.getStorageFormat().getOutputFormatNullable());
        sd.setParameters(ImmutableMap.of());

        if (storage.getBucketProperty().isPresent()) {
            sd.setNumberOfBuckets(storage.getBucketProperty().get().getBucketCount());
            sd.setBucketColumns(storage.getBucketProperty().get().getBucketedBy());
        }

        return sd;
    }

    private static com.amazonaws.services.glue.model.Column convertColumn(Column prestoColumn)
    {
        return new com.amazonaws.services.glue.model.Column()
                .withName(prestoColumn.getName())
                .withType(prestoColumn.getType().toString())
                .withComment(prestoColumn.getComment().orElse(null));
    }
}

@@ -663,32 +663,80 @@
 .DatabaseInput;%0A
+import com.amazonaws.services.glue.model.Order;%0A
 import com.amazo
@@ -988,16 +988,68 @@
 bleMap;%0A
+import io.prestosql.plugin.hive.HiveBucketProperty;%0A
 import i
@@ -1515,16 +1515,43 @@
 il.List;
+%0Aimport java.util.Optional;
 %0A%0Aimport
@@ -4736,39 +4736,109 @@
     
-if (storage.getBucketProperty()
+Optional%3CHiveBucketProperty%3E bucketProperty = storage.getBucketProperty();%0A        if (bucketProperty
 .isP
@@ -4883,28 +4883,17 @@
 Buckets(
-storage.getB
+b
 ucketPro
@@ -4889,34 +4889,32 @@
 s(bucketProperty
-()
 .get().getBucket
@@ -4955,28 +4955,73 @@
 Columns(
-storage.getB
+bucketProperty.get().getBucketedBy());%0A            if (!b
 ucketPro
@@ -5029,34 +5029,322 @@
 erty
+.get
 ().get
-().getBucketedBy());
+SortedBy().isEmpty()) %7B%0A                sd.setSortColumns(bucketProperty.get().getSortedBy().stream()%0A                        .map(column -%3E new Order().withColumn(column.getColumnName()).withSortOrder(column.getOrder().getHiveOrder()))%0A                        .collect(toImmutableList()));%0A            %7D
 %0A   

Add missing support for bucket sort order in Glue