/*
 * Copyright DataStax, Inc.
 *
 * 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.datastax.oss.driver.internal.core.metadata.token;

import com.datastax.oss.driver.api.core.metadata.Node;
import com.datastax.oss.driver.api.core.metadata.token.Token;
import com.datastax.oss.driver.api.core.metadata.token.TokenRange;
import com.datastax.oss.driver.internal.core.util.NanoTime;
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableSetMultimap;
import com.datastax.oss.driver.shaded.guava.common.collect.SetMultimap;
import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Set;
import net.jcip.annotations.Immutable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

The token data for a given replication configuration. It's shared by all keyspaces that use that
configuration.
/**
 * The token data for a given replication configuration. It's shared by all keyspaces that use that
 * configuration.
 */
@Immutable
class KeyspaceTokenMap {

  private static final Logger LOG = LoggerFactory.getLogger(KeyspaceTokenMap.class);

  static KeyspaceTokenMap build(
      Map<String, String> replicationConfig,
      Map<Token, Node> tokenToPrimary,
      List<Token> ring,
      Set<TokenRange> tokenRanges,
      TokenFactory tokenFactory,
      ReplicationStrategyFactory replicationStrategyFactory,
      String logPrefix) {

    long start = System.nanoTime();
    try {
      ReplicationStrategy strategy = replicationStrategyFactory.newInstance(replicationConfig);

      SetMultimap<Token, Node> replicasByToken =
          strategy.computeReplicasByToken(tokenToPrimary, ring);
      SetMultimap<Node, TokenRange> tokenRangesByNode;
      if (ring.size() == 1) {
        // We forced the single range to ]minToken,minToken], make sure to use that instead of
        // relying
        // on the node's token
        ImmutableSetMultimap.Builder<Node, TokenRange> builder = ImmutableSetMultimap.builder();
        for (Node node : tokenToPrimary.values()) {
          builder.putAll(node, tokenRanges);
        }
        tokenRangesByNode = builder.build();
      } else {
        tokenRangesByNode = buildTokenRangesByNode(tokenRanges, replicasByToken);
      }
      return new KeyspaceTokenMap(ring, tokenRangesByNode, replicasByToken, tokenFactory);
    } finally {
      LOG.debug(
          "[{}] Computing keyspace-level data for {} took {}",
          logPrefix,
          replicationConfig,
          NanoTime.formatTimeSince(start));
    }
  }

  private final List<Token> ring;
  private final SetMultimap<Node, TokenRange> tokenRangesByNode;
  private final SetMultimap<Token, Node> replicasByToken;
  private final TokenFactory tokenFactory;

  private KeyspaceTokenMap(
      List<Token> ring,
      SetMultimap<Node, TokenRange> tokenRangesByNode,
      SetMultimap<Token, Node> replicasByToken,
      TokenFactory tokenFactory) {
    this.ring = ring;
    this.tokenRangesByNode = tokenRangesByNode;
    this.replicasByToken = replicasByToken;
    this.tokenFactory = tokenFactory;
  }

  Set<TokenRange> getTokenRanges(Node replica) {
    return tokenRangesByNode.get(replica);
  }

  Set<Node> getReplicas(ByteBuffer partitionKey) {
    return getReplicas(tokenFactory.hash(partitionKey));
  }

  Set<Node> getReplicas(Token token) {
    // If the token happens to be one of the "primary" tokens, get result directly
    Set<Node> nodes = replicasByToken.get(token);
    if (!nodes.isEmpty()) {
      return nodes;
    }
    // Otherwise, find the closest "primary" token on the ring
    int i = Collections.binarySearch(ring, token);
    if (i < 0) {
      i = -i - 1;
      if (i >= ring.size()) {
        i = 0;
      }
    }
    return replicasByToken.get(ring.get(i));
  }

  private static SetMultimap<Node, TokenRange> buildTokenRangesByNode(
      Set<TokenRange> tokenRanges, SetMultimap<Token, Node> replicasByToken) {
    ImmutableSetMultimap.Builder<Node, TokenRange> result = ImmutableSetMultimap.builder();
    for (TokenRange range : tokenRanges) {
      for (Node node : replicasByToken.get(range.getEnd())) {
        result.put(node, range);
      }
    }
    return result.build();
  }
}