/*
 * 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.kylin.cache.memcached;

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

import org.apache.commons.lang3.SerializationUtils;
import org.apache.kylin.common.util.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Charsets;
import com.google.common.base.Function;
import com.google.common.base.Preconditions;
import com.google.common.base.Strings;
import com.google.common.base.Throwables;
import com.google.common.collect.Maps;
import com.google.common.primitives.Ints;
import com.google.common.primitives.Shorts;

import net.spy.memcached.internal.BulkFuture;

/**
 * Subclass of MemcachedCache. It supports storing large objects.  Memcached itself has a limitation to the value size with default value of 1M.
 * This implement extends the limit to 1G and can split huge bytes to multiple chunks. It will take care of the data integrity if part of the chunks lost(due to server restart or other reasons)
 *
 * @author mingmwang
 */
public class MemcachedChunkingCache extends MemcachedCache implements KeyHookLookup {
    private static final Logger logger = LoggerFactory.getLogger(MemcachedChunkingCache.class);

    public MemcachedChunkingCache(MemcachedCache cache) {
        super(cache);
        Preconditions.checkArgument(config.getMaxChunkSize() > 1, "maxChunkSize [%d] must be greater than 1",
                config.getMaxChunkSize());
        Preconditions.checkArgument(config.getMaxObjectSize() > 261, "maxObjectSize [%d] must be greater than 261",
                config.getMaxObjectSize());
    }

    protected static byte[][] splitBytes(final byte[] data, final int nSplit) {
        byte[][] dest = new byte[nSplit][];

        final int splitSize = (data.length - 1) / nSplit + 1;
        for (int i = 0; i < nSplit - 1; i++) {
            dest[i] = Arrays.copyOfRange(data, i * splitSize, (i + 1) * splitSize);
        }
        dest[nSplit - 1] = Arrays.copyOfRange(data, (nSplit - 1) * splitSize, data.length);

        return dest;
    }

    protected static int getValueSplit(MemcachedCacheConfig config, String keyS, int valueBLen) {
        // the number 6 means the chunk number size never exceeds 6 bytes
        final int valueSize = config.getMaxObjectSize() - Shorts.BYTES - Ints.BYTES
                - keyS.getBytes(Charsets.UTF_8).length - 6;
        final int maxValueSize = config.getMaxChunkSize() * valueSize;
        Preconditions.checkArgument(valueBLen <= maxValueSize,
                "the value bytes length [%d] exceeds maximum value size [%d]", valueBLen, maxValueSize);
        return (valueBLen - 1) / valueSize + 1;
    }

    protected static Pair<KeyHook, byte[][]> getKeyValuePair(int nSplit, String keyS, byte[] valueB) {
        KeyHook keyHook;
        byte[][] splitValueB = null;
        if (nSplit > 1) {
            if (logger.isDebugEnabled()) {
                logger.debug("Enable chunking for putting large cached object values, chunk size = " + nSplit
                        + ", original value bytes size = " + valueB.length);
            }
            String[] chunkKeySs = new String[nSplit];
            for (int i = 0; i < nSplit; i++) {
                chunkKeySs[i] = keyS + i;
            }
            keyHook = new KeyHook(chunkKeySs, null);
            splitValueB = splitBytes(valueB, nSplit);
        } else {
            if (logger.isDebugEnabled()) {
                logger.debug(
                        "Chunking not enabled, put the original value bytes to keyhook directly, original value bytes size = "
                                + valueB.length);
            }
            keyHook = new KeyHook(null, valueB);
        }

        return new Pair<>(keyHook, splitValueB);
    }

