/*
 * 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.paimon.flink.procedure;

import org.apache.paimon.flink.globalindex.GenericIndexTopoBuilder;
import org.apache.paimon.flink.globalindex.SortedIndexTopoBuilder;
import org.apache.paimon.globalindex.GlobalIndexer;
import org.apache.paimon.options.Options;
import org.apache.paimon.partition.PartitionPredicate;
import org.apache.paimon.predicate.Predicate;
import org.apache.paimon.table.FileStoreTable;
import org.apache.paimon.types.DataField;
import org.apache.paimon.types.RowType;
import org.apache.paimon.utils.ParameterUtils;

import org.apache.flink.table.annotation.ArgumentHint;
import org.apache.flink.table.annotation.DataTypeHint;
import org.apache.flink.table.annotation.ProcedureHint;
import org.apache.flink.table.procedure.ProcedureContext;

import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

import static org.apache.paimon.utils.ParameterUtils.getPartitions;
import static org.apache.paimon.utils.Preconditions.checkArgument;

/** Procedure to create global index files via Flink. */
public class CreateGlobalIndexProcedure extends ProcedureBase {

    public static final String IDENTIFIER = "create_global_index";

    static Options createUserOptions(FileStoreTable table, String optionString) {
        return createUserOptions(table.options(), optionString);
    }

    static Options createUserOptions(Map<String, String> tableOptions, String optionString) {
        return new Options(tableOptions, optionalConfigMap(optionString));
    }

    @Override
    public String identifier() {
        return IDENTIFIER;
    }

    @ProcedureHint(
            argument = {
                @ArgumentHint(name = "table", type = @DataTypeHint("STRING")),
                @ArgumentHint(name = "index_column", type = @DataTypeHint("STRING")),
                @ArgumentHint(name = "index_type", type = @DataTypeHint("STRING")),
                @ArgumentHint(
                        name = "partitions",
                        type = @DataTypeHint("STRING"),
                        isOptional = true),
                @ArgumentHint(name = "options", type = @DataTypeHint("STRING"), isOptional = true)
            })
    public String[] call(
            ProcedureContext procedureContext,
            String tableId,
            String indexColumn,
            String indexType,
            String partitions,
            String options)
            throws Exception {

        FileStoreTable table = (FileStoreTable) table(tableId);

        // Validate table configuration
        checkArgument(
                table.coreOptions().rowTrackingEnabled(),
                "Table '%s' must enable 'row-tracking.enabled=true' before creating global index.",
                tableId);

        RowType rowType = table.rowType();
        List<String> indexColumns =
                Arrays.stream(indexColumn.split(","))
                        .map(String::trim)
                        .filter(s -> !s.isEmpty())
                        .collect(Collectors.toList());
        checkArgument(!indexColumns.isEmpty(), "At least one column required.");
        checkArgument(
                indexColumns.size() == new HashSet<>(indexColumns).size(),
                "Duplicate index columns are not allowed: %s",
                indexColumns);
        for (String col : indexColumns) {
            checkArgument(
                    rowType.containsField(col),
                    "Column '%s' does not exist in table '%s'.",
                    col,
                    tableId);
        }

        // Parse partition predicate
        PartitionPredicate partitionPredicate = parsePartitionPredicate(table, partitions);

        // Parse options
        Options userOptions = createUserOptions(table, options);

        indexType = indexType.toLowerCase().trim();
        if (indexColumns.size() > 1) {
            // Fail fast before submitting the job: index types that do not support multi-column
            // throw from GlobalIndexerFactory#create, which happens before any indexer side effect.
            DataField indexField = rowType.getField(indexColumns.get(0));
            List<DataField> extraFields =
                    indexColumns.subList(1, indexColumns.size()).stream()
                            .map(rowType::getField)
                            .collect(Collectors.toList());
            try {
                GlobalIndexer.create(indexType, indexField, extraFields, userOptions);
            } catch (UnsupportedOperationException e) {
                throw new IllegalArgumentException(
                        String.format(
                                "Index type '%s' does not support multi-column index, got columns: %s",
                                indexType, indexColumns));
            }
        }
        try {
            if (SortedIndexTopoBuilder.supports(indexType)) {
                SortedIndexTopoBuilder.buildIndexAndExecute(
                        procedureContext.getExecutionEnvironment(),
                        table,
                        indexColumns.get(0),
                        indexType,
                        partitionPredicate,
                        userOptions);
                return new String[] {
                    displayIndexType(indexType)
                            + " global index created successfully for table: "
                            + table.name()
                };
            } else {
                GenericIndexTopoBuilder.buildIndexAndExecute(
                        procedureContext.getExecutionEnvironment(),
                        table,
                        indexColumns.get(0),
                        indexColumns.subList(1, indexColumns.size()),
                        indexType,
                        partitionPredicate,
                        userOptions);
            }
        } catch (Exception e) {
            throw new RuntimeException(
                    String.format(
                            "Failed to create %s index for columns '%s' on table '%s'.",
                            indexType, indexColumns, table.name()),
                    e);
        }
        return new String[] {
            displayIndexType(indexType)
                    + " global index created successfully for table: "
                    + table.name()
        };
    }

    private PartitionPredicate parsePartitionPredicate(FileStoreTable table, String partitions) {
        if (partitions == null || partitions.isEmpty()) {
            return null;
        }

        List<Map<String, String>> partitionList = getPartitions(partitions.split(";"));
        Predicate predicate =
                ParameterUtils.toPartitionPredicate(
                        partitionList,
                        table.schema().logicalPartitionType(),
                        table.coreOptions().partitionDefaultName());
        return PartitionPredicate.fromPredicate(table.schema().logicalPartitionType(), predicate);
    }

    private static String displayIndexType(String indexType) {
        if ("btree".equals(indexType)) {
            return "BTree";
        }
        if ("bitmap".equals(indexType)) {
            return "Bitmap";
        }
        return indexType;
    }
}
