/*
 * 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.flink.runtime.operators.sort;

import org.apache.flink.api.common.typeutils.TypeComparator;
import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.core.memory.MemorySegment;
import org.apache.flink.runtime.io.disk.iomanager.ChannelWriterOutputView;
import org.apache.flink.runtime.memory.AbstractPagedInputView;
import org.apache.flink.runtime.memory.AbstractPagedOutputView;
import org.apache.flink.util.MutableObjectIterator;

import java.io.EOFException;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/** */
public final class FixedLengthRecordSorter<T> implements InMemorySorter<T> {

    private static final int MIN_REQUIRED_BUFFERS = 3;

    // ------------------------------------------------------------------------
    //                               Members
    // ------------------------------------------------------------------------

    private final byte[] swapBuffer;

    private final TypeSerializer<T> serializer;

    private final TypeComparator<T> comparator;

    private final SingleSegmentOutputView outView;

    private final SingleSegmentInputView inView;

    private MemorySegment currentSortBufferSegment;

    private int currentSortBufferOffset;

    private final ArrayList<MemorySegment> freeMemory;

    private final ArrayList<MemorySegment> sortBuffer;

    private long sortBufferBytes;

    private int numRecords;

    private final int numKeyBytes;

    private final int recordSize;

    private final int recordsPerSegment;

    private final int lastEntryOffset;

    private final int segmentSize;

    private final int totalNumBuffers;

    private final boolean useNormKeyUninverted;

    private final T recordInstance;

    // -------------------------------------------------------------------------
    // Constructors / Destructors
    // -------------------------------------------------------------------------

    public FixedLengthRecordSorter(
            TypeSerializer<T> serializer,
            TypeComparator<T> comparator,
            List<MemorySegment> memory) {
        if (serializer == null || comparator == null || memory == null) {
            throw new NullPointerException();
        }

        this.serializer = serializer;
        this.comparator = comparator;
        this.useNormKeyUninverted = !comparator.invertNormalizedKey();

        // check the size of the first buffer and record it. all further buffers must have the same
        // size.
        // the size must also be a power of 2
        this.totalNumBuffers = memory.size();
        if (this.totalNumBuffers < MIN_REQUIRED_BUFFERS) {
            throw new IllegalArgumentException(
                    "Normalized-Key sorter requires at least "
                            + MIN_REQUIRED_BUFFERS
                            + " memory buffers.");
        }
        this.segmentSize = memory.get(0).size();
        this.recordSize = serializer.getLength();
        this.numKeyBytes = this.comparator.getNormalizeKeyLen();

        // check that the serializer and comparator allow our operations
        if (this.recordSize <= 0) {
            throw new IllegalArgumentException(
                    "This sorter works only for fixed-length data types.");
        } else if (this.recordSize > this.segmentSize) {
            throw new IllegalArgumentException(
                    "This sorter works only for record lengths below the memory segment size.");
        } else if (!comparator.supportsSerializationWithKeyNormalization()) {
            throw new IllegalArgumentException(
                    "This sorter requires a comparator that supports serialization with key normalization.");
        }

        // compute the entry size and limits
        this.recordsPerSegment = segmentSize / this.recordSize;
        this.lastEntryOffset = (this.recordsPerSegment - 1) * this.recordSize;
        this.swapBuffer = new byte[this.recordSize];

        this.freeMemory = new ArrayList<MemorySegment>(memory);

        // create the buffer collections
        this.sortBuffer = new ArrayList<MemorySegment>(16);
        this.outView = new SingleSegmentOutputView(this.segmentSize);
        this.inView = new SingleSegmentInputView(this.lastEntryOffset + this.recordSize);
        this.currentSortBufferSegment = nextMemorySegment();
        this.sortBuffer.add(this.currentSortBufferSegment);
        this.outView.set(this.currentSortBufferSegment);

        this.recordInstance = this.serializer.createInstance();
    }

    @Override
    public int recordSize() {
        return recordSize;
    }

