/*
 * 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 com.datastax.driver.core;

import com.datastax.driver.core.utils.Bytes;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.Lists;
import com.google.common.primitives.UnsignedBytes;
import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.List;

/** A token on the Cassandra ring. */
public abstract class Token implements Comparable<Token> {

  /**
   * Returns the data type of this token's value.
   *
   * @return the datatype.
   */
  public abstract DataType getType();

  /**
   * Returns the raw value of this token.
   *
   * @return the value.
   */
  public abstract Object getValue();

  /**
   * Returns the serialized form of the current token, using the appropriate codec depending on the
   * partitioner in use and the CQL datatype for the token.
   *
   * @param protocolVersion the protocol version in use.
   * @return the serialized form of the current token
   */
  public abstract ByteBuffer serialize(ProtocolVersion protocolVersion);

  static Token.Factory getFactory(String partitionerName) {
    if (partitionerName.endsWith("Murmur3Partitioner")) return M3PToken.FACTORY;
    else if (partitionerName.endsWith("RandomPartitioner")) return RPToken.FACTORY;
    else if (partitionerName.endsWith("OrderedPartitioner")) return OPPToken.FACTORY;
    else return null;
  }

  abstract static class Factory {
    abstract Token fromString(String tokenStr);

    abstract DataType getTokenType();

    abstract Token deserialize(ByteBuffer buffer, ProtocolVersion protocolVersion);

    /**
     * The minimum token is a special value that no key ever hashes to, it's used both as lower and
     * upper bound.
     */
    abstract Token minToken();

    abstract Token hash(ByteBuffer partitionKey);

    abstract List<Token> split(Token startToken, Token endToken, int numberOfSplits);

    // Base implementation for split
    protected List<BigInteger> split(
        BigInteger start,
        BigInteger range,
        BigInteger ringEnd,
        BigInteger ringLength,
        int numberOfSplits) {
      BigInteger[] tmp = range.divideAndRemainder(BigInteger.valueOf(numberOfSplits));
      BigInteger divider = tmp[0];
      int remainder = tmp[1].intValue();

      List<BigInteger> results = Lists.newArrayListWithExpectedSize(numberOfSplits - 1);
      BigInteger current = start;
      BigInteger dividerPlusOne =
          (remainder == 0)
              ? null // won't be used
              : divider.add(BigInteger.ONE);

      for (int i = 1; i < numberOfSplits; i++) {
        current = current.add(remainder-- > 0 ? dividerPlusOne : divider);
        if (ringEnd != null && current.compareTo(ringEnd) > 0)
          current = current.subtract(ringLength);
        results.add(current);
      }
      return results;
    }
  }

  // Murmur3Partitioner tokens
  static class M3PToken extends Token {
    private final long value;

    public static final Factory FACTORY = new M3PTokenFactory();

    private static class M3PTokenFactory extends Factory {

      private static final BigInteger RING_END = BigInteger.valueOf(Long.MAX_VALUE);
      private static final BigInteger RING_LENGTH =
          RING_END.subtract(BigInteger.valueOf(Long.MIN_VALUE));
      static final M3PToken MIN_TOKEN = new M3PToken(Long.MIN_VALUE);
      static final M3PToken MAX_TOKEN = new M3PToken(Long.MAX_VALUE);

      private long getblock(ByteBuffer key, int offset, int index) {
        int i_8 = index << 3;
        int blockOffset = offset + i_8;
        return ((long) key.get(blockOffset + 0) & 0xff)
            + (((long) key.get(blockOffset + 1) & 0xff) << 8)
            + (((long) key.get(blockOffset + 2) & 0xff) << 16)
            + (((long) key.get(blockOffset + 3) & 0xff) << 24)
            + (((long) key.get(blockOffset + 4) & 0xff) << 32)
            + (((long) key.get(blockOffset + 5) & 0xff) << 40)
            + (((long) key.get(blockOffset + 6) & 0xff) << 48)
            + (((long) key.get(blockOffset + 7) & 0xff) << 56);
      }

      private long rotl64(long v, int n) {
        return ((v << n) | (v >>> (64 - n)));
      }

