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

import org.apache.paimon.CoreOptions;
import org.apache.paimon.data.GenericRow;
import org.apache.paimon.flink.FlinkRowWrapper;
import org.apache.paimon.flink.action.SortCompactAction;
import org.apache.paimon.sort.zorder.ZIndexer;
import org.apache.paimon.types.DataField;
import org.apache.paimon.types.DataTypes;
import org.apache.paimon.types.RowType;
import org.apache.paimon.utils.SortUtil;

import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.data.RowData;

import java.util.Arrays;
import java.util.Collections;

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

/**
 * This is a table sorter which will sort the records by the z-order of specified columns. It works
 * on stream api. It computes the z-order-index by {@link ZIndexer}. After add the column of
 * z-order, it does the range shuffle and sort. Finally, {@link SortCompactAction} will remove the
 * "z-order" column and insert sorted record to overwrite the origin table.
 */
public class ZorderSorter extends TableSorter {

    private static final RowType KEY_TYPE =
            new RowType(Collections.singletonList(new DataField(0, "Z_INDEX", DataTypes.BYTES())));

    private final TableSortInfo sortInfo;

    public ZorderSorter(
            StreamExecutionEnvironment batchTEnv,
            DataStream<RowData> origin,
            CoreOptions options,
            RowType rowType,
            TableSortInfo sortInfo) {
        super(batchTEnv, origin, options, rowType, sortInfo.getSortColumns());
        this.sortInfo = sortInfo;
    }

    @Override
    public DataStream<RowData> sort() {
        return sortStreamByZOrder(origin, options, rowType);
    }

    /**
     * Sort the input stream by the given order columns with z-order.
     *
     * @param inputStream the stream waited to be sorted
     * @return the sorted data stream
     */
    private DataStream<RowData> sortStreamByZOrder(
            DataStream<RowData> inputStream, CoreOptions options, RowType rowType) {
        final ZIndexer zIndexer = new ZIndexer(rowType, orderColNames, options.varTypeSize());
        return SortUtils.sortStreamByKey(
                inputStream,
                options,
                rowType,
                KEY_TYPE,
                TypeInformation.of(byte[].class),
                () ->
                        (b1, b2) -> {
                            checkArgument(
                                    b1.length == b2.length,
                                    "Z-order curve keys must have the same length, but got %d and %d",
                                    b1.length,
                                    b2.length);
                            return SortUtil.compareBinary(b1, b2);
                        },
                new SortUtils.KeyAbstract<byte[]>() {
                    @Override
                    public void open() {
                        zIndexer.open();
                    }

                    @Override
                    public byte[] apply(RowData value) {
                        byte[] zorder = zIndexer.index(new FlinkRowWrapper(value));
                        // we can just return the reused bytes zorder, because the sample operator
                        // will remember the record to sample.
                        return Arrays.copyOf(zorder, zorder.length);
                    }
                },
                GenericRow::of,
                sortInfo);
    }
}