    @Override
    public int recordsPerSegment() {
        return recordsPerSegment;
    }

    // -------------------------------------------------------------------------
    // Memory Segment
    // -------------------------------------------------------------------------

    /**
     * Resets the sort buffer back to the state where it is empty. All contained data is discarded.
     */
    @Override
    public void reset() {
        // reset all offsets
        this.numRecords = 0;
        this.currentSortBufferOffset = 0;
        this.sortBufferBytes = 0;

        // return all memory
        this.freeMemory.addAll(this.sortBuffer);
        this.sortBuffer.clear();

        // grab first buffers
        this.currentSortBufferSegment = nextMemorySegment();
        this.sortBuffer.add(this.currentSortBufferSegment);
        this.outView.set(this.currentSortBufferSegment);
    }

    /**
     * Checks whether the buffer is empty.
     *
     * @return True, if no record is contained, false otherwise.
     */
    @Override
    public boolean isEmpty() {
        return this.numRecords == 0;
    }

    @Override
    public void dispose() {
        this.freeMemory.clear();
        this.sortBuffer.clear();
    }

    @Override
    public long getCapacity() {
        return ((long) this.totalNumBuffers) * this.segmentSize;
    }

    @Override
    public long getOccupancy() {
        return this.sortBufferBytes;
    }

    // -------------------------------------------------------------------------
    // Retrieving and Writing
    // -------------------------------------------------------------------------

    @Override
    public T getRecord(int logicalPosition) throws IOException {
        return getRecord(serializer.createInstance(), logicalPosition);
    }

    @Override
    public T getRecord(T reuse, int logicalPosition) throws IOException {
        final int buffer = logicalPosition / this.recordsPerSegment;
        final int inBuffer = (logicalPosition % this.recordsPerSegment) * this.recordSize;
        this.inView.set(this.sortBuffer.get(buffer), inBuffer);
        return this.comparator.readWithKeyDenormalization(reuse, this.inView);
    }

    /**
     * Writes a given record to this sort buffer. The written record will be appended and take the
     * last logical position.
     *
     * @param record The record to be written.
     * @return True, if the record was successfully written, false, if the sort buffer was full.
     * @throws IOException Thrown, if an error occurred while serializing the record into the
     *     buffers.
     */
    @Override
    public boolean write(T record) throws IOException {
        // check whether we need a new memory segment for the sort index
        if (this.currentSortBufferOffset > this.lastEntryOffset) {
            if (memoryAvailable()) {
                this.currentSortBufferSegment = nextMemorySegment();
                this.sortBuffer.add(this.currentSortBufferSegment);
                this.outView.set(this.currentSortBufferSegment);
                this.currentSortBufferOffset = 0;
                this.sortBufferBytes += this.segmentSize;
            } else {
                return false;
            }
        }

        // serialize the record into the data buffers
        try {
            this.comparator.writeWithKeyNormalization(record, this.outView);
            this.numRecords++;
            this.currentSortBufferOffset += this.recordSize;
            return true;
        } catch (EOFException eofex) {
            throw new IOException(
                    "Error: Serialization consumes more bytes than announced by the serializer.");
        }
    }

    // ------------------------------------------------------------------------
    //                           Access Utilities
    // ------------------------------------------------------------------------

    private boolean memoryAvailable() {
        return !this.freeMemory.isEmpty();
    }

    private MemorySegment nextMemorySegment() {
        return this.freeMemory.remove(this.freeMemory.size() - 1);
    }

    // -------------------------------------------------------------------------
    // Sorting
    // -------------------------------------------------------------------------

    @Override
    public int compare(int i, int j) {
        final int segmentNumberI = i / this.recordsPerSegment;
        final int segmentOffsetI = (i % this.recordsPerSegment) * this.recordSize;

        final int segmentNumberJ = j / this.recordsPerSegment;
        final int segmentOffsetJ = (j % this.recordsPerSegment) * this.recordSize;

        return compare(segmentNumberI, segmentOffsetI, segmentNumberJ, segmentOffsetJ);
    }