      private long fmix(long k) {
        k ^= k >>> 33;
        k *= 0xff51afd7ed558ccdL;
        k ^= k >>> 33;
        k *= 0xc4ceb9fe1a85ec53L;
        k ^= k >>> 33;
        return k;
      }

      // This is an adapted version of the MurmurHash.hash3_x64_128 from Cassandra used
      // for M3P. Compared to that methods, there's a few inlining of arguments and we
      // only return the first 64-bits of the result since that's all M3P uses.
      @SuppressWarnings("fallthrough")
      private long murmur(ByteBuffer data) {
        int offset = data.position();
        int length = data.remaining();

        int nblocks = length >> 4; // Process as 128-bit blocks.

        long h1 = 0;
        long h2 = 0;

        long c1 = 0x87c37b91114253d5L;
        long c2 = 0x4cf5ad432745937fL;

        // ----------
        // body

        for (int i = 0; i < nblocks; i++) {
          long k1 = getblock(data, offset, i * 2 + 0);
          long k2 = getblock(data, offset, i * 2 + 1);

          k1 *= c1;
          k1 = rotl64(k1, 31);
          k1 *= c2;
          h1 ^= k1;
          h1 = rotl64(h1, 27);
          h1 += h2;
          h1 = h1 * 5 + 0x52dce729;
          k2 *= c2;
          k2 = rotl64(k2, 33);
          k2 *= c1;
          h2 ^= k2;
          h2 = rotl64(h2, 31);
          h2 += h1;
          h2 = h2 * 5 + 0x38495ab5;
        }

        // ----------
        // tail

        // Advance offset to the unprocessed tail of the data.
        offset += nblocks * 16;

        long k1 = 0;
        long k2 = 0;

        switch (length & 15) {
          case 15:
            k2 ^= ((long) data.get(offset + 14)) << 48;
          case 14:
            k2 ^= ((long) data.get(offset + 13)) << 40;
          case 13:
            k2 ^= ((long) data.get(offset + 12)) << 32;
          case 12:
            k2 ^= ((long) data.get(offset + 11)) << 24;
          case 11:
            k2 ^= ((long) data.get(offset + 10)) << 16;
          case 10:
            k2 ^= ((long) data.get(offset + 9)) << 8;
          case 9:
            k2 ^= ((long) data.get(offset + 8)) << 0;
            k2 *= c2;
            k2 = rotl64(k2, 33);
            k2 *= c1;
            h2 ^= k2;

          case 8:
            k1 ^= ((long) data.get(offset + 7)) << 56;
          case 7:
            k1 ^= ((long) data.get(offset + 6)) << 48;
          case 6:
            k1 ^= ((long) data.get(offset + 5)) << 40;
          case 5:
            k1 ^= ((long) data.get(offset + 4)) << 32;
          case 4:
            k1 ^= ((long) data.get(offset + 3)) << 24;
          case 3:
            k1 ^= ((long) data.get(offset + 2)) << 16;
          case 2:
            k1 ^= ((long) data.get(offset + 1)) << 8;
          case 1:
            k1 ^= ((long) data.get(offset));
            k1 *= c1;
            k1 = rotl64(k1, 31);
            k1 *= c2;
            h1 ^= k1;
        }

        // ----------
        // finalization

        h1 ^= length;
        h2 ^= length;

        h1 += h2;
        h2 += h1;

        h1 = fmix(h1);
        h2 = fmix(h2);

        h1 += h2;
        h2 += h1;

        return h1;
      }

      @Override
      M3PToken fromString(String tokenStr) {
        return new M3PToken(Long.parseLong(tokenStr));
      }

      @Override
      DataType getTokenType() {
        return DataType.bigint();
      }

      @Override
      Token deserialize(ByteBuffer buffer, ProtocolVersion protocolVersion) {
        return new M3PToken(TypeCodec.bigint().deserialize(buffer, protocolVersion));
      }

      @Override
      Token minToken() {
        return MIN_TOKEN;
      }

      @Override
      M3PToken hash(ByteBuffer partitionKey) {
        long v = murmur(partitionKey);
        return new M3PToken(v == Long.MIN_VALUE ? Long.MAX_VALUE : v);
      }

