/*
 * 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 com.facebook.presto.operator.aggregation;

import com.facebook.presto.array.ObjectBigArray;
import com.facebook.presto.spi.block.Block;
import com.facebook.presto.spi.block.BlockBuilder;
import com.facebook.presto.spi.block.BlockBuilderStatus;
import com.facebook.presto.spi.block.InterleavedBlockBuilder;
import com.facebook.presto.spi.type.Type;
import com.facebook.presto.type.ArrayType;
import com.facebook.presto.type.RowType;
import com.google.common.collect.ImmutableList;
import org.openjdk.jol.info.ClassLayout;

import java.util.Optional;

import static com.facebook.presto.type.TypeUtils.expectedValueSize;
import static java.util.Objects.requireNonNull;

public class MultiKeyValuePairs
{
    private static final int INSTANCE_SIZE = ClassLayout.parseClass(MultiKeyValuePairs.class).instanceSize();
    private static final int EXPECTED_ENTRIES = 10;
    private static final int EXPECTED_ENTRY_SIZE = 16;

    private final BlockBuilder keyBlockBuilder;
    private final Type keyType;

    private final BlockBuilder valueBlockBuilder;
    private final Type valueType;

    private final RowType serializedRowType;

    public MultiKeyValuePairs(Type keyType, Type valueType)
    {
        this.keyType = requireNonNull(keyType, "keyType is null");
        this.valueType = requireNonNull(valueType, "valueType is null");
        keyBlockBuilder = this.keyType.createBlockBuilder(new BlockBuilderStatus(), EXPECTED_ENTRIES, expectedValueSize(keyType, EXPECTED_ENTRY_SIZE));
        valueBlockBuilder = this.valueType.createBlockBuilder(new BlockBuilderStatus(), EXPECTED_ENTRIES, expectedValueSize(valueType, EXPECTED_ENTRY_SIZE));
        serializedRowType = new RowType(ImmutableList.of(keyType, valueType), Optional.empty());
    }

    public MultiKeyValuePairs(Block serialized, Type keyType, Type valueType)
    {
        this(keyType, valueType);
        deserialize(requireNonNull(serialized, "serialized is null"));
    }

    public Block getKeys()
    {
        return keyBlockBuilder.build();
    }

    public Block getValues()
    {
        return valueBlockBuilder.build();
    }

    private void deserialize(Block block)
    {
        for (int i = 0; i < block.getPositionCount(); i++) {
            Block entryBlock = block.getObject(i, Block.class);
            add(entryBlock, entryBlock, 0, 1);
        }
    }

    public void serialize(BlockBuilder out)
    {
        BlockBuilder arrayBlockBuilder = out.beginBlockEntry();
        for (int i = 0; i < keyBlockBuilder.getPositionCount(); i++) {
            BlockBuilder rowBlockBuilder = arrayBlockBuilder.beginBlockEntry();
            keyType.appendTo(keyBlockBuilder, i, rowBlockBuilder);
            valueType.appendTo(valueBlockBuilder, i, rowBlockBuilder);
            arrayBlockBuilder.closeEntry();
        }
        out.closeEntry();
    }

    /**
     * Serialize as a multimap: map(key, array(value)), each key can be associated with multiple values
     */
    public Block toMultimapNativeEncoding()
    {
        Block keys = keyBlockBuilder.build();
        Block values = valueBlockBuilder.build();

        // Merge values of the same key into an array
        BlockBuilder distinctKeyBlockBuilder = keyType.createBlockBuilder(new BlockBuilderStatus(), keys.getPositionCount(), expectedValueSize(keyType, EXPECTED_ENTRY_SIZE));
        ObjectBigArray<BlockBuilder> valueArrayBlockBuilders = new ObjectBigArray<>();
        valueArrayBlockBuilders.ensureCapacity(keys.getPositionCount());
        TypedSet keySet = new TypedSet(keyType, keys.getPositionCount());
        for (int keyValueIndex = 0; keyValueIndex < keys.getPositionCount(); keyValueIndex++) {
            if (!keySet.contains(keys, keyValueIndex)) {
                keySet.add(keys, keyValueIndex);
                keyType.appendTo(keys, keyValueIndex, distinctKeyBlockBuilder);
                BlockBuilder valueArrayBuilder = valueType.createBlockBuilder(new BlockBuilderStatus(), 10, expectedValueSize(valueType, EXPECTED_ENTRY_SIZE));
                valueArrayBlockBuilders.set(keySet.positionOf(keys, keyValueIndex), valueArrayBuilder);
            }
            valueType.appendTo(values, keyValueIndex, valueArrayBlockBuilders.get(keySet.positionOf(keys, keyValueIndex)));
        }

        // Write keys and value arrays into one Block
        Block distinctKeys = distinctKeyBlockBuilder.build();
        Type valueArrayType = new ArrayType(valueType);
        BlockBuilder multimapBlockBuilder = new InterleavedBlockBuilder(ImmutableList.of(keyType, valueArrayType), new BlockBuilderStatus(), distinctKeyBlockBuilder.getPositionCount());
        for (int i = 0; i < distinctKeys.getPositionCount(); i++) {
            keyType.appendTo(distinctKeys, i, multimapBlockBuilder);
            valueArrayType.writeObject(multimapBlockBuilder, valueArrayBlockBuilders.get(i).build());
        }

        return multimapBlockBuilder.build();
    }

    public long estimatedInMemorySize()
    {
        long size = INSTANCE_SIZE;
        size += keyBlockBuilder.getRetainedSizeInBytes();
        size += valueBlockBuilder.getRetainedSizeInBytes();
        return size;
    }

    public void add(Block key, Block value, int keyPosition, int valuePosition)
    {
        keyType.appendTo(key, keyPosition, keyBlockBuilder);
        if (value.isNull(valuePosition)) {
            valueBlockBuilder.appendNull();
        }
        else {
            valueType.appendTo(value, valuePosition, valueBlockBuilder);
        }
    }
}