    @Override
    public int compare(
            int segmentNumberI, int segmentOffsetI, int segmentNumberJ, int segmentOffsetJ) {
        final MemorySegment segI = this.sortBuffer.get(segmentNumberI);
        final MemorySegment segJ = this.sortBuffer.get(segmentNumberJ);

        int val = segI.compare(segJ, segmentOffsetI, segmentOffsetJ, this.numKeyBytes);
        return this.useNormKeyUninverted ? val : -val;
    }

    @Override
    public void swap(int i, int j) {
        final int segmentNumberI = i / this.recordsPerSegment;
        final int segmentOffsetI = (i % this.recordsPerSegment) * this.recordSize;

        final int segmentNumberJ = j / this.recordsPerSegment;
        final int segmentOffsetJ = (j % this.recordsPerSegment) * this.recordSize;

        swap(segmentNumberI, segmentOffsetI, segmentNumberJ, segmentOffsetJ);
    }

    @Override
    public void swap(
            int segmentNumberI, int segmentOffsetI, int segmentNumberJ, int segmentOffsetJ) {
        final MemorySegment segI = this.sortBuffer.get(segmentNumberI);
        final MemorySegment segJ = this.sortBuffer.get(segmentNumberJ);

        segI.swapBytes(this.swapBuffer, segJ, segmentOffsetI, segmentOffsetJ, this.recordSize);
    }

    @Override
    public int size() {
        return this.numRecords;
    }

    // -------------------------------------------------------------------------

    /**
     * Gets an iterator over all records in this buffer in their logical order.
     *
     * @return An iterator returning the records in their logical order.
     */
    @Override
    public final MutableObjectIterator<T> getIterator() {
        final SingleSegmentInputView startIn =
                new SingleSegmentInputView(this.recordsPerSegment * this.recordSize);
        startIn.set(this.sortBuffer.get(0), 0);

        return new MutableObjectIterator<T>() {

            private final SingleSegmentInputView in = startIn;
            private final TypeComparator<T> comp = comparator;

            private final int numTotal = size();
            private final int numPerSegment = recordsPerSegment;

            private int currentTotal = 0;
            private int currentInSegment = 0;
            private int currentSegmentIndex = 0;

            @Override
            public T next(T reuse) {
                if (this.currentTotal < this.numTotal) {

                    if (this.currentInSegment >= this.numPerSegment) {
                        this.currentInSegment = 0;
                        this.currentSegmentIndex++;
                        this.in.set(sortBuffer.get(this.currentSegmentIndex), 0);
                    }

                    this.currentTotal++;
                    this.currentInSegment++;

                    try {
                        return this.comp.readWithKeyDenormalization(reuse, this.in);
                    } catch (IOException ioe) {
                        throw new RuntimeException(ioe);
                    }
                } else {
                    return null;
                }
            }

            @Override
            public T next() {
                if (this.currentTotal < this.numTotal) {

                    if (this.currentInSegment >= this.numPerSegment) {
                        this.currentInSegment = 0;
                        this.currentSegmentIndex++;
                        this.in.set(sortBuffer.get(this.currentSegmentIndex), 0);
                    }

                    this.currentTotal++;
                    this.currentInSegment++;

                    try {
                        return this.comp.readWithKeyDenormalization(
                                serializer.createInstance(), this.in);
                    } catch (IOException ioe) {
                        throw new RuntimeException(ioe);
                    }
                } else {
                    return null;
                }
            }
        };
    }

    // ------------------------------------------------------------------------
    //                Writing to a DataOutputView
    // ------------------------------------------------------------------------