      @Override
      List<Token> split(Token startToken, Token endToken, int numberOfSplits) {
        // edge case: ]min, min] means the whole ring
        if (startToken.equals(endToken) && startToken.equals(MIN_TOKEN)) endToken = MAX_TOKEN;

        BigInteger start = BigInteger.valueOf(((M3PToken) startToken).value);
        BigInteger end = BigInteger.valueOf(((M3PToken) endToken).value);

        BigInteger range = end.subtract(start);
        if (range.compareTo(BigInteger.ZERO) < 0) range = range.add(RING_LENGTH);

        List<BigInteger> values = super.split(start, range, RING_END, RING_LENGTH, numberOfSplits);
        List<Token> tokens = Lists.newArrayListWithExpectedSize(values.size());
        for (BigInteger value : values) tokens.add(new M3PToken(value.longValue()));
        return tokens;
      }
    }

    private M3PToken(long value) {
      this.value = value;
    }

    @Override
    public DataType getType() {
      return FACTORY.getTokenType();
    }

    @Override
    public Object getValue() {
      return value;
    }

    @Override
    public ByteBuffer serialize(ProtocolVersion protocolVersion) {
      return TypeCodec.bigint().serialize(value, protocolVersion);
    }

    @Override
    public int compareTo(Token other) {
      assert other instanceof M3PToken;
      long otherValue = ((M3PToken) other).value;
      return value < otherValue ? -1 : (value == otherValue) ? 0 : 1;
    }

    @Override
    public boolean equals(Object obj) {
      if (this == obj) return true;
      if (obj == null || this.getClass() != obj.getClass()) return false;

      return value == ((M3PToken) obj).value;
    }

    @Override
    public int hashCode() {
      return (int) (value ^ (value >>> 32));
    }

    @Override
    public String toString() {
      return Long.toString(value);
    }
  }

  // OPPartitioner tokens
  static class OPPToken extends Token {

    private final ByteBuffer value;

    public static final Factory FACTORY = new OPPTokenFactory();

    private static class OPPTokenFactory extends Factory {
      private static final BigInteger TWO = BigInteger.valueOf(2);
      private static final Token MIN_TOKEN = new OPPToken(ByteBuffer.allocate(0));

      @Override
      public OPPToken fromString(String tokenStr) {
        // This method must be able to parse the contents of system.peers.tokens, which do not have
        // the "0x" prefix.
        // On the other hand, OPPToken#toString has the "0x" because it should be usable in a CQL
        // query, and it's
        // nice to have fromString and toString symetrical.
        // So handle both cases:
        if (!tokenStr.startsWith("0x")) {
          String prefix = (tokenStr.length() % 2 == 0) ? "0x" : "0x0";
          tokenStr = prefix + tokenStr;
        }
        ByteBuffer value = Bytes.fromHexString(tokenStr);
        return new OPPToken(value);
      }

      @Override
      DataType getTokenType() {
        return DataType.blob();
      }

      @Override
      Token deserialize(ByteBuffer buffer, ProtocolVersion protocolVersion) {
        return new OPPToken(buffer);
      }

      @Override
      Token minToken() {
        return MIN_TOKEN;
      }

      @Override
      OPPToken hash(ByteBuffer partitionKey) {
        return new OPPToken(partitionKey);
      }