    /**
     * This method overrides the parent getBinary(), it gets the KeyHook from the Cache first and check the KeyHook that whether chunking is enabled or not.
     */
    @Override
    public byte[] getBinary(String keyS) {
        if (Strings.isNullOrEmpty(keyS)) {
            return null;
        }
        KeyHook keyHook = lookupKeyHook(keyS);
        if (keyHook == null) {
            return null;
        }

        if (keyHook.getChunkskey() == null || keyHook.getChunkskey().length == 0) {
            if (logger.isDebugEnabled()) {
                logger.debug("Chunking not enabled, return the value bytes in the keyhook directly, value bytes size = "
                        + keyHook.getValues().length);
            }
            return keyHook.getValues();
        }

        BulkFuture<Map<String, Object>> bulkFuture;
        long start = System.currentTimeMillis();

        if (logger.isDebugEnabled()) {
            logger.debug("Chunking enabled, chunk size = " + keyHook.getChunkskey().length);
        }

        Map<String, String> keyLookup = computeKeyHash(Arrays.asList(keyHook.getChunkskey()));
        try {
            bulkFuture = client.asyncGetBulk(keyLookup.keySet());
        } catch (IllegalStateException e) {
            // operation did not get queued in time (queue is full)
            errorCount.incrementAndGet();
            logger.error("Unable to queue cache operation.", e);
            return null;
        } catch (Throwable t) {
            errorCount.incrementAndGet();
            logger.error("Unable to queue cache operation.", t);
            return null;
        }

        try {
            Map<String, Object> bulkResult = bulkFuture.get(config.getTimeout(), TimeUnit.MILLISECONDS);
            cacheGetTime.addAndGet(System.currentTimeMillis() - start);
            if (bulkResult.size() != keyHook.getChunkskey().length) {
                missCount.incrementAndGet();
                logger.warn("Some paritial chunks missing for query key:" + keyS);
                //remove all the partital chunks here.
                for (String partitalKey : bulkResult.keySet()) {
                    client.delete(partitalKey);
                }
                deleteKeyHook(keyS);
                return null;
            }
            hitCount.getAndAdd(keyHook.getChunkskey().length);
            byte[][] bytesArray = new byte[keyHook.getChunkskey().length][];
            for (Map.Entry<String, Object> entry : bulkResult.entrySet()) {
                byte[] bytes = (byte[]) entry.getValue();
                readBytes.addAndGet(bytes.length);
                String originalKeyS = keyLookup.get(entry.getKey());
                int idx = Integer.parseInt(originalKeyS.substring(keyS.length()));
                bytesArray[idx] = decodeValue(originalKeyS.getBytes(Charsets.UTF_8), bytes);
            }
            return concatBytes(bytesArray);
        } catch (TimeoutException e) {
            timeoutCount.incrementAndGet();
            bulkFuture.cancel(false);
            return null;
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw Throwables.propagate(e);
        } catch (ExecutionException e) {
            errorCount.incrementAndGet();
            logger.error("ExecutionException when pulling item from cache.", e);
            return null;
        }
    }

    /**
     * This method overrides the parent putBinary() method. It will split the large value bytes into multiple chunks to fit into the internal Cache.
     * It generates a KeyHook to store the splitted chunked keys.
     */
    @Override
    public void putBinary(String keyS, byte[] valueB, int expiration) {
        if (Strings.isNullOrEmpty(keyS)) {
            return;
        }
        int nSplit = getValueSplit(config, keyS, valueB.length);
        Pair<KeyHook, byte[][]> keyValuePair = getKeyValuePair(nSplit, keyS, valueB);
        KeyHook keyHook = keyValuePair.getFirst();
        byte[][] splitValueB = keyValuePair.getSecond();

        if (logger.isDebugEnabled()) {
            logger.debug("put key hook:{} to cache for hash key", keyHook);
        }
        super.putBinary(keyS, serializeValue(keyHook), expiration);
        if (nSplit > 1) {
            for (int i = 0; i < nSplit; i++) {
                if (logger.isDebugEnabled()) {
                    logger.debug("Chunk[" + i + "] bytes size before encoding  = " + splitValueB[i].length);
                }
                super.putBinary(keyHook.getChunkskey()[i], splitValueB[i], expiration);
            }
        }
    }

    public void evict(String keyS) {
        if (Strings.isNullOrEmpty(keyS)) {
            return;
        }
        KeyHook keyHook = lookupKeyHook(keyS);
        if (keyHook == null) {
            return;
        }

        if (keyHook.getChunkskey() != null && keyHook.getChunkskey().length > 0) {
            String[] chunkKeys = keyHook.getChunkskey();
            for (String chunkKey : chunkKeys) {
                super.evict(chunkKey);
            }
        }
        super.evict(keyS);
    }

    protected Map<String, String> computeKeyHash(List<String> keySList) {
        return Maps.uniqueIndex(keySList, new Function<String, String>() {
            @Override
            public String apply(String keyS) {
                return computeKeyHash(keyS);
            }
        });
    }

    private void deleteKeyHook(String keyS) {
        try {
            super.evict(keyS);
        } catch (IllegalStateException e) {
            // operation did not get queued in time (queue is full)
            errorCount.incrementAndGet();
            logger.error("Unable to queue cache operation: ", e);
        }
    }

    private byte[] concatBytes(byte[]... bytesArray) {
        int length = 0;
        for (byte[] bytes : bytesArray) {
            length += bytes.length;
        }
        byte[] result = new byte[length];
        int destPos = 0;
        for (byte[] bytes : bytesArray) {
            System.arraycopy(bytes, 0, result, destPos, bytes.length);
            destPos += bytes.length;
        }
        if (logger.isDebugEnabled()) {
            logger.debug("Original value bytes size for all chunks  = " + result.length);
        }

        return result;
    }

    @Override
    public KeyHook lookupKeyHook(String keyS) {
        byte[] bytes = super.getBinary(keyS);
        if (bytes == null) {
            return null;
        }
        return (KeyHook) SerializationUtils.deserialize(bytes);
    }
}