    /**
     * Writes the records in this buffer in their logical order to the given output.
     *
     * @param output The output view to write the records to.
     * @throws IOException Thrown, if an I/O exception occurred writing to the output view.
     */
    @Override
    public void writeToOutput(final ChannelWriterOutputView output) throws IOException {
        final TypeComparator<T> comparator = this.comparator;
        final TypeSerializer<T> serializer = this.serializer;
        T record = this.recordInstance;

        final SingleSegmentInputView inView = this.inView;

        final int recordsPerSegment = this.recordsPerSegment;
        int recordsLeft = this.numRecords;
        int currentMemSeg = 0;

        while (recordsLeft > 0) {
            final MemorySegment currentIndexSegment = this.sortBuffer.get(currentMemSeg++);
            inView.set(currentIndexSegment, 0);

            // check whether we have a full or partially full segment
            if (recordsLeft >= recordsPerSegment) {
                // full segment
                for (int numInMemSeg = 0; numInMemSeg < recordsPerSegment; numInMemSeg++) {
                    record = comparator.readWithKeyDenormalization(record, inView);
                    serializer.serialize(record, output);
                }
                recordsLeft -= recordsPerSegment;
            } else {
                // partially filled segment
                for (; recordsLeft > 0; recordsLeft--) {
                    record = comparator.readWithKeyDenormalization(record, inView);
                    serializer.serialize(record, output);
                }
            }
        }
    }

    @Override
    public void writeToOutput(
            ChannelWriterOutputView output, LargeRecordHandler<T> largeRecordsOutput)
            throws IOException {
        writeToOutput(output);
    }

    /**
     * Writes a subset of the records in this buffer in their logical order to the given output.
     *
     * @param output The output view to write the records to.
     * @param start The logical start position of the subset.
     * @param num The number of elements to write.
     * @throws IOException Thrown, if an I/O exception occurred writing to the output view.
     */
    @Override
    public void writeToOutput(final ChannelWriterOutputView output, final int start, int num)
            throws IOException {
        final TypeComparator<T> comparator = this.comparator;
        final TypeSerializer<T> serializer = this.serializer;
        T record = this.recordInstance;

        final SingleSegmentInputView inView = this.inView;

        final int recordsPerSegment = this.recordsPerSegment;
        int currentMemSeg = start / recordsPerSegment;
        int offset = (start % recordsPerSegment) * this.recordSize;

        while (num > 0) {
            final MemorySegment currentIndexSegment = this.sortBuffer.get(currentMemSeg++);
            inView.set(currentIndexSegment, offset);

            // check whether we have a full or partially full segment
            if (num >= recordsPerSegment && offset == 0) {
                // full segment
                for (int numInMemSeg = 0; numInMemSeg < recordsPerSegment; numInMemSeg++) {
                    record = comparator.readWithKeyDenormalization(record, inView);
                    serializer.serialize(record, output);
                }
                num -= recordsPerSegment;
            } else {
                // partially filled segment
                for (;
                        num > 0 && offset <= this.lastEntryOffset;
                        num--, offset += this.recordSize) {
                    record = comparator.readWithKeyDenormalization(record, inView);
                    serializer.serialize(record, output);
                }
            }

            offset = 0;
        }
    }

    private static final class SingleSegmentOutputView extends AbstractPagedOutputView {

        SingleSegmentOutputView(int segmentSize) {
            super(segmentSize, 0);
        }

        void set(MemorySegment segment) {
            seekOutput(segment, 0);
        }

        @Override
        protected MemorySegment nextSegment(MemorySegment current, int positionInCurrent)
                throws IOException {
            throw new EOFException();
        }
    }

    private static final class SingleSegmentInputView extends AbstractPagedInputView {

        private final int limit;

        SingleSegmentInputView(int limit) {
            super(0);
            this.limit = limit;
        }

        protected void set(MemorySegment segment, int offset) {
            seekInput(segment, offset, this.limit);
        }

        @Override
        protected MemorySegment nextSegment(MemorySegment current) throws EOFException {
            throw new EOFException();
        }

        @Override
        protected int getLimitForSegment(MemorySegment segment) {
            return this.limit;
        }
    }
}
