/*
 * 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.cassandra.service.accord.api;

import java.io.IOException;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

import com.google.common.annotations.VisibleForTesting;

import accord.api.RoutingKey;
import accord.local.ShardDistributor;
import accord.primitives.Range;
import accord.primitives.RangeFactory;
import accord.primitives.Ranges;
import accord.utils.Invariants;
import accord.utils.VIntCoding;

import org.apache.cassandra.db.marshal.ByteBufferAccessor;
import org.apache.cassandra.db.marshal.ValueAccessor;
import org.apache.cassandra.dht.IPartitioner;
import org.apache.cassandra.dht.Token;
import org.apache.cassandra.io.ParameterisedUnversionedSerializer;
import org.apache.cassandra.io.util.DataInputPlus;
import org.apache.cassandra.io.util.DataOutputPlus;
import org.apache.cassandra.schema.TableId;
import org.apache.cassandra.service.accord.TokenRange;
import org.apache.cassandra.utils.ObjectSizes;

import static org.apache.cassandra.config.DatabaseDescriptor.getPartitioner;

public final class TokenKey extends AccordRoutableKey implements RoutingKey, RangeFactory
{
    private static final long EMPTY_SIZE = ObjectSizes.measure(new TokenKey(null, null));

    @Override
    public Range asRange()
    {
        return TokenRange.create(before(), this);
    }

    // we use the first 2 bits as a prefix, and the last 6 bits as a postfix comparison
    final byte sentinel;
    final Token token;
    private TokenKey(TableId tableId, byte sentinel, Token token)
    {
        super(tableId);
        this.sentinel = sentinel;
        this.token = token;
    }

    public TokenKey(TableId tableId, Token token)
    {
        this(tableId, NORMAL_SENTINEL, token);
    }

    public TokenKey withToken(Token token)
    {
        return new TokenKey(table, sentinel, token);
    }

    @Override
    public Token token()
    {
        return token;
    }

    @Override
    public byte sentinel()
    {
        return sentinel;
    }

    public byte prefixSentinel()
    {
        return (byte) (sentinel & PREFIX_MASK);
    }

    public byte suffixSentinel()
    {
        return (byte) (sentinel & SUFFIX_MASK);
    }

    // this can be invoked to a depth of 3 from a real token
    @VisibleForTesting
    public TokenKey before()
    {
        Invariants.require(!isTokenSentinel(), "Unable to call .before() when already a token sentinel: %s", this);
        int lowestBit = Integer.lowestOneBit(sentinel);
        Invariants.require(lowestBit != 1);
        byte newSentinel = (byte)((sentinel ^ lowestBit) | (lowestBit >>> 1));
        return new TokenKey(table, newSentinel, token);
    }

    // this can be invoked to a depth of 2 from a real token
    @VisibleForTesting
    public TokenKey after()
    {
        Invariants.require(!isTokenSentinel(), "Unable to call .after() when already a token sentinel: %s", this);
        int lowestBit = Integer.lowestOneBit(sentinel);
        // we can't use 0xf as we would not be able to disambiguate with variable length byte encoding escape
        Invariants.require((lowestBit != 1) && (sentinel & 0xf) != 0xe);
        byte newSentinel = (byte)(sentinel | (lowestBit >>> 1));
        return new TokenKey(table, newSentinel, token);
    }

    @Override
    public Object suffix()
    {
        return token;
    }

    public Object printableSuffix()
    {
        Object suffix = suffix();
        if (isSentinel())
        {
            if (isTableSentinel()) suffix = isMin() ? "-Inf" : "+Inf";
            if (isTokenSentinel()) suffix = (isBefore() ? "before(" : "after(") + suffix + ')';
        }
        return suffix;
    }

    @Override
    public String toString()
    {
        return prefix() + ":" + printableSuffix();
    }

    public static TokenKey parse(String str, IPartitioner partitioner)
    {

        int split = str.indexOf(':', str.startsWith("tid:") ? 4 : 0);
        TableId tableId = TableId.fromString(str.substring(0, split));
        str = str.substring(split + 1);
        return parse(tableId, str, partitioner);
    }

    public static TokenKey parse(TableId tableId, String str, IPartitioner partitioner)
    {
        if (str.endsWith("Inf"))
        {
            return new TokenKey(tableId, str.charAt(0) == '-' ? MIN_TABLE_SENTINEL : MAX_TABLE_SENTINEL, partitioner.getMinimumToken());
        }
        if (str.endsWith(")"))
        {
            boolean isBefore = str.startsWith("before(");
            boolean isAfter = str.startsWith("after(");
            if (isBefore || isAfter)
            {
                byte sentinel = isBefore ? BEFORE_TOKEN_SENTINEL : AFTER_TOKEN_SENTINEL;
                str = str.substring(isBefore ? 7 : 6, str.length() - 1);
                return new TokenKey(tableId, sentinel, partitioner.getTokenFactory().fromString(str));
            }
        }
        return new TokenKey(tableId, partitioner.getTokenFactory().fromString(str));
    }

    public long estimatedSizeOnHeap()
    {
        return EMPTY_SIZE + token().getHeapSize();
    }

    public TokenKey withTable(TableId table)
    {
        return new TokenKey(table, sentinel, token);
    }

    @Override
    public RangeFactory rangeFactory()
    {
        return this;
    }

    @Override
    public Range newRange(RoutingKey start, RoutingKey end)
    {
        return TokenRange.create((TokenKey) start, (TokenKey) end);
    }

    @Override
    public Range newAntiRange(RoutingKey start, RoutingKey end)
    {
        return TokenRange.createUnsafe((TokenKey) start, (TokenKey) end);
    }

    @Override
    public RoutingKey toUnseekable()
    {
        return this;
    }

    public boolean isMin()
    {
        return (sentinel & PREFIX_MASK) == (MIN_TABLE_SENTINEL & PREFIX_MASK);
    }

    public boolean isMax()
    {
        return (sentinel & PREFIX_MASK) == (MAX_TABLE_SENTINEL & PREFIX_MASK);
    }

    public boolean isSentinel()
    {
        return sentinel != NORMAL_SENTINEL;
    }

    public boolean isTableSentinel()
    {
        return (sentinel & PREFIX_MASK) != (NORMAL_SENTINEL & PREFIX_MASK);
    }

    public boolean isTokenSentinel()
    {
        return (sentinel & SUFFIX_MASK) != (NORMAL_SENTINEL & SUFFIX_MASK);
    }

    public boolean isBefore()
    {
        return (sentinel & SUFFIX_MASK) == (BEFORE_TOKEN_SENTINEL & SUFFIX_MASK);
    }

    public boolean isAfter()
    {
        return (sentinel & SUFFIX_MASK) == (AFTER_TOKEN_SENTINEL & SUFFIX_MASK);
    }

    public static TokenKey min(TableId table, IPartitioner partitioner)
    {
        return new TokenKey(table, MIN_TABLE_SENTINEL, partitioner.getMinimumToken());
    }

    public static TokenKey max(TableId table, IPartitioner partitioner)
    {
        return new TokenKey(table, MAX_TABLE_SENTINEL, partitioner.getMinimumToken());
    }

    public static TokenKey before(TableId table, Token token)
    {
        return new TokenKey(table, BEFORE_TOKEN_SENTINEL, token);
    }

    public static TokenKey after(TableId table, Token token)
    {
        return new TokenKey(table, AFTER_TOKEN_SENTINEL, token);
    }

    public static final NoTableSerializer noTableSerializer = new NoTableSerializer();


    public static class NoTableSerializer implements ParameterisedUnversionedSerializer<TokenKey, TableId>
    {
        @Override
        public void serialize(TokenKey key, TableId tableId, DataOutputPlus out) throws IOException
        {
            IPartitioner partitioner = key.token.getPartitioner();
            int fixedLength = partitioner.accordFixedLength();
            if (fixedLength < 0)
            {
                int len = partitioner.accordSerializedSize(key.token);
                out.writeUnsignedVInt32(len);
            }
            serializer.serializeWithoutPrefixOrLength(key, out);
        }

        public void serialize(TokenKey key, DataOutputPlus out) throws IOException
        {
            serialize(key, key.table, out);
        }

        @Override
        public long serializedSize(TokenKey key, TableId tableId)
        {
            IPartitioner partitioner = key.token.getPartitioner();
            int tokenSize = partitioner.accordFixedLength();
            if (tokenSize >= 0)
                return 2 + tokenSize;
            tokenSize = partitioner.accordSerializedSize(key.token);
            return 2 + tokenSize + VIntCoding.sizeOfUnsignedVInt(tokenSize);
        }

        public long serializedSize(TokenKey key)
        {
            return serializedSize(key, key.table);
        }

        @Override
        public TokenKey deserialize(TableId tableId, DataInputPlus in) throws IOException
        {
            IPartitioner partitioner = getPartitioner();
            int len = partitioner.accordFixedLength();
            if (len < 0) len = in.readUnsignedVInt32();
            return serializer.deserializeWithPrefix(tableId, len + 2, in, partitioner);
        }
    }

    public static final class Serializer implements AccordSearchableKeySerializer<TokenKey>
    {
        private Serializer() {}

        // stream serialization methods - including a dynamic length for variable size tokens
        // types are byte comparable only after any length component

        @Override
        public long serializedSize(TokenKey key)
        {
            IPartitioner partitioner = key.token.getPartitioner();
            int size = 2 + key.table.serializedCompactComparableSize();
            int tokenSize = partitioner.accordFixedLength();
            if (tokenSize >= 0)
                return size + tokenSize;
            tokenSize = partitioner.accordSerializedSize(key.token);
            return size + tokenSize + VIntCoding.sizeOfUnsignedVInt(tokenSize);
        }

        @Override
        public void serialize(TokenKey key, DataOutputPlus out) throws IOException
        {
            IPartitioner partitioner = key.token.getPartitioner();
            int fixedLength = partitioner.accordFixedLength();
            if (fixedLength < 0)
            {
                int len = partitioner.accordSerializedSize(key.token);
                out.writeUnsignedVInt32(len);
            }
            key.table.serializeCompactComparable(out);
            serializeWithoutPrefixOrLength(key, out);
        }

        @Override
        public TokenKey deserialize(DataInputPlus in) throws IOException
        {
            return deserialize(in, getPartitioner());
        }

        public TokenKey deserialize(DataInputPlus in, IPartitioner partitioner) throws IOException
        {
            int len = partitioner.accordFixedLength();
            if (len < 0) len = in.readUnsignedVInt32();
            TableId tableId = deserializePrefix(in);
            return deserializeWithPrefix(tableId, len + 2, in, partitioner);
        }

        @Override
        public void skip(DataInputPlus in) throws IOException
        {
            skip(in, getPartitioner());
        }

        public void skip(DataInputPlus in, IPartitioner partitioner) throws IOException
        {
            int len = partitioner.accordFixedLength();
            if (len < 0) len = in.readUnsignedVInt32();
            TableId.skipCompact(in);
            in.skipBytesFully(len + 2);
        }

        // methods for encoding/decoding a single ByteBuffer value

        public ByteBuffer serialize(TokenKey key)
        {
            int size = key.table.serializedCompactComparableSize() + serializedSizeWithoutPrefixOrLength(key);
            ByteBuffer result = ByteBuffer.allocate(size);
            result.position(key.table.serializeCompactComparable(result, ByteBufferAccessor.instance, 0));
            serializeWithoutPrefixOrLength(key, result);
            result.flip();
            return result;
        }

        // WARNING: consumes buffer!
        public TokenKey deserialize(ByteBuffer buffer)
        {
            return deserialize(buffer, getPartitioner());
        }

        public TokenKey deserialize(ByteBuffer buffer, IPartitioner partitioner)
        {
            TableId tableId = TableId.deserializeCompactComparable(buffer, ByteBufferAccessor.instance, 0);
            int offset = tableId.serializedCompactComparableSize();
            return deserializeWithPrefix(tableId, buffer.remaining() - offset, buffer, ByteBufferAccessor.instance, offset, partitioner);
        }

        // WARNING: consumes buffer!
        public TokenKey deserializeAndConsume(ByteBuffer buffer, IPartitioner partitioner)
        {
            TableId tableId = TableId.deserializeCompactComparable(buffer, ByteBufferAccessor.instance, 0);
            int offset = buffer.position();
            buffer.position(offset + tableId.serializedCompactComparableSize());
            return deserializeWithPrefix(tableId, buffer.remaining(), buffer, partitioner);
        }

        // methods for encoding searchable tokens separately from tableIds

        @Override
        public int fixedKeyLengthForPrefix(Object prefix)
        {
            int size = getPartitioner().accordFixedLength();
            if (size < 0)
                return size;
            return 2 + size;
        }

        @Override
        public int serializedSizeWithoutPrefixOrLength(TokenKey key)
        {
            return 2 + key.token.getPartitioner().accordSerializedSize(key.token);
        }

        @Override
        public int serializedSizeOfPrefix(Object prefix)
        {
            return ((TableId) prefix).serializedCompactComparableSize();
        }

        @Override
        public void serializePrefix(Object prefix, DataOutputPlus out) throws IOException
        {
            ((TableId)prefix).serializeCompactComparable(out);
        }

        @Override
        public void serializeWithoutPrefixOrLength(TokenKey key, DataOutputPlus out) throws IOException
        {
            out.write(key.prefixSentinel());
            key.token.getPartitioner().accordSerialize(key.token, out);
            out.write(key.suffixSentinel());
        }

        public ByteBuffer serializeWithoutPrefixOrLength(TokenKey key)
        {
            IPartitioner partitioner = key.token.getPartitioner();
            ByteBuffer result = ByteBuffer.allocate(serializedSizeWithoutPrefixOrLength(key));
            serializeWithoutPrefixOrLength(key, result, partitioner);
            result.flip();
            return result;
        }

        public void serializeWithoutPrefixOrLength(TokenKey key, ByteBuffer out)
        {
            serializeWithoutPrefixOrLength(key, out, key.token.getPartitioner());
        }

        private static void serializeWithoutPrefixOrLength(TokenKey key, ByteBuffer out, IPartitioner partitioner)
        {
            out.put(key.prefixSentinel());
            partitioner.accordSerialize(key.token, out);
            out.put(key.suffixSentinel());
        }

        @Override
        public TableId deserializePrefix(DataInputPlus in) throws IOException
        {
            return TableId.deserializeCompactComparable(in);
        }

        @Override
        public TokenKey deserializeWithPrefix(Object tableId, int length, DataInputPlus in) throws IOException
        {
            return deserializeWithPrefix(tableId, length, in, getPartitioner());
        }

        public TokenKey deserializeWithPrefix(Object tableId, int length, DataInputPlus in, IPartitioner partitioner) throws IOException
        {
            byte sentinel = in.readByte();
            Token token = partitioner.accordDeserialize(in, length - 2);
            sentinel |= in.readByte();
            return new TokenKey((TableId) tableId, sentinel, token);
        }

        public <V> TokenKey deserializeWithPrefixAndImpliedLength(Object tableId, V src, ValueAccessor<V> accessor, int offset)
        {
            return deserializeWithPrefixAndImpliedLength(tableId, src, accessor, offset, getPartitioner());
        }

        public <V> TokenKey deserializeWithPrefixAndImpliedLength(Object tableId, V src, ValueAccessor<V> accessor, int offset, IPartitioner partitioner)
        {
            return deserializeWithPrefix(tableId, accessor.remaining(src, offset), src, accessor, offset, partitioner);
        }

        public <V> TokenKey deserializeWithPrefix(Object tableId, int length, V src, ValueAccessor<V> accessor, int offset, IPartitioner partitioner)
        {
            byte sentinel = accessor.getByte(src, offset++);
            Token token = partitioner.accordDeserialize(src, accessor, offset, length - 2);
            offset += partitioner.accordSerializedSize(token);
            sentinel |= accessor.getByte(src, offset);
            return new TokenKey((TableId) tableId, sentinel, token);
        }

        // WARNING: consumes buffer!
        public TokenKey deserializeWithPrefixAndImpliedLength(Object tableId, ByteBuffer buffer, IPartitioner partitioner)
        {
            return deserializeWithPrefix(tableId, buffer.remaining(), buffer, partitioner);
        }

        // WARNING: consumes buffer!
        public TokenKey deserializeWithPrefix(Object tableId, int length, ByteBuffer buffer, IPartitioner partitioner)
        {
            byte sentinel = buffer.get();
            Token token = partitioner.accordDeserialize(buffer, length - 2);
            sentinel |= buffer.get();
            return new TokenKey((TableId) tableId, sentinel, token);
        }

        public static final byte ESCAPE_BYTE = 0x0f;
        private static final byte[] ESCAPE_BYTES = new byte[] { ESCAPE_BYTE };
        private static final int UNESCAPE = ESCAPE_BYTE;
        private static final int UNESCAPE_MASK = 0xffff;

        public static int countEscapes(byte[] bytes)
        {
            int escapeLimit = escapeLimit(bytes);
            int i = 0;
            int count = 0;
            while ((i = nextEscape(bytes, i, escapeLimit)) >= 0)
            {
                ++count;
                ++i;
            }
            return count;
        }

        public static int serializedSize(byte[] bytes)
        {
            return 1 + bytes.length + countEscapes(bytes);
        }

        private static int escapeLimit(byte[] bytes)
        {
            return bytes.length - 1;
        }

        private static int nextEscape(byte[] bytes, int index, int escapeLimit)
        {
            while (index <= escapeLimit)
            {
                if (bytes[index] == 0 && (index == escapeLimit || (bytes[index + 1] & 0xff) <= ESCAPE_BYTE))
                    return index;
                ++index;
            }
            return -1;
        }

        public static void serializeWithEscapes(byte[] bytes, ByteBuffer out)
        {
            serializeWithEscapesInternal(bytes, out, ByteBuffer::put);
            out.put(trailingByte(bytes));
        }

        public static void serializeWithEscapes(byte[] bytes, DataOutputPlus out) throws IOException
        {
            serializeWithEscapesInternal(bytes, out, DataOutputPlus::write);
            out.writeByte(trailingByte(bytes));
        }

        interface WriteBytes<V, T extends Throwable>
        {
            void write(V out, byte[] bytes, int offset, int length) throws T;
        }

        private static byte trailingByte(byte[] bytes)
        {
            return 0;
        }

        private static <V, T extends Throwable> void serializeWithEscapesInternal(byte[] bytes, V out, WriteBytes<V, T> write) throws T
        {
            int i = 0, escapeLimit = escapeLimit(bytes);
            while (true)
            {
                int nexti = nextEscape(bytes, i, escapeLimit);
                if (nexti < 0)
                    break;
                write.write(out, bytes, i, 1 + nexti - i);
                write.write(out, ESCAPE_BYTES, 0, 1);
                i = nexti + 1;
            }
            write.write(out, bytes, i, bytes.length - i);
        }

        public static byte[] deserializeWithEscapes(ByteBuffer in, int escapedLength)
        {
            Invariants.require(escapedLength >= 1);
            --escapedLength;
            byte[] bytes = new byte[escapedLength];
            in.get(bytes, 0, Math.min(escapedLength, in.remaining()));
            byte[] result = removeEscapes(bytes);
            byte trailingEscape = in.get();
            Invariants.require(trailingEscape == trailingByte(result));
            return result;
        }

        public static <V> byte[] deserializeWithEscapes(V src, ValueAccessor<V> accessor, int offset, int escapedLength)
        {
            Invariants.require(--escapedLength >= 0);
            byte[] result = removeEscapes(accessor.toArray(src, offset, Math.min(escapedLength, accessor.remaining(src, offset))));
            Invariants.require(trailingByte(result) == accessor.getByte(src, offset + escapedLength));
            return result;
        }

        public static byte[] deserializeWithEscapes(DataInputPlus in, int escapedLength) throws IOException
        {
            byte[] result = new byte[escapedLength - 1];
            in.readFully(result);
            result = removeEscapes(result);
            byte trailingEscape = in.readByte();
            Invariants.require(trailingEscape == trailingByte(result));
            return result;
        }

        private static byte[] removeEscapes(byte[] bytes)
        {
            if (bytes.length == 0)
                return bytes;

            int count = 1;
            int escapeMatcher = bytes[0];
            for (int i = 1; i < bytes.length ; ++i)
            {
                byte next = bytes[i];
                escapeMatcher = (escapeMatcher << 8) | next;
                if ((escapeMatcher & UNESCAPE_MASK) != UNESCAPE)
                {
                    if (count != i)
                        bytes[count] = next;
                    count++;
                }
            }

            if (bytes.length != count)
                bytes = Arrays.copyOf(bytes, count);
            return bytes;
        }

    }

    public static final Serializer serializer = new Serializer();

    public static class KeyspaceSplitter implements ShardDistributor
    {
        final EvenSplit<BigInteger> subSplitter;
        public KeyspaceSplitter(EvenSplit<BigInteger> subSplitter)
        {
            this.subSplitter = subSplitter;
        }

        @Override
        public List<Ranges> split(Ranges ranges)
        {
            Map<TableId, List<Range>> byTable = new TreeMap<>();
            for (Range range : ranges)
            {
                byTable.computeIfAbsent(((AccordRoutableKey)range.start()).table, ignore -> new ArrayList<>())
                          .add(range);
            }

            List<Ranges> results = new ArrayList<>();
            for (List<Range> keyspaceRanges : byTable.values())
                results.addAll(subSplitter.split(Ranges.ofSortedAndDeoverlapped(keyspaceRanges.toArray(new Range[0]))));
            return results;
        }

        @Override
        public Range splitRange(Range range, int from, int to, int numSplits)
        {
            return subSplitter.splitRange(range, from, to, numSplits);
        }

        @Override
        public Ranges selectFirstSubRanges(Range range, Ranges subRanges, int totalSplits)
        {
            return subSplitter.selectFirstSubRanges(range, subRanges, totalSplits);
        }

        @Override
        public int numberOfSplitsPossible(Range range)
        {
            return subSplitter.numberOfSplitsPossible(range);
        }
    }
}