      @Override
      List<Token> split(Token startToken, Token endToken, int numberOfSplits) {
        int tokenOrder = startToken.compareTo(endToken);

        // ]min,min] means the whole ring. However, since there is no "max token" with this
        // partitioner, we can't come up
        // with a magic end value that would cover the whole ring
        if (tokenOrder == 0 && startToken.equals(MIN_TOKEN))
          throw new IllegalArgumentException("Cannot split whole ring with ordered partitioner");

        OPPToken oppStartToken = (OPPToken) startToken;
        OPPToken oppEndToken = (OPPToken) endToken;

        int significantBytes;
        BigInteger start, end, range, ringEnd, ringLength;
        BigInteger bigNumberOfSplits = BigInteger.valueOf(numberOfSplits);
        if (tokenOrder < 0) {
          // Since tokens are compared lexicographically, convert to integers using the largest
          // length
          // (ex: given 0x0A and 0x0BCD, switch to 0x0A00 and 0x0BCD)
          significantBytes = Math.max(oppStartToken.value.capacity(), oppEndToken.value.capacity());

          // If the number of splits does not fit in the difference between the two integers, use
          // more bytes
          // (ex: cannot fit 4 splits between 0x01 and 0x03, so switch to 0x0100 and 0x0300)
          // At most 4 additional bytes will be needed, since numberOfSplits is an integer.
          int addedBytes = 0;
          while (true) {
            start = toBigInteger(oppStartToken.value, significantBytes);
            end = toBigInteger(oppEndToken.value, significantBytes);
            range = end.subtract(start);
            if (addedBytes == 4 || start.equals(end) || range.compareTo(bigNumberOfSplits) >= 0)
              break;
            significantBytes += 1;
            addedBytes += 1;
          }
          ringEnd = ringLength = null; // won't be used
        } else {
          // Same logic except that we wrap around the ring
          significantBytes = Math.max(oppStartToken.value.capacity(), oppEndToken.value.capacity());
          int addedBytes = 0;
          while (true) {
            start = toBigInteger(oppStartToken.value, significantBytes);
            end = toBigInteger(oppEndToken.value, significantBytes);
            ringLength = TWO.pow(significantBytes * 8);
            ringEnd = ringLength.subtract(BigInteger.ONE);
            range = end.subtract(start).add(ringLength);
            if (addedBytes == 4 || range.compareTo(bigNumberOfSplits) >= 0) break;
            significantBytes += 1;
            addedBytes += 1;
          }
        }

        List<BigInteger> values = super.split(start, range, ringEnd, ringLength, numberOfSplits);
        List<Token> tokens = Lists.newArrayListWithExpectedSize(values.size());
        for (BigInteger value : values) tokens.add(new OPPToken(toBytes(value, significantBytes)));
        return tokens;
      }

      // Convert a token's byte array to a number in order to perform computations.
      // This depends on the number of "significant bytes" that we use to normalize all tokens to
      // the same size.
      // For example if the token is 0x01 but significantBytes is 2, the result is 8 (0x0100).
      private BigInteger toBigInteger(ByteBuffer bb, int significantBytes) {
        byte[] bytes = Bytes.getArray(bb);
        byte[] target;
        if (significantBytes != bytes.length) {
          target = new byte[significantBytes];
          System.arraycopy(bytes, 0, target, 0, bytes.length);
        } else target = bytes;
        return new BigInteger(1, target);
      }

      // Convert a numeric representation back to a byte array.
      // Again, the number of significant bytes matters: if the input value is 1 but
      // significantBytes is 2, the
      // expected result is 0x0001 (a simple conversion would produce 0x01).
      protected ByteBuffer toBytes(BigInteger value, int significantBytes) {
        byte[] rawBytes = value.toByteArray();
        byte[] result;
        if (rawBytes.length == significantBytes) result = rawBytes;
        else {
          result = new byte[significantBytes];
          int start, length;
          if (rawBytes[0] == 0) { // that's a sign byte, ignore (it can cause rawBytes.length ==
            // significantBytes + 1)
            start = 1;
            length = rawBytes.length - 1;
          } else {
            start = 0;
            length = rawBytes.length;
          }
          System.arraycopy(rawBytes, start, result, significantBytes - length, length);
        }
        return ByteBuffer.wrap(result);
      }
    }

    @VisibleForTesting
    OPPToken(ByteBuffer value) {
      this.value = value;
    }

    @Override
    public DataType getType() {
      return FACTORY.getTokenType();
    }

    @Override
    public Object getValue() {
      return value;
    }

    @Override
    public ByteBuffer serialize(ProtocolVersion protocolVersion) {
      return TypeCodec.blob().serialize(value, protocolVersion);
    }

    @Override
    public int compareTo(Token other) {
      assert other instanceof OPPToken;
      return UnsignedBytes.lexicographicalComparator()
          .compare(Bytes.getArray(value), Bytes.getArray(((OPPToken) other).value));
    }

    @Override
    public boolean equals(Object obj) {
      if (this == obj) return true;
      if (obj == null || this.getClass() != obj.getClass()) return false;

      return value.equals(((OPPToken) obj).value);
    }

