/*
 * 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.hyracks.storage.am.common;

import static org.junit.Assert.fail;

import java.io.ByteArrayInputStream;
import java.io.DataInput;
import java.io.DataInputStream;
import java.io.DataOutput;
import java.util.Collection;
import java.util.Iterator;
import java.util.Random;
import java.util.logging.Level;
import java.util.logging.Logger;

import org.apache.commons.lang3.tuple.MutablePair;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.hyracks.api.dataflow.value.*;
import org.apache.hyracks.api.exceptions.ErrorCode;
import org.apache.hyracks.api.exceptions.HyracksDataException;
import org.apache.hyracks.dataflow.common.comm.io.ArrayTupleBuilder;
import org.apache.hyracks.dataflow.common.comm.io.ArrayTupleReference;
import org.apache.hyracks.dataflow.common.data.accessors.ITupleReference;
import org.apache.hyracks.dataflow.common.utils.TupleUtils;
import org.apache.hyracks.storage.am.common.api.ITreeIndexAccessor;
import org.apache.hyracks.storage.am.common.api.ITreeIndexCursor;
import org.apache.hyracks.storage.common.IIndexBulkLoader;
import org.apache.hyracks.storage.common.ISearchPredicate;

@SuppressWarnings("rawtypes")
public abstract class TreeIndexTestUtils {
    private static final Logger LOGGER = Logger.getLogger(TreeIndexTestUtils.class.getName());

    protected abstract CheckTuple createCheckTuple(int numFields, int numKeyFields);

    protected abstract ISearchPredicate createNullSearchPredicate();

    public abstract void checkExpectedResults(ITreeIndexCursor cursor, Collection checkTuples,
            ISerializerDeserializer[] fieldSerdes, int keyFieldCount, Iterator<CheckTuple> checkIter) throws Exception;

    protected abstract CheckTuple createIntCheckTuple(int[] fieldValues, int numKeyFields);

    protected abstract void setIntKeyFields(int[] fieldValues, int numKeyFields, int maxValue, Random rnd);

    protected abstract void setIntPayloadFields(int[] fieldValues, int numKeyFields, int numFields);

    protected abstract Collection createCheckTuplesCollection();

    protected abstract ArrayTupleBuilder createDeleteTupleBuilder(IIndexTestContext ctx);

    // See if tuple with corresponding checkTuple exists in ctx.checkTuples.
    protected abstract boolean checkDiskOrderScanResult(ITupleReference tuple, CheckTuple checkTuple,
            IIndexTestContext ctx) throws HyracksDataException;

    public static int compareFilterTuples(ITupleReference lhs, ITupleReference rhs, IBinaryComparator comp)
            throws HyracksDataException {
        return comp.compare(lhs.getFieldData(0), lhs.getFieldStart(0), lhs.getFieldLength(0), rhs.getFieldData(0),
                rhs.getFieldStart(0), rhs.getFieldLength(0));
    }

    public static void createTupleFromCheckTuple(CheckTuple checkTuple, ArrayTupleBuilder tupleBuilder,
            ArrayTupleReference tuple, ISerializerDeserializer[] fieldSerdes) throws HyracksDataException {
        createTupleFromCheckTuple(checkTuple, tupleBuilder, tuple, fieldSerdes, false);
    }

    @SuppressWarnings("unchecked")
    public static void createTupleFromCheckTuple(CheckTuple checkTuple, ArrayTupleBuilder tupleBuilder,
            ArrayTupleReference tuple, ISerializerDeserializer[] fieldSerdes, boolean filtered)
            throws HyracksDataException {
        int fieldCount = tupleBuilder.getFieldEndOffsets().length;
        DataOutput dos = tupleBuilder.getDataOutput();
        tupleBuilder.reset();
        for (int i = 0; i < (filtered ? fieldCount - 1 : fieldCount); i++) {
            fieldSerdes[i].serialize(checkTuple.getField(i), dos);
            tupleBuilder.addFieldEndOffset();
        }
        if (filtered) {
            fieldSerdes[0].serialize(checkTuple.getField(0), dos);
            tupleBuilder.addFieldEndOffset();
        }
        tuple.reset(tupleBuilder.getFieldEndOffsets(), tupleBuilder.getByteArray());
    }

    @SuppressWarnings("unchecked")
    public CheckTuple createCheckTupleFromTuple(ITupleReference tuple, ISerializerDeserializer[] fieldSerdes,
            int numKeys) throws HyracksDataException {
        CheckTuple checkTuple = createCheckTuple(fieldSerdes.length, numKeys);
        int fieldCount = Math.min(fieldSerdes.length, tuple.getFieldCount());
        for (int i = 0; i < fieldCount; i++) {
            ByteArrayInputStream inStream = new ByteArrayInputStream(tuple.getFieldData(i), tuple.getFieldStart(i),
                    tuple.getFieldLength(i));
            DataInput dataIn = new DataInputStream(inStream);
            Comparable fieldObj = (Comparable) fieldSerdes[i].deserialize(dataIn);
            checkTuple.appendField(fieldObj);
        }
        return checkTuple;
    }

    @SuppressWarnings("unchecked")
    public void checkScan(IIndexTestContext ctx) throws Exception {
        if (LOGGER.isLoggable(Level.INFO)) {
            LOGGER.info("Testing Scan.");
        }
        ITreeIndexCursor scanCursor = (ITreeIndexCursor) ctx.getIndexAccessor().createSearchCursor(false);
        ISearchPredicate nullPred = createNullSearchPredicate();
        ctx.getIndexAccessor().search(scanCursor, nullPred);
        Iterator<CheckTuple> checkIter = ctx.getCheckTuples().iterator();
        checkExpectedResults(scanCursor, ctx.getCheckTuples(), ctx.getFieldSerdes(), ctx.getKeyFieldCount(), checkIter);
    }

    public void checkDiskOrderScan(IIndexTestContext ctx) throws Exception {
        try {
            if (LOGGER.isLoggable(Level.INFO)) {
                LOGGER.info("Testing Disk-Order Scan.");
            }
            ITreeIndexAccessor treeIndexAccessor = (ITreeIndexAccessor) ctx.getIndexAccessor();
            ITreeIndexCursor diskOrderCursor = treeIndexAccessor.createDiskOrderScanCursor();
            treeIndexAccessor.diskOrderScan(diskOrderCursor);
            int actualCount = 0;
            try {
                while (diskOrderCursor.hasNext()) {
                    diskOrderCursor.next();
                    ITupleReference tuple = diskOrderCursor.getTuple();
                    CheckTuple checkTuple = createCheckTupleFromTuple(tuple, ctx.getFieldSerdes(),
                            ctx.getKeyFieldCount());
                    if (!checkDiskOrderScanResult(tuple, checkTuple, ctx)) {
                        fail("Disk-order scan returned unexpected answer: " + checkTuple.toString());
                    }
                    actualCount++;
                }
                if (actualCount < ctx.getCheckTuples().size()) {
                    fail("Disk-order scan returned fewer answers than expected.\nExpected: "
                            + ctx.getCheckTuples().size() + "\nActual  : " + actualCount);
                }
                if (actualCount > ctx.getCheckTuples().size()) {
                    fail("Disk-order scan returned more answers than expected.\nExpected: "
                            + ctx.getCheckTuples().size() + "\nActual  : " + actualCount);
                }
            } finally {
                try {
                    diskOrderCursor.close();
                } catch (Exception ex) {
                    LOGGER.log(Level.WARNING, "Error during scan cursor close", ex);
                }
            }
        } catch (UnsupportedOperationException e) {
            // Ignore exception because some indexes, e.g. the LSMTrees, don't
            // support disk-order scan.
            if (LOGGER.isLoggable(Level.INFO)) {
                LOGGER.info("Ignoring disk-order scan since it's not supported.");
            }
        } catch (ClassCastException e) {
            // Ignore exception because IIndexAccessor sometimes isn't
            // an ITreeIndexAccessor, e.g., for the LSMBTree.
            if (LOGGER.isLoggable(Level.INFO)) {
                LOGGER.info("Ignoring disk-order scan since it's not supported.");
            }
        }
    }

    @SuppressWarnings("unchecked")

    public Pair<ITupleReference, ITupleReference> insertIntTuples(IIndexTestContext ctx, int numTuples, Random rnd)
            throws Exception {
        return insertIntTuples(ctx, numTuples, false, rnd);
    }

    public Pair<ITupleReference, ITupleReference> insertIntTuples(IIndexTestContext ctx, int numTuples,
            boolean filtered, Random rnd) throws Exception {
        int fieldCount = ctx.getFieldCount();
        int numKeyFields = ctx.getKeyFieldCount();
        int[] fieldValues = new int[ctx.getFieldCount()];
        // Scale range of values according to number of keys.
        // For example, for 2 keys we want the square root of numTuples, for 3
        // keys the cube root of numTuples, etc.
        int maxValue = (int) Math.ceil(Math.pow(numTuples, 1.0 / numKeyFields));
        MutablePair<ITupleReference, ITupleReference> minMax = null;
        for (int i = 0; i < numTuples; i++) {
            // Set keys.
            setIntKeyFields(fieldValues, numKeyFields, maxValue, rnd);
            // Set values.
            setIntPayloadFields(fieldValues, numKeyFields, fieldCount);
            TupleUtils.createIntegerTuple(ctx.getTupleBuilder(), ctx.getTuple(), filtered, fieldValues);
            if (LOGGER.isLoggable(Level.INFO)) {
                if ((i + 1) % (numTuples / Math.min(10, numTuples)) == 0) {
                    LOGGER.info("Inserting Tuple " + (i + 1) + "/" + numTuples);
                }
            }
            try {
                ctx.getIndexAccessor().insert(ctx.getTuple());
                ctx.insertCheckTuple(createIntCheckTuple(fieldValues, ctx.getKeyFieldCount()), ctx.getCheckTuples());
                if (filtered) {
                    addFilterField(ctx, minMax);
                }
            } catch (HyracksDataException e) {
                // We set expected values only after insertion succeeds because
                // we ignore duplicate keys.
                if (e.getErrorCode() != ErrorCode.DUPLICATE_KEY) {
                    throw e;
                }
            }
        }
        return minMax;
    }

    protected void addFilterField(IIndexTestContext ctx, MutablePair<ITupleReference, ITupleReference> minMax)
            throws HyracksDataException {
        //Duplicate the PK field as a filter field at the end of the tuple to be inserted.
        int filterField = ctx.getFieldCount();
        ITupleReference currTuple = ctx.getTuple();
        ArrayTupleBuilder filterBuilder = new ArrayTupleBuilder(1);
        filterBuilder.addField(currTuple.getFieldData(filterField), currTuple.getFieldStart(filterField),
                currTuple.getFieldLength(filterField));
        IBinaryComparator comparator = ctx.getComparatorFactories()[0].createBinaryComparator();
        ArrayTupleReference filterOnlyTuple = new ArrayTupleReference();
        filterOnlyTuple.reset(filterBuilder.getFieldEndOffsets(), filterBuilder.getByteArray());
        if (minMax == null) {
            minMax = MutablePair.of(filterOnlyTuple, filterOnlyTuple);
        } else if (compareFilterTuples(minMax.getLeft(), filterOnlyTuple, comparator) > 0) {
            minMax.setLeft(filterOnlyTuple);
        } else if (compareFilterTuples(minMax.getRight(), filterOnlyTuple, comparator) < 0) {
            minMax.setRight(filterOnlyTuple);
        }
    }

    @SuppressWarnings("unchecked")
    public void upsertIntTuples(IIndexTestContext ctx, int numTuples, Random rnd) throws Exception {
        int fieldCount = ctx.getFieldCount();
        int numKeyFields = ctx.getKeyFieldCount();
        int[] fieldValues = new int[ctx.getFieldCount()];
        // Scale range of values according to number of keys.
        // For example, for 2 keys we want the square root of numTuples, for 3
        // keys the cube root of numTuples, etc.
        int maxValue = (int) Math.ceil(Math.pow(numTuples, 1.0 / numKeyFields));
        for (int i = 0; i < numTuples; i++) {
            // Set keys.
            setIntKeyFields(fieldValues, numKeyFields, maxValue, rnd);
            // Set values.
            setIntPayloadFields(fieldValues, numKeyFields, fieldCount);
            TupleUtils.createIntegerTuple(ctx.getTupleBuilder(), ctx.getTuple(), fieldValues);
            if (LOGGER.isLoggable(Level.INFO)) {
                if ((i + 1) % (numTuples / Math.min(10, numTuples)) == 0) {
                    LOGGER.info("Inserting Tuple " + (i + 1) + "/" + numTuples);
                }
            }
            try {
                ctx.getIndexAccessor().upsert(ctx.getTuple());
                ctx.insertCheckTuple(createIntCheckTuple(fieldValues, ctx.getKeyFieldCount()), ctx.getCheckTuples());
            } catch (HyracksDataException e) {
                // We set expected values only after insertion succeeds because
                // we ignore duplicate keys.
                if (e.getErrorCode() != ErrorCode.DUPLICATE_KEY) {
                    throw e;
                }
            }
        }
    }

    public void bulkLoadIntTuples(IIndexTestContext ctx, int numTuples, Random rnd) throws Exception {
        bulkLoadIntTuples(ctx, numTuples, false, rnd);
    }

    @SuppressWarnings("unchecked")
    public void bulkLoadIntTuples(IIndexTestContext ctx, int numTuples, boolean filter, Random rnd) throws Exception {
        int fieldCount = ctx.getFieldCount();
        int numKeyFields = ctx.getKeyFieldCount();
        int[] fieldValues = new int[ctx.getFieldCount()];
        int maxValue = (int) Math.ceil(Math.pow(numTuples, 1.0 / numKeyFields));
        Collection<CheckTuple> tmpCheckTuples = createCheckTuplesCollection();
        for (int i = 0; i < numTuples; i++) {
            // Set keys.
            setIntKeyFields(fieldValues, numKeyFields, maxValue, rnd);
            // Set values.
            setIntPayloadFields(fieldValues, numKeyFields, fieldCount);

            // Set expected values. (We also use these as the pre-sorted stream
            // for ordered indexes bulk loading).
            ctx.insertCheckTuple(createIntCheckTuple(fieldValues, ctx.getKeyFieldCount()), tmpCheckTuples);
        }
        bulkLoadCheckTuples(ctx, tmpCheckTuples, filter);

        // Add tmpCheckTuples to ctx check tuples for comparing searches.
        for (CheckTuple checkTuple : tmpCheckTuples) {
            ctx.insertCheckTuple(checkTuple, ctx.getCheckTuples());
        }
    }

    public static void bulkLoadCheckTuples(IIndexTestContext ctx, Collection<CheckTuple> checkTuples)
            throws HyracksDataException {
        bulkLoadCheckTuples(ctx, checkTuples, false);
    }

    public static void bulkLoadCheckTuples(IIndexTestContext ctx, Collection<CheckTuple> checkTuples, boolean filtered)
            throws HyracksDataException {
        int fieldCount = ctx.getFieldCount();
        int numTuples = checkTuples.size();
        ArrayTupleBuilder tupleBuilder = filtered ? new ArrayTupleBuilder(fieldCount + 1)
                : new ArrayTupleBuilder(fieldCount);
        ArrayTupleReference tuple = new ArrayTupleReference();
        // Perform bulk load.
        IIndexBulkLoader bulkLoader = ctx.getIndex().createBulkLoader(0.7f, false, numTuples, false);
        int c = 1;
        for (CheckTuple checkTuple : checkTuples) {
            if (LOGGER.isLoggable(Level.INFO)) {
                //if (c % (numTuples / 10) == 0) {
                LOGGER.info("Bulk Loading Tuple " + c + "/" + numTuples);
                //}
            }
            createTupleFromCheckTuple(checkTuple, tupleBuilder, tuple, ctx.getFieldSerdes(), filtered);
            bulkLoader.add(tuple);
            c++;
        }
        bulkLoader.end();
    }

    @SuppressWarnings("unchecked")
    public void deleteTuples(IIndexTestContext ctx, int numTuples, Random rnd) throws Exception {
        ArrayTupleBuilder deleteTupleBuilder = createDeleteTupleBuilder(ctx);
        ArrayTupleReference deleteTuple = new ArrayTupleReference();
        int numCheckTuples = ctx.getCheckTuples().size();
        // Copy CheckTuple references into array, so we can randomly pick from
        // there.
        CheckTuple[] checkTuples = new CheckTuple[numCheckTuples];
        int idx = 0;
        Iterator<CheckTuple> iter = ctx.getCheckTuples().iterator();
        while (iter.hasNext()) {
            CheckTuple checkTuple = iter.next();
            checkTuples[idx++] = checkTuple;
        }

        for (int i = 0; i < numTuples && numCheckTuples > 0; i++) {
            if (LOGGER.isLoggable(Level.INFO)) {
                if ((i + 1) % (numTuples / Math.min(10, numTuples)) == 0) {
                    LOGGER.info("Deleting Tuple " + (i + 1) + "/" + numTuples);
                }
            }
            int checkTupleIdx = Math.abs(rnd.nextInt() % numCheckTuples);
            CheckTuple checkTuple = checkTuples[checkTupleIdx];
            createTupleFromCheckTuple(checkTuple, deleteTupleBuilder, deleteTuple, ctx.getFieldSerdes());
            ctx.getIndexAccessor().delete(deleteTuple);

            // Remove check tuple from expected results.
            ctx.deleteCheckTuple(checkTuple, ctx.getCheckTuples());

            // Swap with last "valid" CheckTuple.
            CheckTuple tmp = checkTuples[numCheckTuples - 1];
            checkTuples[numCheckTuples - 1] = checkTuple;
            checkTuples[checkTupleIdx] = tmp;
            numCheckTuples--;
        }
    }

}