    @Override
    public int hashCode() {
      return value.hashCode();
    }

    @Override
    public String toString() {
      return Bytes.toHexString(value);
    }
  }

  // RandomPartitioner tokens
  static class RPToken extends Token {

    private final BigInteger value;

    public static final Factory FACTORY = new RPTokenFactory();

    private static class RPTokenFactory extends Factory {

      private static final BigInteger MIN_VALUE = BigInteger.ONE.negate();
      private static final BigInteger MAX_VALUE = BigInteger.valueOf(2).pow(127);
      private static final BigInteger RING_LENGTH = MAX_VALUE.add(BigInteger.ONE);
      private static final Token MIN_TOKEN = new RPToken(MIN_VALUE);
      private static final Token MAX_TOKEN = new RPToken(MAX_VALUE);

      private final MessageDigest prototype;
      private final boolean supportsClone;

      private RPTokenFactory() {
        prototype = createMessageDigest();
        boolean supportsClone;
        try {
          prototype.clone();
          supportsClone = true;
        } catch (CloneNotSupportedException e) {
          supportsClone = false;
        }
        this.supportsClone = supportsClone;
      }

      private static MessageDigest createMessageDigest() {
        try {
          return MessageDigest.getInstance("MD5");
        } catch (NoSuchAlgorithmException e) {
          throw new RuntimeException("MD5 doesn't seem to be available on this JVM", e);
        }
      }

      private MessageDigest newMessageDigest() {
        if (supportsClone) {
          try {
            return (MessageDigest) prototype.clone();
          } catch (CloneNotSupportedException ignored) {
          }
        }
        return createMessageDigest();
      }

      private BigInteger md5(ByteBuffer data) {
        MessageDigest digest = newMessageDigest();
        digest.update(data.duplicate());
        return new BigInteger(digest.digest()).abs();
      }

      @Override
      RPToken fromString(String tokenStr) {
        return new RPToken(new BigInteger(tokenStr));
      }

      @Override
      DataType getTokenType() {
        return DataType.varint();
      }

      @Override
      Token deserialize(ByteBuffer buffer, ProtocolVersion protocolVersion) {
        return new RPToken(TypeCodec.varint().deserialize(buffer, protocolVersion));
      }

      @Override
      Token minToken() {
        return MIN_TOKEN;
      }

      @Override
      RPToken hash(ByteBuffer partitionKey) {
        return new RPToken(md5(partitionKey));
      }

      @Override
      List<Token> split(Token startToken, Token endToken, int numberOfSplits) {
        // edge case: ]min, min] means the whole ring
        if (startToken.equals(endToken) && startToken.equals(MIN_TOKEN)) endToken = MAX_TOKEN;

        BigInteger start = ((RPToken) startToken).value;
        BigInteger end = ((RPToken) endToken).value;

        BigInteger range = end.subtract(start);
        if (range.compareTo(BigInteger.ZERO) < 0) range = range.add(RING_LENGTH);

        List<BigInteger> values = super.split(start, range, MAX_VALUE, RING_LENGTH, numberOfSplits);
        List<Token> tokens = Lists.newArrayListWithExpectedSize(values.size());
        for (BigInteger value : values) tokens.add(new RPToken(value));
        return tokens;
      }
    }

    private RPToken(BigInteger value) {
      this.value = value;
    }

    @Override
    public DataType getType() {
      return FACTORY.getTokenType();
    }

    @Override
    public Object getValue() {
      return value;
    }

    @Override
    public ByteBuffer serialize(ProtocolVersion protocolVersion) {
      return TypeCodec.varint().serialize(value, protocolVersion);
    }

    @Override
    public int compareTo(Token other) {
      assert other instanceof RPToken;
      return value.compareTo(((RPToken) other).value);
    }

    @Override
    public boolean equals(Object obj) {
      if (this == obj) return true;
      if (obj == null || this.getClass() != obj.getClass()) return false;

      return value.equals(((RPToken) obj).value);
    }

    @Override
    public int hashCode() {
      return value.hashCode();
    }

    @Override
    public String toString() {
      return value.toString();
    }
  }
}
